GR00T-WholeBodyControl/gear_sonic/scripts/real_robot_deployment/check_checkpoint.py

#!/usr/bin/env python3
"""Dependency-free structural preflight for a GR00T SONIC checkpoint."""

from __future__ import annotations

import argparse
import json
from pathlib import Path
import struct
import sys


EXPECTED_MODALITIES = {
    "video": ["ego_view", "left_wrist", "right_wrist"],
    "state": [
        "left_leg",
        "right_leg",
        "waist",
        "left_arm",
        "right_arm",
        "left_gripper",
        "right_gripper",
        "head_joints",
        "projected_gravity",
    ],
    "action": [
        "motion_token",
        "left_hand_joints",
        "right_hand_joints",
        "head_joints",
    ],
    "language": ["annotation.human.task_description"],
}
EXPECTED_DIMS = {
    "state": {
        "left_leg": 6,
        "right_leg": 6,
        "waist": 3,
        "left_arm": 7,
        "right_arm": 7,
        "left_gripper": 1,
        "right_gripper": 1,
        "head_joints": 2,
        "projected_gravity": 3,
    },
    "action": {
        "motion_token": 64,
        "left_hand_joints": 1,
        "right_hand_joints": 1,
        "head_joints": 2,
    },
}
REQUIRED_FILES = (
    "config.json",
    "embodiment_id.json",
    "model.safetensors.index.json",
    "processor_config.json",
    "statistics.json",
)


def load_json(path: Path):
    try:
        with path.open("r", encoding="utf-8") as stream:
            return json.load(stream)
    except (OSError, json.JSONDecodeError) as error:
        raise ValueError(f"cannot read JSON {path}: {error}") from error


def validate_safetensors(path: Path) -> int:
    """Return tensor payload bytes after verifying the complete shard layout."""
    size = path.stat().st_size
    if size < 10:
        raise ValueError(f"safetensors shard is too small: {path} ({size} bytes)")
    with path.open("rb") as stream:
        header_size_raw = stream.read(8)
        if len(header_size_raw) != 8:
            raise ValueError(f"truncated safetensors header length: {path}")
        header_size = struct.unpack("<Q", header_size_raw)[0]
        if header_size <= 1 or 8 + header_size > size:
            raise ValueError(f"invalid safetensors header size in {path}: {header_size}")
        header_raw = stream.read(header_size)
    try:
        header = json.loads(header_raw.decode("utf-8"))
    except (UnicodeDecodeError, json.JSONDecodeError) as error:
        raise ValueError(f"invalid safetensors JSON header in {path}: {error}") from error

    max_end = 0
    tensor_count = 0
    for name, metadata in header.items():
        if name == "__metadata__":
            continue
        offsets = metadata.get("data_offsets") if isinstance(metadata, dict) else None
        if not isinstance(offsets, list) or len(offsets) != 2:
            raise ValueError(f"tensor {name!r} has invalid offsets in {path}")
        start, end = offsets
        if not isinstance(start, int) or not isinstance(end, int) or start < 0 or end < start:
            raise ValueError(f"tensor {name!r} has invalid offsets {offsets} in {path}")
        max_end = max(max_end, end)
        tensor_count += 1

    expected_size = 8 + header_size + max_end
    if tensor_count == 0 or size != expected_size:
        raise ValueError(
            f"incomplete/invalid safetensors shard {path}: size={size}, expected={expected_size}"
        )
    return max_end


def stat_width(statistics: dict, modality: str, key: str) -> int:
    values = statistics[modality][key]
    for field in ("q01", "mean", "min"):
        value = values.get(field)
        if isinstance(value, list):
            return len(value)
    raise ValueError(f"cannot determine {modality}.{key} width from statistics")


def check_checkpoint(checkpoint: Path, embodiment: str) -> None:
    if not checkpoint.is_dir():
        raise ValueError(f"checkpoint directory does not exist: {checkpoint}")
    for filename in REQUIRED_FILES:
        path = checkpoint / filename
        if not path.is_file() or path.stat().st_size == 0:
            raise ValueError(f"missing or empty checkpoint file: {path}")

    config = load_json(checkpoint / "config.json")
    if config.get("action_horizon") != 40:
        raise ValueError(f"expected action_horizon=40, got {config.get('action_horizon')}")

    processor = load_json(checkpoint / "processor_config.json")
    try:
        modalities = processor["processor_kwargs"]["modality_configs"][embodiment]
    except KeyError as error:
        raise ValueError(f"checkpoint does not contain embodiment {embodiment!r}") from error
    for modality, expected_keys in EXPECTED_MODALITIES.items():
        actual = modalities[modality]["modality_keys"]
        if actual != expected_keys:
            raise ValueError(
                f"{embodiment}.{modality} keys differ: expected {expected_keys}, got {actual}"
            )

    all_statistics = load_json(checkpoint / "statistics.json")
    if embodiment not in all_statistics:
        raise ValueError(f"statistics do not contain embodiment {embodiment!r}")
    statistics = all_statistics[embodiment]
    for modality, expected in EXPECTED_DIMS.items():
        for key, width in expected.items():
            actual_width = stat_width(statistics, modality, key)
            if actual_width != width:
                raise ValueError(
                    f"{modality}.{key} width differs: expected {width}, got {actual_width}"
                )

    index = load_json(checkpoint / "model.safetensors.index.json")
    shards = sorted(set(index.get("weight_map", {}).values()))
    if not shards:
        raise ValueError("model.safetensors.index.json contains no shards")
    payload_total = 0
    for filename in shards:
        path = checkpoint / filename
        if not path.is_file():
            raise ValueError(f"checkpoint shard is still missing: {path}")
        payload = validate_safetensors(path)
        payload_total += payload
        print(f"  valid shard: {filename} ({path.stat().st_size:,} bytes)")

    indexed_total = index.get("metadata", {}).get("total_size")
    if isinstance(indexed_total, int) and payload_total != indexed_total:
        raise ValueError(
            f"tensor payload total differs: index={indexed_total}, shards={payload_total}"
        )

    print("Checkpoint preflight passed")
    print(f"  checkpoint: {checkpoint}")
    print(f"  embodiment: {embodiment}")
    print("  video:      ego_view + left_wrist + right_wrist (separate)")
    print("  state dim:  36")
    print("  action dim: 68 = 64 + 1 + 1 + 2")
    print("  horizon:    40")


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--checkpoint", type=Path, required=True)
    parser.add_argument("--embodiment", default="unitree_g1_sonic")
    args = parser.parse_args()
    try:
        check_checkpoint(args.checkpoint.expanduser(), args.embodiment)
    except (OSError, ValueError, KeyError, TypeError) as error:
        print(f"Checkpoint preflight failed: {error}", file=sys.stderr)
        return 1
    return 0


if __name__ == "__main__":
    raise SystemExit(main())