#!/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(" 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())