File size: 2,648 Bytes
1ee3857 | 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 | #!/usr/bin/env python
"""Strip torch.compile's `_orig_mod.` prefix from a saved checkpoint.
The physical_ai_ft run was launched with `compile: true` and saved the COMPILED
module, so 167 of its 489 tensors are named `expert._orig_mod.*` while the model
class expects `expert.*`.
This matters more than a normal key mismatch. LeRobot's `PreTrainedPolicy.
from_pretrained` loads with `strict=False` and merely LOGS missing keys, and
`FlowMatchingExpert.__init__` zero-initialises `action_out` — so a model loaded
from the raw checkpoint predicts velocity 0, and `predict_action_chunk` returns
its initial Gaussian noise unchanged. That is a plausible-looking bad number, not
a crash: nothing anywhere in the stack tells you the action expert is random.
Usage:
python scripts/rekey_checkpoint.py <src_ckpt_dir> <dst_ckpt_dir>
"""
from __future__ import annotations
import shutil
import sys
from pathlib import Path
from safetensors.torch import load_file, save_file
PREFIX = "expert._orig_mod."
def main():
src = Path(sys.argv[1] if len(sys.argv) > 1
else "/home/alexw/tinyvla_data/b200/outputs/physical_ai_ft/final")
dst = Path(sys.argv[2] if len(sys.argv) > 2
else "/home/alexw/tinyvla/outputs/physical_ai_ft_fixed")
dst.mkdir(parents=True, exist_ok=True)
for f in src.iterdir():
if f.suffix == ".json" or f.name.startswith("README"):
shutil.copy(f, dst / f.name)
sd = load_file(str(src / "model.safetensors"))
renamed = {k.replace(PREFIX, "expert."): v for k, v in sd.items()}
n_hit = sum(1 for k in sd if k.startswith(PREFIX))
assert len(renamed) == len(sd), "rename collided — two keys mapped to one name"
print(f"renamed {n_hit}/{len(sd)} tensors ({PREFIX} -> expert.)")
save_file(renamed, str(dst / "model.safetensors"), metadata={"format": "pt"})
# Hard gate: build the model from the checkpoint's own config and demand a
# STRICT load. from_pretrained would silently accept a missing expert.
from tinyvla.modeling_tinyvla import TinyVLAPolicy
pol = TinyVLAPolicy.from_pretrained(str(dst))
fresh = load_file(str(dst / "model.safetensors"))
missing, unexpected = pol.load_state_dict(fresh, strict=True), None
n_model = len(pol.state_dict())
aow = pol.expert.action_out.weight.abs().mean().item()
print(f"{len(fresh)}/{n_model} tensors, strict=True OK")
print(f"expert.action_out |w| = {aow:.4f} (zero at init -> nonzero proves training)")
assert aow > 1e-3, "action_out is at its zero init — the expert never trained"
print(f"wrote {dst}")
if __name__ == "__main__":
main()
|