| |
| """Few-shot adaptation grid — the DECISIVE TinyVLA-2 experiment. |
| |
| Zero-shot offline is confounded by A/B normalization scale (plan Q5). Few-shot is |
| fair: adapt each pretrained variant on N held-out-robot episodes from an equal |
| start, measure endpoint error on a disjoint test split. Hypothesis: C adapts with |
| the fewest examples (write a descriptor), then B (learn a fresh ID row), then A. |
| |
| Protocol per (variant, N): |
| - resume the pretrained final; fast path (expert + projectors + conditioning) is |
| trainable, backbone frozen (as in pretraining). A/B get a fresh embodiment id |
| (3, an unused row); C uses jaco's morphology descriptor (no new parameter). |
| - FT on jaco episodes [0:N], 5k steps, lr 5e-5. |
| - eval endpoint error on jaco episodes [test_start:test_start+20] (disjoint). |
| |
| Usage: python scripts/fewshot_grid.py (runs the whole grid, writes results json) |
| """ |
|
|
| from __future__ import annotations |
|
|
| import json |
| import subprocess |
| from pathlib import Path |
|
|
| ROOT = Path("/home/alexw/tinyvla") |
| PY = str(Path.home() / "tinyvla_venv/bin/python") |
| JACO = "heldout_jaco_play" |
| JACO_ROOT = str(Path.home() / "tinyvla_data/heldout/heldout_jaco_play") |
| TEST_START = 956 |
| SIZES = [100, 500] |
| VARIANTS = { |
| "A": {"dir": "tv2_A_id_native", "cond": "id", "space": "native", "emb": 3, "mkey": None}, |
| "B": {"dir": "tv2_B_id_canon", "cond": "id", "space": "canonical", "emb": 3, "mkey": None}, |
| "C": {"dir": "tv2_C_morph_canon", "cond": "morph", "space": "canonical", "emb": 0, "mkey": "jaco"}, |
| } |
| OUT = ROOT / "outputs" / "fewshot_results.json" |
|
|
|
|
| def write_config(v, spec, n): |
| cfg = ROOT / "outputs" / f"fewshot_{v}_{n}.yaml" |
| ds_extra = f" morph_key: {spec['mkey']}\n" if spec["mkey"] else "" |
| cfg.write_text(f"""output_dir: {ROOT}/outputs/fewshot_{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: {JACO} |
| root: {JACO_ROOT} |
| episodes: {n} |
| weight: 1.0 |
| embodiment_id: {spec['emb']} |
| {ds_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: 5000 |
| grad_clip: 10.0 |
| log_freq: 500 |
| save_freq: 5000 |
| 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 main(): |
| results = {} |
| for v, spec in VARIANTS.items(): |
| for n in SIZES: |
| tag = f"{v}_{n}" |
| cfg = write_config(v, spec, n) |
| ft_log = ROOT / "outputs" / f"fewshot_{tag}_ft.log" |
| print(f"[FT] {tag} ...", flush=True) |
| run([PY, str(ROOT / "scripts/train.py"), "--config", str(cfg)], ft_log) |
| ckpt = ROOT / "outputs" / f"fewshot_{tag}" / "final" |
| if not ckpt.exists(): |
| results[tag] = {"error": "FT produced no final ckpt"} |
| continue |
| ev_log = ROOT / "outputs" / f"fewshot_{tag}_eval.log" |
| cmd = [PY, str(ROOT / "scripts/eval_canonical.py"), |
| "--checkpoint", str(ckpt), "--dataset", JACO, "--root", JACO_ROOT, |
| "--ep-start", str(TEST_START), "--episodes", "20", "--oracle-ids", "4"] |
| if spec["mkey"]: |
| cmd += ["--morph-key", spec["mkey"]] |
| print(f"[EVAL] {tag} ...", flush=True) |
| run(cmd, ev_log) |
| line = [l for l in ev_log.read_text().splitlines() if "BEST" in l] |
| results[tag] = {"eval": line[-1] if line else "no BEST line"} |
| print(f" {tag}: {results[tag]['eval']}", flush=True) |
| OUT.write_text(json.dumps(results, indent=1)) |
| print("\n=== FEW-SHOT GRID DONE ===") |
| for k, v in results.items(): |
| print(f"{k}: {v}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
