File size: 3,786 Bytes
5a2e445 | 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 | #!/usr/bin/env python
"""TidyBot (wheeled holonomic mobile manipulator) few-shot: does morphology
conditioning (C) adapt to a genuinely MOBILE morphology better than learned-ID (B)?
jaco (another arm) gave C≈B because its descriptor ≈ training arms. TidyBot is
is_mobile=1 / base_holonomic=1 — far from the fixed-arm training cluster, so this
is a sharper test. Its whole-body action is already EE-space (pos+quat+gripper),
so it fits the canonical head directly.
TidyBot has 110 episodes → FT on first N in {30, 80}, eval on episodes [95:110].
"""
from __future__ import annotations
import json, subprocess
from pathlib import Path
ROOT = Path("/home/alexw/tinyvla")
PY = str(Path.home() / "tinyvla_venv/bin/python")
DS = "heldout_tidybot"
DS_ROOT = str(Path.home() / "tinyvla_data/heldout/heldout_tidybot")
TEST_START = 95 # eval on eps [95:110]; FT on first N (<= 80), disjoint
SIZES = [30, 80]
VARIANTS = {
"B": {"dir": "tv2_B_id_canon", "cond": "id", "space": "canonical", "emb": 5, "mkey": None},
"C": {"dir": "tv2_C_morph_canon", "cond": "morph", "space": "canonical", "emb": 0, "mkey": "tidybot"},
}
OUT = ROOT / "outputs" / "tidybot_fewshot_results.json"
def write_config(v, spec, n):
cfg = ROOT / "outputs" / f"tidybot_{v}_{n}.yaml"
extra = f" morph_key: {spec['mkey']}\n" if spec["mkey"] else ""
cfg.write_text(f"""output_dir: {ROOT}/outputs/tidybot_{v}_{n}
wandb: null
seed: 42
morphology_descriptors: {ROOT}/configs/morphology/descriptors.yaml
resume_from: {ROOT}/outputs/{spec['dir']}/final
resume_step: 0
datasets:
- repo_id: {DS}
root: {DS_ROOT}
episodes: {n}
weight: 1.0
embodiment_id: {spec['emb']}
{extra}policy:
chunk_size: 50
n_action_steps: 50
image_size: 256
freeze_lm: true
freeze_vision_encoder: true
num_embodiments: 16
conditioning: {spec['cond']}
action_space: {spec['space']}
batch_size: 32
grad_accum: 1
num_workers: 8
lr: 5.0e-5
warmup_steps: 100
steps: 4000
grad_clip: 10.0
log_freq: 500
save_freq: 4000
mixed_precision: bf16
""")
return cfg
def run(cmd, log):
with open(log, "w") as f:
return subprocess.run(cmd, stdout=f, stderr=subprocess.STDOUT).returncode
def evalcmd(spec, ckpt_dir):
c = [PY, str(ROOT / "scripts/eval_canonical.py"), "--checkpoint", str(ckpt_dir),
"--dataset", DS, "--root", DS_ROOT, "--ep-start", str(TEST_START), "--episodes", "15"]
return c + (["--morph-key", "tidybot"] if spec["mkey"] else ["--oracle-ids", "5"])
def main():
res = {}
for v, spec in VARIANTS.items(): # zero-shot first
ev = ROOT / "outputs" / f"tidybot_{v}_0_eval.log"
run(evalcmd(spec, ROOT / f"outputs/{spec['dir']}/final"), ev)
line = [l for l in ev.read_text().splitlines() if "BEST" in l]
res[f"{v}_0"] = line[-1] if line else "?"
print(f"{v}_0: {res[f'{v}_0']}", flush=True); OUT.write_text(json.dumps(res, indent=1))
for v, spec in VARIANTS.items():
for n in SIZES:
tag = f"{v}_{n}"; cfg = write_config(v, spec, n)
print(f"[FT] {tag}", flush=True)
run([PY, str(ROOT / "scripts/train.py"), "--config", str(cfg)], ROOT / f"outputs/tidybot_{tag}_ft.log")
ckpt = ROOT / f"outputs/tidybot_{tag}/final"
if not ckpt.exists():
res[tag] = "no ckpt"; continue
ev = ROOT / f"outputs/tidybot_{tag}_eval.log"
run(evalcmd(spec, ckpt), ev)
line = [l for l in ev.read_text().splitlines() if "BEST" in l]
res[tag] = line[-1] if line else "?"
print(f"{tag}: {res[tag]}", flush=True); OUT.write_text(json.dumps(res, indent=1))
print("\n=== TIDYBOT FEW-SHOT DONE ===")
for k, val in res.items():
print(k, val)
if __name__ == "__main__":
main()
|