File size: 4,169 Bytes
0497fe1 11ae309 0497fe1 | 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 | #!/usr/bin/env python
"""Grow a checkpoint's embodiment/state/action dimensions without losing what it learned.
Adding the unitree + navigation packs widens three things:
max_state_dim 16 -> 32 (G1 Dex3 has 28 joint channels)
max_action_dim 8 -> 32 (same, and nav actions are 3-dim)
num_embodiments 16 -> 24 (packed ids reach 19)
Every affected weight is grown by ZERO padding, never re-initialised:
state_proj.weight (d, 16) -> (d, 32) new input columns = 0
action_in.weight (d, 8) -> (d, 32) new input columns = 0
action_out.weight (8, d) -> (32, d) new output rows = 0
action_out.bias (8,) -> (32,) new outputs = 0
embodiment emb (16, d) -> (24, d) new rows = 0
Zeroed input columns receive zero-padded inputs anyway, and zeroed output rows
predict 0 for dims no old robot had — so on the OLD data the expanded model is
bit-for-bit the old model, and the new dims start from a clean slate. That makes
"did the new data help or hurt?" answerable, which a re-initialised head would not.
Usage:
python scripts/expand_checkpoint.py outputs/tv2_C_scaled/final outputs/tv2_C_scaled_x32
"""
from __future__ import annotations
import json
import shutil
import sys
from pathlib import Path
import torch
from safetensors.torch import load_file, save_file
# param name suffix -> axis that must grow (0 = rows/out, 1 = cols/in)
GROW = {
"state_proj.weight": 1,
"action_in.weight": 1,
"action_out.weight": 0,
"action_out.bias": 0,
}
def main():
src = Path(sys.argv[1])
dst = Path(sys.argv[2])
new_state = int(sys.argv[3]) if len(sys.argv) > 3 else 32
new_action = int(sys.argv[4]) if len(sys.argv) > 4 else 32
new_emb = int(sys.argv[5]) if len(sys.argv) > 5 else 24
dst.mkdir(parents=True, exist_ok=True)
for f in src.iterdir():
if f.suffix in (".json", ".txt") or f.name.startswith("README"):
shutil.copy(f, dst / f.name)
cfg_path = dst / "config.json"
cfg = json.loads(cfg_path.read_text())
old_state = cfg.get("max_state_dim", 16)
old_action = cfg.get("max_action_dim", 8)
old_emb = cfg.get("num_embodiments", 16)
cfg["max_state_dim"] = new_state
cfg["max_action_dim"] = new_action
cfg["num_embodiments"] = new_emb
# input/output_features carry their own copy of the dims. Leaving them stale
# makes the checkpoint load with a 32-wide expert but an 8-wide declared
# action, so anything reading the feature spec silently truncates.
if "action" in cfg.get("output_features", {}):
cfg["output_features"]["action"]["shape"] = [new_action]
if "observation.state" in cfg.get("input_features", {}):
cfg["input_features"]["observation.state"]["shape"] = [new_state]
cfg_path.write_text(json.dumps(cfg, indent=2))
print(f"config: state {old_state}->{new_state} action {old_action}->{new_action} "
f"embodiments {old_emb}->{new_emb}")
weights = src / "model.safetensors"
sd = load_file(str(weights))
grown = []
for k, v in list(sd.items()):
target = None
for suffix, axis in GROW.items():
if k.endswith(suffix):
want = new_state if "state_proj" in k else new_action
if v.shape[axis] < want:
target = (axis, want)
break
# embodiment table: a 2-D embedding whose row count equals num_embodiments
if target is None and v.dim() == 2 and v.shape[0] == old_emb and "emb" in k.lower():
target = (0, new_emb)
if target is None:
continue
axis, want = target
shape = list(v.shape)
pad = want - shape[axis]
if pad <= 0:
continue
shape[axis] = pad
sd[k] = torch.cat([v, v.new_zeros(*shape)], dim=axis)
grown.append(f" {k}: {tuple(v.shape)} -> {tuple(sd[k].shape)}")
save_file(sd, str(dst / "model.safetensors"), metadata={"format": "pt"})
print("grown tensors:" if grown else "nothing to grow (already wide enough)")
print("\n".join(grown))
print(f"wrote {dst}")
if __name__ == "__main__":
main()
|