| |
| """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 |
| 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(): |
| 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() |
|
|
