lerobot/scripts/backfill_pi052_action_tokenizer.py

#!/usr/bin/env python

# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Audit or backfill checkpoint-local FAST artifacts for PI052 model repositories."""

from __future__ import annotations

import argparse
import hashlib
import io
import json
from pathlib import Path, PurePosixPath
from typing import Any

from huggingface_hub import CommitOperationAdd, HfApi, hf_hub_download

DEFAULT_REPOSITORIES = (
    "pepijn223/pi052_atomic4_01_baseline",
    "pepijn223/pi052_atomic4_02_lr_1e5",
    "pepijn223/pi052_atomic4_03_recipe_50_50",
    "pepijn223/pi052_atomic4_04_flow_weight_10",
    "pepijn223/pi052_atomic4_05_flow_repeat_1",
    "pepijn223/pi052_atomic4_06_ki_off",
)
CHECKPOINT_DIRECTORIES = (
    "",
    "checkpoints/003000/pretrained_model",
    "checkpoints/006000/pretrained_model",
    "checkpoints/009000/pretrained_model",
    "checkpoints/012000/pretrained_model",
)
TOKENIZER_DIRECTORY = "action_tokenizer"


def artifact_fingerprint(files: list[tuple[str, bytes]]) -> str:
    digest = hashlib.sha256()
    for relative_path, content in sorted(files):
        encoded_path = relative_path.encode()
        digest.update(len(encoded_path).to_bytes(8, "big"))
        digest.update(encoded_path)
        digest.update(len(content).to_bytes(8, "big"))
        digest.update(content)
    return digest.hexdigest()


def tokenizer_files(tokenizer_path: Path) -> list[tuple[str, Path]]:
    return [
        (path.relative_to(tokenizer_path).as_posix(), path)
        for path in sorted(tokenizer_path.rglob("*"))
        if path.is_file()
    ]


def _repo_path(directory: str, filename: str) -> str:
    return (PurePosixPath(directory) / filename).as_posix() if directory else filename


def _download_json(repo_id: str, path_in_repo: str, revision: str | None = None) -> dict[str, Any]:
    path = hf_hub_download(repo_id, path_in_repo, repo_type="model", revision=revision)
    return json.loads(Path(path).read_text())


def make_portable_preprocessor(config: dict[str, Any]) -> dict[str, Any]:
    config = json.loads(json.dumps(config))
    action_steps = [
        step for step in config["steps"] if step.get("registry_name") == "action_tokenizer_processor"
    ]
    if len(action_steps) != 1:
        raise ValueError(f"Expected one action tokenizer step, found {len(action_steps)}")
    action_step = action_steps[0]
    action_step["config"]["action_tokenizer_name"] = TOKENIZER_DIRECTORY
    action_step["artifacts"] = {"action_tokenizer_name": TOKENIZER_DIRECTORY}

    recipe_steps = [
        step for step in config["steps"] if step.get("registry_name") == "render_messages_processor"
    ]
    if len(recipe_steps) != 1 or not recipe_steps[0].get("config", {}).get("recipe"):
        raise ValueError("PI052 preprocessor does not contain an embedded training recipe")
    return config


def _json_operation(path_in_repo: str, content: dict[str, Any]) -> CommitOperationAdd:
    serialized = (json.dumps(content, indent=2) + "\n").encode()
    return CommitOperationAdd(path_in_repo=path_in_repo, path_or_fileobj=io.BytesIO(serialized))


def prepare_operations(
    repo_id: str,
    tokenizer_path: Path,
    revision: str | None = None,
) -> list[CommitOperationAdd]:
    operations: list[CommitOperationAdd] = []
    files = tokenizer_files(tokenizer_path)
    for directory in CHECKPOINT_DIRECTORIES:
        preprocessor_path = _repo_path(directory, "policy_preprocessor.json")
        operations.append(
            _json_operation(
                preprocessor_path,
                make_portable_preprocessor(_download_json(repo_id, preprocessor_path, revision)),
            )
        )
        for relative_path, local_path in files:
            operations.append(
                CommitOperationAdd(
                    path_in_repo=_repo_path(
                        directory,
                        f"{TOKENIZER_DIRECTORY}/{relative_path}",
                    ),
                    path_or_fileobj=str(local_path),
                )
            )
    return operations


