File size: 3,998 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 | #!/usr/bin/env python
"""UAV (quadrotor) few-shot: does morphology conditioning (C) beat learned-ID (B)
when the held-out robot is a genuinely different MORPHOLOGY (flying free-body, not
another arm)? The jaco result (C≈B) is expected — jaco's descriptor ≈ training
arms. The uav descriptor is far outside the training cluster, so this is the real
test of C>B.
For each variant in {B, C} and N in {100, 500}: FT the pretrained final on uav[0:N],
eval endpoint error on uav[test:test+20]. A is skipped (native already shown to
diverge; and uav native == canonical 6-DOF so it'd be redundant).
"""
from __future__ import annotations
import json, subprocess
from pathlib import Path
ROOT = Path("/home/alexw/tinyvla")
PY = str(Path.home() / "tinyvla_venv/bin/python")
UAV = "heldout_uav"
UAV_ROOT = str(Path.home() / "tinyvla_data/heldout/heldout_uav")
TEST_START = 1250 # 1275 eps; test on last 20 (leave a gap), FT on first N
SIZES = [100, 500]
VARIANTS = {
"B": {"dir": "tv2_B_id_canon", "cond": "id", "space": "canonical", "emb": 4, "mkey": None},
"C": {"dir": "tv2_C_morph_canon", "cond": "morph", "space": "canonical", "emb": 0, "mkey": "uav"},
}
OUT = ROOT / "outputs" / "uav_fewshot_results.json"
def write_config(v, spec, n):
cfg = ROOT / "outputs" / f"uav_{v}_{n}.yaml"
extra = f" morph_key: {spec['mkey']}\n" if spec["mkey"] else ""
cfg.write_text(f"""output_dir: {ROOT}/outputs/uav_{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: {UAV}
root: {UAV_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: 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():
res = {}
# zero-shot (N=0) baselines first
for v, spec in VARIANTS.items():
ev = ROOT / "outputs" / f"uav_{v}_0_eval.log"
cmd = [PY, str(ROOT/"scripts/eval_canonical.py"), "--checkpoint",
str(ROOT/f"outputs/{spec['dir']}/final"), "--dataset", UAV, "--root", UAV_ROOT,
"--ep-start", str(TEST_START), "--episodes", "20"]
cmd += ["--morph-key", "uav"] if spec["mkey"] else ["--oracle-ids", "4"]
print(f"[ZS] {v} N=0", flush=True); run(cmd, ev)
line=[l for l in ev.read_text().splitlines() if "BEST" in l]
res[f"{v}_0"]={"eval":line[-1] if line else "?"}; print(" ",res[f"{v}_0"]); OUT.write_text(json.dumps(res,indent=1))
# few-shot
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/uav_{tag}_ft.log")
ckpt=ROOT/f"outputs/uav_{tag}/final"
if not ckpt.exists(): res[tag]={"error":"no ckpt"}; continue
ev=ROOT/f"outputs/uav_{tag}_eval.log"
cmd=[PY,str(ROOT/"scripts/eval_canonical.py"),"--checkpoint",str(ckpt),"--dataset",UAV,"--root",UAV_ROOT,"--ep-start",str(TEST_START),"--episodes","20"]
cmd += ["--morph-key","uav"] if spec["mkey"] else ["--oracle-ids","4"]
print(f"[EVAL] {tag}", flush=True); run(cmd, ev)
line=[l for l in ev.read_text().splitlines() if "BEST" in l]
res[tag]={"eval":line[-1] if line else "?"}; print(" ",res[tag]); OUT.write_text(json.dumps(res,indent=1))
print("\n=== UAV FEW-SHOT DONE ==="); [print(k,v) for k,v in res.items()]
if __name__ == "__main__":
main()
|