File size: 7,059 Bytes
700dd75 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 | #!/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())
|