def audit_repository(
    api: HfApi,
    repo_id: str,
    expected_tokenizer_fingerprint: str,
    revision: str | None = None,
) -> None:
    info = api.model_info(repo_id, revision=revision)
    repository_files = {sibling.rfilename for sibling in info.siblings or []}

    for directory in CHECKPOINT_DIRECTORIES:
        preprocessor_path = _repo_path(directory, "policy_preprocessor.json")
        policy_config_path = _repo_path(directory, "config.json")
        postprocessor_path = _repo_path(directory, "policy_postprocessor.json")
        for required_path in (preprocessor_path, policy_config_path, postprocessor_path):
            if required_path not in repository_files:
                raise FileNotFoundError(f"{repo_id}@{revision or 'main'} is missing {required_path}")

        preprocessor = _download_json(repo_id, preprocessor_path, revision)
        portable_preprocessor = make_portable_preprocessor(preprocessor)
        if preprocessor != portable_preprocessor:
            raise ValueError(f"{repo_id}:{preprocessor_path} is not portable")

        normalizer_steps = [
            step for step in preprocessor["steps"] if step.get("registry_name") == "normalizer_processor"
        ]
        if len(normalizer_steps) != 1 or "state_file" not in normalizer_steps[0]:
            raise ValueError(f"{repo_id}:{preprocessor_path} is missing normalizer state metadata")
        normalizer_path = _repo_path(directory, normalizer_steps[0]["state_file"])
        if normalizer_path not in repository_files:
            raise FileNotFoundError(f"{repo_id} is missing {normalizer_path}")

        remote_tokenizer_files: list[tuple[str, bytes]] = []
        for relative_path in _tokenizer_relative_paths(repository_files, directory):
            path_in_repo = _repo_path(directory, f"{TOKENIZER_DIRECTORY}/{relative_path}")
            downloaded = hf_hub_download(repo_id, path_in_repo, repo_type="model", revision=revision)
            remote_tokenizer_files.append((relative_path, Path(downloaded).read_bytes()))
        fingerprint = artifact_fingerprint(remote_tokenizer_files)
        if fingerprint != expected_tokenizer_fingerprint:
            raise ValueError(
                f"{repo_id}:{_repo_path(directory, TOKENIZER_DIRECTORY)} fingerprint "
                f"{fingerprint} != {expected_tokenizer_fingerprint}"
            )


def _tokenizer_relative_paths(repository_files: set[str], directory: str) -> list[str]:
    prefix = _repo_path(directory, TOKENIZER_DIRECTORY).rstrip("/") + "/"
    paths = sorted(path.removeprefix(prefix) for path in repository_files if path.startswith(prefix))
    if not paths:
        raise FileNotFoundError(f"Missing tokenizer artifact directory {prefix.rstrip('/')}")
    return paths


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--tokenizer-path", type=Path, required=True)
    parser.add_argument("--repo-id", action="append", dest="repo_ids")
    parser.add_argument("--revision")
    parser.add_argument("--dry-run", action="store_true")
    parser.add_argument("--audit-only", action="store_true")
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    tokenizer_path = args.tokenizer_path.resolve()
    if not tokenizer_path.is_dir():
        raise FileNotFoundError(f"Tokenizer directory does not exist: {tokenizer_path}")

    files = tokenizer_files(tokenizer_path)
    fingerprint = artifact_fingerprint([(relative_path, path.read_bytes()) for relative_path, path in files])
    api = HfApi()
    repositories = tuple(args.repo_ids or DEFAULT_REPOSITORIES)
    print(f"Tokenizer fingerprint: {fingerprint}")

    for repo_id in repositories:
        if args.audit_only:
            audit_repository(api, repo_id, fingerprint, args.revision)
            print(f"AUDIT OK {repo_id}@{args.revision or 'main'}")
            continue

        operations = prepare_operations(repo_id, tokenizer_path, args.revision)
        if args.dry_run:
            print(f"DRY RUN {repo_id}: {len(operations)} files")
            for operation in operations:
                print(f"  {operation.path_in_repo}")
            continue

        commit = api.create_commit(
            repo_id=repo_id,
            repo_type="model",
            operations=operations,
            commit_message="Embed fitted FAST tokenizer for portable PI052 checkpoints",
            revision=args.revision,
        )
        audit_repository(api, repo_id, fingerprint, commit.oid)
        print(f"BACKFILLED {repo_id}@{commit.oid}")


if __name__ == "__main__":
    main()