diff --git a/.gitattributes b/.gitattributes index 92e08d7ac8aa20258069506f45b3340b42eaeae5..7b1cef52b188704a13fa2de49a3ae718f1429423 100644 --- a/.gitattributes +++ b/.gitattributes @@ -39,3 +39,15 @@ videos/ep02_task2_succ.mp4 filter=lfs diff=lfs merge=lfs -text videos/ep03_task3_fail.mp4 filter=lfs diff=lfs merge=lfs -text videos/ep04_task4_succ.mp4 filter=lfs diff=lfs merge=lfs -text videos/ep05_task5_fail.mp4 filter=lfs diff=lfs merge=lfs -text +tinyvla2/assets/so_arm/assets/base_motor_holder_so101_v1.stl filter=lfs diff=lfs merge=lfs -text +tinyvla2/assets/so_arm/assets/base_so101_v2.stl filter=lfs diff=lfs merge=lfs -text +tinyvla2/assets/so_arm/assets/motor_holder_so101_base_v1.stl filter=lfs diff=lfs merge=lfs -text +tinyvla2/assets/so_arm/assets/motor_holder_so101_wrist_v1.stl filter=lfs diff=lfs merge=lfs -text +tinyvla2/assets/so_arm/assets/moving_jaw_so101_v1.stl filter=lfs diff=lfs merge=lfs -text +tinyvla2/assets/so_arm/assets/rotation_pitch_so101_v1.stl filter=lfs diff=lfs merge=lfs -text +tinyvla2/assets/so_arm/assets/sts3215_03a_no_horn_v1.stl filter=lfs diff=lfs merge=lfs -text +tinyvla2/assets/so_arm/assets/sts3215_03a_v1.stl filter=lfs diff=lfs merge=lfs -text +tinyvla2/assets/so_arm/assets/under_arm_so101_v1.stl filter=lfs diff=lfs merge=lfs -text +tinyvla2/assets/so_arm/assets/upper_arm_so101_v1.stl filter=lfs diff=lfs merge=lfs -text +tinyvla2/assets/so_arm/assets/wrist_roll_follower_so101_v1.stl filter=lfs diff=lfs merge=lfs -text +tinyvla2/assets/so_arm/assets/wrist_roll_pitch_so101_v2.stl filter=lfs diff=lfs merge=lfs -text diff --git a/tinyvla2/CLAUDE.md b/tinyvla2/CLAUDE.md new file mode 100644 index 0000000000000000000000000000000000000000..0827014d53a4cd168d95629e22484d2832161298 --- /dev/null +++ b/tinyvla2/CLAUDE.md @@ -0,0 +1,153 @@ +# TinyVLA — project context + +Small vision-language-action models. Two lines of work: + +- **v1 (shipped)** — dual-rate VLA for LIBERO: slow Qwen3.5 semantic path (cacheable, + 0.5–2 Hz) + fast flow-matching expert (10–20 Hz). **86% LIBERO-spatial, 84% 4-suite + average** — beats OpenVLA-7B on every suite at 0.64B params / 140M online. +- **v2 (research)** — cross-embodiment: canonical action space + morphology + conditioning across 10 robots. Best: `tv2_C_scaled` (ratio-to-baseline **0.81** + in-training / **0.88** held-out; was 0.99/1.17 before scaling). + +Everything (checkpoints, results, wiring) lives in the **private** HF repo +`AlexWortega/tinyvla` (v1 at the root, v2 under `tinyvla2/`). + +## Environment + +```bash +source ~/tinyvla_venv/bin/activate # NOT physics_venv +``` + +Pinned, non-obvious, and hard-won: +- lerobot **0.6.1**, torch **2.10.0+cu128**, transformers **5.5.4** (lerobot's + `transformers-dep` extra; older breaks `eo1`). +- **torchcodec 0.10 CPU build from PyPI** — the cu128 wheel clashes with the system + NPP libs. Video decode is the training bottleneck, use 12+ dataloader workers. +- `placo` for SO-101 forward kinematics; needs `assets/so_arm/so101_new_calib.urdf` + **plus the 13 STL meshes** (placo loads collision geometry, it won't run without them). +- `fla-core` + `causal-conv1d` give Qwen3.5 linear-attention its fast path (~5×); + they require CUDA — any CPU-only smoke test of the model will crash inside + `causal_conv1d`. +- Data lives in `~/tinyvla_data` (gitignored), HF cache in `~/.cache/huggingface`. + +## Architecture + +``` +SLOW (cacheable @1Hz) FAST (10–20 Hz, ~140M online) +image → Qwen3.5-0.8B vision tower ─┬────► 64 spatial tokens/cam ──┐ +instruction ──────────────────────┤ │ + ▼ proprio ┤ + LM (first 12 of 24 layers) morphology/ID ┤ + ▼ ▼ + 8–16 readout latents ────────────────► flow-matching expert + (38M, chunk 50, 10 Euler) +``` + +Variant switches in `TinyVLAConfig`: +- `conditioning`: `id` (learned embedding) | `morph` (16-number descriptor → MLP) | + `morph_qwen` (robot text through the shared LM) +- `action_space`: `native` | `canonical` (base-frame EE deltas) +- `vlm_native`: demos + robot text go INTO the LM sequence (dual-rate preserved) +- `use_demo_conditioning`, `demo_only`: in-context exemplar channels (all dead, see below) + +## Data pipeline + +``` +LeRobot dataset → scripts/build_ee_sidecar.py (offline, once, NO video decode) + SO-101: FK(joints) via placo; others: EE pose from state + → ~/tinyvla_data/ee_sidecar/.parquet + → src/tinyvla/data/canonical.py (10 Hz resample, base-frame + left-delta rotvec, gripper q1–q99, quantile norm to [-1,1]) + → CanonicalSource → WeightedMixtureDataset → policy +``` + +`docs/WIRING.md` has the full embodiment registry (id ↔ morph_key ↔ dataset ↔ +canonicalization source ↔ weight ↔ train/held-out). + +Robot pool: SO-101 (40 community sets, FK), Bridge/WidowX 53K eps, RT-1/Google 87K eps, +UR5, xArm, DLR-EDAN, TidyBot, Stretch, Jaco, LeKiwi (held out), RoboCasa365/PandaOmron. + +## Running things + +```bash +# training (any variant — the config picks the mode) +PYTHONUNBUFFERED=1 python scripts/train.py --config configs/tv2_C_scaled.yaml + +# physical-space eval across robots (the metric that matters) +python scripts/eval_similar_robots.py # ratio = endpoint err ÷ zero-floor + +# LIBERO closed-loop (v1) +MUJOCO_GL=egl python scripts/eval_libero.py --checkpoint \ + --suite libero_spatial --episodes 100 --n-action-steps 10 --action-repeat 1 +``` + +Long runs: launch with `nohup`, watch via checkpoints (stdout is block-buffered +under nohup — use `PYTHONUNBUFFERED=1` or you see nothing for hours). + +## Evaluation — read this before trusting any number + +1. **Always run the zero-prediction control.** Per-step action error is degenerate + (jaco moves ~3.6 mm/step; predicting zeros scores ~4 mm). Use the + **integrated-trajectory endpoint error over the 50-step chunk**, and divide by the + per-robot **zero-floor** (what "predict mean motion" achieves). Only the ratio is + comparable across robots; ratio < 1.0 means the model beats the trivial baseline. +2. **Check both held-out robots, not one.** "canonical ≫ native" looked decisive on + jaco and reversed on stretch. +3. Absolute mm are meaningless across robots — different motion scales. + +## Results + +### v1 — LIBERO (closed-loop success) +| suite | TinyVLA | OpenVLA-7B | +|---|---|---| +| spatial | **86%** (dual-rate @1 Hz; 79% full-rate) | 84.7% | +| object | **94%** | 88.4% | +| goal | **86%** | 79.2% | +| long | **70%** | 53.7% | +| **avg** | **84.0%** | 76.5% | + +Dual-rate (semantic latent refreshed at 1 Hz) matched or beat full-rate — the core +architectural claim, verified closed-loop. + +### v2 — cross-embodiment (ratio to zero-floor, lower is better) +| | in-training | held-out | +|---|---|---| +| C-diverse (8 robots, 20k steps, frozen LM) | 0.99 | 1.17 | +| **C-scaled (9.1M frames, 60k steps, LM unfrozen)** | **0.81** | **0.88** | + +Per-robot final: jaco 0.69, xarm 0.64, dlr_edan 0.69, stretch 0.69, LeKiwi (held-out) +1.07, ur5 1.22 (hardest — the fastest robot, 9.9 mm/step; its canonicalization is fine). + +**Negative results (all documented in `tinyvla2/RESULTS.md`):** +- Robot description as text in the prompt (frozen LM) — hurts. +- Morphology via the shared Qwen + unfrozen LM — 3× more content-sensitive, no + absolute gain. +- In-context demo conditioning — **five architectures**, all dead (real vs garbage + exemplars indistinguishable): MLP same-task → MLP cross-task → 24× capacity → + demos inside the LM attention stream → demos as the only identity channel. + +## Lessons + +1. **Check against the trivial baseline before running ablations.** Five conditioning + mechanisms moved nothing while the model sat at ratio 0.99 — we were comparing + variants of an underfit model and measuring noise. Data + steps + unfreezing the LM + moved everything. +2. **Silent data-convention bugs cost the most time.** LIBERO went 0% → 60% purely by + fixing three of them: robosuite renders are **180°-rotated** (a vertical flip alone + leaves a mirror), eef rotvec is antipodal vs the dataset, and the dataset's + `fps: 10` metadata is wrong (frames are native 20 Hz). A **GT-action replay + harness** (`scripts/replay_libero.py`) diagnoses all of these independently of the + policy — build it first. +3. **Unfreezing the LM is the single biggest lever** (LIBERO 60% → 86%). +4. Cross-morphology transfer converges slowest — LeKiwi was still improving at 60k + (1.24 → 1.13 → 1.07) when in-training had plateaued at 40k. +5. Some community datasets are broken at source (meta frame count ≠ parquet rows); + validate `len(reader.hf_dataset) == meta.total_frames` and quarantine mismatches. + +## State + +Best checkpoints: **v1** `outputs/libero_ft4/final` (86% LIBERO) · **v2** +`outputs/tv2_C_scaled/final` (0.81/0.88). Currently training: `tv2_robocasa` — +continue-train from C-scaled on RoboCasa365 (sim kitchens, PandaOmron mobile +manipulator, 1.5M frames) at weight 0.5 with the old robots replayed at 0.5, to +separate new-domain integration from catastrophic forgetting. diff --git a/tinyvla2/REPORT.md b/tinyvla2/REPORT.md new file mode 100644 index 0000000000000000000000000000000000000000..b649d5dc3218bdb656aaf72baeeb290ddfb5d35b --- /dev/null +++ b/tinyvla2/REPORT.md @@ -0,0 +1,171 @@ +--- +license: apache-2.0 +base_model: Qwen/Qwen3.5-0.8B +tags: + - robotics + - vla + - lerobot + - libero + - dual-rate +pipeline_tag: robotics +--- + +# TinyVLA: a dual-rate VLA that beats OpenVLA-7B on LIBERO-Spatial at 0.64B params + +**Headline result: 86% success on LIBERO-Spatial (n=100) in dual-rate mode (semantic latent +refreshed at 1 Hz), 79% full-rate — vs OpenVLA-7B's 84.7% — with 140M parameters on the +online control path.** Built and trained end-to-end in ~5 days on a single RTX A6000. + +## 1. Architecture + +``` + SLOW / SMART path (0.5–2 Hz, cached) +instruction ─────────────────────┐ +camera → Qwen3.5-0.8B vision tower → Qwen3.5 language stack (first 12 of 24 layers) + │ │ + │ 8 readout tokens → linear → semantic latent (8×512) + │ │ cached between refreshes +──────────────────┼───────────────────────┼────────────────────────── + │ FAST / MOTOR path (10–20 Hz, ~140M online) + └── 64 spatial tokens per camera ─┐ +proprio (padded 16d) ───────────────────────────────┤ +embodiment-id embedding ────────────────────────────┤ + ▼ + flow-matching action expert (38M, 12 blocks, + interleaved cross-/self-attention, AdaLN-free) + ▼ + 50-step action chunk (padded 8d + validity mask) +``` + +- **Single shared encoder**: the Qwen3.5-0.8B native ViT (100.6M) feeds both paths. +- **Slow path**: image + instruction + 8 trailing pad-token "readout" slots through the + truncated language stack (L/2 layer-skip a la SmolVLA); their hidden states → latent 8×512. +- **Fast path**: raw vision-tower spatial tokens (bypassing the LM) + proprio + embodiment + token, cross-attended by the flow expert. 10 Euler steps at inference. +- **Flow matching**: linear-OT path, t ~ Beta(1.5, 1), velocity target, masked MSE. +- Totals: 643M stored / 39.5M–575M trainable depending on stage / ~140M online per tick. +- Implemented as an **out-of-tree LeRobot policy plugin** + (`--policy.type=tinyvla --policy.discover_packages_path=tinyvla`, lerobot 0.6.1). + +## 2. Training recipe + +| Stage | Data | Steps | Trainable | Result | +|---|---|---|---|---| +| 1. Overfit smoke | svla_so100_stacking (56 eps) | 5k | expert+proj (40M) | loss 1.885→0.099 | +| 2. Mixture pretrain | 40 SO100/101 community datasets (~10.7K eps, weight 0.65, per-dataset proportional) + Bridge V2 10K eps (0.35) | 40k, bs 64×2, lr 1e-4 | 39.5M (backbone frozen) | held-out RMSE@t0 0.30σ; **stale-latent 2s == fresh** | +| 2b. Staleness aug | same | from step 20k | — | latent from frame t−k, k~U(0,2s), p=0.5 | +| 3. Spatial distillation | + 132K teacher labels (Qwen3.5-4B, center-point → 32×32 grid), CE ×0.15 on cam0 tokens | 20k, lr 5e-5 | same + head | head top-1 20.8%, top-5 44.8% (random 0.1%) | +| 4. LIBERO FT (ft4) | HuggingFaceVLA/libero (1693 eps, all suites) | 40k, bs 48, lr 1e-4 / backbone 0.1× | **575M** (LM layers + vision unfrozen; vocab embeddings frozen) | loss 0.066 | + +Ablation FTs: ft1 (all frozen) → 0% closed-loop; ft2 (vision unfrozen) → 60%; +ft3 (= ft2 from distilled ckpt) → 61%; **ft4 (LM+vision unfrozen) → 79/86%**. +Unfreezing the LM was the single largest gain; distillation was neutral on LIBERO +(likely overwritten by vision unfreezing during FT — bake it into pretraining next time). + +Cross-embodiment handling: canonical schema (cam0=scene/cam1=wrist via regex mapping, +per-dataset mean/std normalization, action pad-to-8 + validity mask, state pad-to-16, +learned embodiment-id token). Weighted mixture sampler with per-source retry on corrupt +frames (IndexError/KeyError/decode RuntimeError). + +## 3. Results + +### LIBERO — all four suites (dual-rate, latent @ 1 Hz) + +| Suite | TinyVLA | OpenVLA-7B (published) | +|---|---|---| +| Spatial | **86%** (n=100; full-rate 79%) | 84.7% | +| Object | **94%** (n=50) | 88.4% | +| Goal | **86%** (n=50) | 79.2% | +| Long (libero_10) | **70%** (n=50) | 53.7% | +| **Average** | **84.0%** | 76.5% | + +Beats OpenVLA-7B on every suite at 1/11 the parameters and ~1/50 the online compute. +XS-VLA (90.3 avg) remains ahead — their latent flow matching is future work for us. + +### LIBERO-Spatial (100 episodes, 10 per task, init states 0-9) + +| Mode | Success | +|---|---| +| Full-rate (latent every chunk) | 79/100 = **79%** | +| **Dual-rate (latent @ 1 Hz)** | 86/100 = **86%** | + +Dual-rate ≥ full-rate consistently across all evals (also 65% vs 58% at ft3, n=50) — +a stable cached latent appears to act as temporal regularization. + +### Context (published numbers, their own eval protocols) + +| Model | Params | LIBERO-Spatial | 4-suite avg | +|---|---|---|---| +| MolmoAct2-LIBERO | 5B | 98.4% (**97.2% re-measured on our rig**, n=500) | **98.25%** | +| π0 | 3.3B | 98% | — | +| SmolVLA-2.25B | 2.25B | 93% | — | +| XS-VLA | 0.26B | ~93% | 90.3% | +| MolmoAct-7B (v1) | 7B | 87.0% | 86.6% | +| **TinyVLA (this)** | **0.64B (140M online)** | **86% dual-rate** | **84.0%** | +| OpenVLA | 7B | 84.7% | 76.5% | + +**MolmoAct2 context** (from their cards/repo, no numeric tables published outside the +LIBERO-LeRobot card): 5B params on Molmo2-ER backbone + flow expert; claims to beat π0.5; +Molmo2-ER reportedly surpasses GPT-5 / Gemini-Robotics-ER-1.5 on 13 embodied-reasoning +benchmarks; #1 VLA on the MolmoSpace leaderboard. Inference needs ~26GB fp32 (16GB bf16), +cuda-graphs for ~2× speedup — H100-class serving vs our Raspberry-Pi-class 140M fast path. +Notably they ship **MolmoAct2-SO100_101** (absolute joint control) — a 5B competitor on our +target embodiment; a head-to-head offline eval on SO101 held-out data is an obvious next test. + +Honesty notes: our n=100 (±8pp) vs their 500; single suite; eval covers init states 0–9 +per task (demos of same tasks/init states are in training data — this is the standard +LIBERO protocol used by all listed models). Multi-suite evals (object/goal/long) and an +identically-protocoled SmolVLA baseline run are in progress. + +### Dual-rate validation (the architectural claim) + +- Offline: held-out chunk MSE with latent from a frame 2s earlier == fresh (0.772 vs 0.772). +- Closed-loop: 86% @ 1 Hz refresh vs 79% full-rate. +- Deployment implication (Raspberry Pi 5): the 0.5B slow path can run at 0.5–1 Hz + (llama.cpp GGUF Q4 territory), only the 140M fast path needs real-time rates. + +## 4. Pitfalls we hit (read before reproducing) + +1. **robosuite/LIBERO renders are 180°-rotated** vs recorded datasets. A vertical flip + alone leaves a horizontal MIRROR — the policy reaches to the wrong side and scores 0%. + Use `img[::-1, ::-1]`. (lerobot's `LiberoEnv.render()` flips for visualization but + `_format_raw_obs` feeds policies raw rotated frames — check your eval path.) +2. **eef orientation antipodal ambiguity**: env quat → rotvec gives −π-ish x-component, + datasets store +π-ish. Canonicalize (flip antipodal representation when rotvec_x < 0) + or your normalized state input is off by ~17σ. +3. **HuggingFaceVLA/libero meta says fps=10 but frames are native 20 Hz** (~161 + frames/episode ≈ 8s demos). Do NOT action-repeat at eval. +4. **GT-replay harness is the decisive convention test**: replay dataset actions in the env + (first episode per task ↔ init_state 0). 4/5 success = conventions right; policy-independent. +5. HuggingFaceVLA community_dataset v1/v3 are aggregates of v2.1 sub-datasets → convert with + lerobot's `convert_dataset_v21_to_v30.py`; 4/44 sources have meta≠parquet row counts + (broken at source) — validate `len(reader.hf_dataset) == meta.total_frames`. +6. Qwen3.5 custom forwards need `mm_token_type_ids` for M-RoPE; patchify must match the + official processor bit-exact (verify!). Linear-attention needs fla-core + causal-conv1d + or you get a 5× slower torch fallback. +7. Flow-matching loss can drop while conditioning is ignored — A/B with shuffled + observations to verify the model actually uses them. + +## 5. Repository layout + +``` +src/tinyvla/ # LeRobot plugin: configuration/modeling/processor + modules +scripts/train.py # accelerate loop: weighted mixture, staleness aug, distill loss, resume +scripts/eval_libero.py # closed-loop LIBERO eval (all convention fixes, dual-rate, video) +scripts/replay_libero.py, diag_obs_gap.py # convention diagnostics +scripts/prepare_community_v3.py, label_spatial.py, eval_offline.py +configs/ # stage2/stage3/libero_ft* yaml +checkpoints/ # stage2_final, stage3_final, libero_ft4_final (main) +videos/ # sample eval episodes +``` + +## 6. Quick start + +```python +from tinyvla.modeling_tinyvla import TinyVLAPolicy +policy = TinyVLAPolicy.from_pretrained("AlexWortega/tinyvla") # ft4 checkpoint at root +# closed-loop LIBERO: python scripts/eval_libero.py --checkpoint \ +# --suite libero_spatial --episodes 100 --n-action-steps 10 --refresh-s 1.0 +``` + +Trained 2026-08-14 → 2026-08-19 on 1× RTX A6000 (48GB), lerobot 0.6.1, torch 2.10, transformers 5.5.4. diff --git a/tinyvla2/assets/so_arm/assets/base_motor_holder_so101_v1.stl b/tinyvla2/assets/so_arm/assets/base_motor_holder_so101_v1.stl new file mode 100644 index 0000000000000000000000000000000000000000..ac9c38076fe1036517faf0bccadea5de9dce0097 --- /dev/null +++ b/tinyvla2/assets/so_arm/assets/base_motor_holder_so101_v1.stl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8cd2f241037ea377af1191fffe0dd9d9006beea6dcc48543660ed41647072424 +size 1877084 diff --git a/tinyvla2/assets/so_arm/assets/base_so101_v2.stl b/tinyvla2/assets/so_arm/assets/base_so101_v2.stl new file mode 100644 index 0000000000000000000000000000000000000000..503d30be06a91e401ba8d46ebb7e650866229550 --- /dev/null +++ b/tinyvla2/assets/so_arm/assets/base_so101_v2.stl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:bb12b7026575e1f70ccc7240051f9d943553bf34e5128537de6cd86fae33924d +size 471584 diff --git a/tinyvla2/assets/so_arm/assets/motor_holder_so101_base_v1.stl b/tinyvla2/assets/so_arm/assets/motor_holder_so101_base_v1.stl new file mode 100644 index 0000000000000000000000000000000000000000..f8e3d75c027f28bb672f830ec6e0795567c1b7c9 --- /dev/null +++ b/tinyvla2/assets/so_arm/assets/motor_holder_so101_base_v1.stl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:31242ae6fb59d8b15c66617b88ad8e9bded62d57c35d11c0c43a70d2f4caa95b +size 1129384 diff --git a/tinyvla2/assets/so_arm/assets/motor_holder_so101_wrist_v1.stl b/tinyvla2/assets/so_arm/assets/motor_holder_so101_wrist_v1.stl new file mode 100644 index 0000000000000000000000000000000000000000..e55b7194683c6ac301504c6f59137362f0ebd13e --- /dev/null +++ b/tinyvla2/assets/so_arm/assets/motor_holder_so101_wrist_v1.stl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:887f92e6013cb64ea3a1ab8675e92da1e0beacfd5e001f972523540545e08011 +size 1052184 diff --git a/tinyvla2/assets/so_arm/assets/moving_jaw_so101_v1.stl b/tinyvla2/assets/so_arm/assets/moving_jaw_so101_v1.stl new file mode 100644 index 0000000000000000000000000000000000000000..eb17d253df8a84a88472ecc7f859d3b8b4d78884 --- /dev/null +++ b/tinyvla2/assets/so_arm/assets/moving_jaw_so101_v1.stl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:785a9dded2f474bc1d869e0d3dae398a3dcd9c0c345640040472210d2861fa9d +size 1413584 diff --git a/tinyvla2/assets/so_arm/assets/rotation_pitch_so101_v1.stl b/tinyvla2/assets/so_arm/assets/rotation_pitch_so101_v1.stl new file mode 100644 index 0000000000000000000000000000000000000000..b536cb4100c1f204f8a9d9b182acdc4a3afbc66c --- /dev/null +++ b/tinyvla2/assets/so_arm/assets/rotation_pitch_so101_v1.stl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9be900cc2a2bf718102841ef82ef8d2873842427648092c8ed2ca1e2ef4ffa34 +size 883684 diff --git a/tinyvla2/assets/so_arm/assets/sts3215_03a_no_horn_v1.stl b/tinyvla2/assets/so_arm/assets/sts3215_03a_no_horn_v1.stl new file mode 100644 index 0000000000000000000000000000000000000000..18e9335673f6d46ea8fd0a03a791516203eb6f4c --- /dev/null +++ b/tinyvla2/assets/so_arm/assets/sts3215_03a_no_horn_v1.stl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:75ef3781b752e4065891aea855e34dc161a38a549549cd0970cedd07eae6f887 +size 865884 diff --git a/tinyvla2/assets/so_arm/assets/sts3215_03a_v1.stl b/tinyvla2/assets/so_arm/assets/sts3215_03a_v1.stl new file mode 100644 index 0000000000000000000000000000000000000000..a14c57b9033b82f1daa38f45e7e3c91343702df4 --- /dev/null +++ b/tinyvla2/assets/so_arm/assets/sts3215_03a_v1.stl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a37c871fb502483ab96c256baf457d36f2e97afc9205313d9c5ab275ef941cd0 +size 954084 diff --git a/tinyvla2/assets/so_arm/assets/under_arm_so101_v1.stl b/tinyvla2/assets/so_arm/assets/under_arm_so101_v1.stl new file mode 100644 index 0000000000000000000000000000000000000000..47b611ef939e452f791ae749756f717317922cfd --- /dev/null +++ b/tinyvla2/assets/so_arm/assets/under_arm_so101_v1.stl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d01d1f2de365651dcad9d6669e94ff87ff7652b5bb2d10752a66a456a86dbc71 +size 1975884 diff --git a/tinyvla2/assets/so_arm/assets/upper_arm_so101_v1.stl b/tinyvla2/assets/so_arm/assets/upper_arm_so101_v1.stl new file mode 100644 index 0000000000000000000000000000000000000000..8832740f9540065e6006907a9a826b01f96cd122 --- /dev/null +++ b/tinyvla2/assets/so_arm/assets/upper_arm_so101_v1.stl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:475056e03a17e71919b82fd88ab9a0b898ab50164f2a7943652a6b2941bb2d4f +size 1303484 diff --git a/tinyvla2/assets/so_arm/assets/waveshare_mounting_plate_so101_v2.stl b/tinyvla2/assets/so_arm/assets/waveshare_mounting_plate_so101_v2.stl new file mode 100644 index 0000000000000000000000000000000000000000..4dba1f22b01ab9afef2466872fcb3e1efca89e38 Binary files /dev/null and b/tinyvla2/assets/so_arm/assets/waveshare_mounting_plate_so101_v2.stl differ diff --git a/tinyvla2/assets/so_arm/assets/wrist_roll_follower_so101_v1.stl b/tinyvla2/assets/so_arm/assets/wrist_roll_follower_so101_v1.stl new file mode 100644 index 0000000000000000000000000000000000000000..9a5fa8fe2d7d8e59cd4a30d4dba0ca337513ab4a --- /dev/null +++ b/tinyvla2/assets/so_arm/assets/wrist_roll_follower_so101_v1.stl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:4b17b410a12d64ec39554abc3e8054d8a97384b2dc4a8d95a5ecb2a93670f5f4 +size 1439884 diff --git a/tinyvla2/assets/so_arm/assets/wrist_roll_pitch_so101_v2.stl b/tinyvla2/assets/so_arm/assets/wrist_roll_pitch_so101_v2.stl new file mode 100644 index 0000000000000000000000000000000000000000..2f531712f88ec01d09824ee8e27c791a4616516f --- /dev/null +++ b/tinyvla2/assets/so_arm/assets/wrist_roll_pitch_so101_v2.stl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:6c7ec5525b4d8b9e397a30ab4bb0037156a5d5f38a4adf2c7d943d6c56eda5ae +size 2699784 diff --git a/tinyvla2/assets/so_arm/so101_new_calib.urdf b/tinyvla2/assets/so_arm/so101_new_calib.urdf new file mode 100644 index 0000000000000000000000000000000000000000..9552a231d8b23bed68ec15779eba620c5d875ec4 --- /dev/null +++ b/tinyvla2/assets/so_arm/so101_new_calib.urdf @@ -0,0 +1,453 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + transmission_interface/SimpleTransmission + + hardware_interface/PositionJointInterface + + + hardware_interface/PositionJointInterface + 1 + + + + + + + + + + + + + + transmission_interface/SimpleTransmission + + hardware_interface/PositionJointInterface + + + hardware_interface/PositionJointInterface + 1 + + + + + + + + + + + + + + transmission_interface/SimpleTransmission + + hardware_interface/PositionJointInterface + + + hardware_interface/PositionJointInterface + 1 + + + + + + + + + + + + + + + transmission_interface/SimpleTransmission + + hardware_interface/PositionJointInterface + + + hardware_interface/PositionJointInterface + 1 + + + + + + + + + + + + + + transmission_interface/SimpleTransmission + + hardware_interface/PositionJointInterface + + + hardware_interface/PositionJointInterface + 1 + + + + + + + + + + + + + + transmission_interface/SimpleTransmission + + hardware_interface/PositionJointInterface + + + hardware_interface/PositionJointInterface + 1 + + + + \ No newline at end of file diff --git a/tinyvla2/assets/so_arm/so101_nomesh.urdf b/tinyvla2/assets/so_arm/so101_nomesh.urdf new file mode 100644 index 0000000000000000000000000000000000000000..235283af9dd15e896ca828b391c307cffc74a319 --- /dev/null +++ b/tinyvla2/assets/so_arm/so101_nomesh.urdf @@ -0,0 +1,453 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + transmission_interface/SimpleTransmission + + hardware_interface/PositionJointInterface + + + hardware_interface/PositionJointInterface + 1 + + + + + + + + + + + + + + transmission_interface/SimpleTransmission + + hardware_interface/PositionJointInterface + + + hardware_interface/PositionJointInterface + 1 + + + + + + + + + + + + + + transmission_interface/SimpleTransmission + + hardware_interface/PositionJointInterface + + + hardware_interface/PositionJointInterface + 1 + + + + + + + + + + + + + + + transmission_interface/SimpleTransmission + + hardware_interface/PositionJointInterface + + + hardware_interface/PositionJointInterface + 1 + + + + + + + + + + + + + + transmission_interface/SimpleTransmission + + hardware_interface/PositionJointInterface + + + hardware_interface/PositionJointInterface + 1 + + + + + + + + + + + + + + transmission_interface/SimpleTransmission + + hardware_interface/PositionJointInterface + + + hardware_interface/PositionJointInterface + 1 + + + + \ No newline at end of file diff --git a/tinyvla2/configs/ablation_nolatent.yaml b/tinyvla2/configs/ablation_nolatent.yaml new file mode 100644 index 0000000000000000000000000000000000000000..6516fdaa74765a5c3927946c292ad25e7c025b04 --- /dev/null +++ b/tinyvla2/configs/ablation_nolatent.yaml @@ -0,0 +1,25 @@ +# Stage-2 gate baseline: same mixture, 5k steps, nolatent ablation +output_dir: /home/alexw/tinyvla/outputs/ablation_nolatent +wandb: null +datasets: + - root_glob: /home/alexw/tinyvla_data/so101_v3/* + weight: 0.65 + - repo_id: nvidia/BridgeData2_LeRobot_v3 + root: /home/alexw/tinyvla_data/bridge_v3_10k + revision: main + episodes: 10000 + weight: 0.35 +policy: + chunk_size: 50 + n_action_steps: 50 + image_size: 256 + use_semantic_latent: false +batch_size: 64 +grad_accum: 2 +num_workers: 12 +lr: 1.0e-4 +warmup_steps: 500 +steps: 5000 +log_freq: 50 +save_freq: 2500 +mixed_precision: bf16 diff --git a/tinyvla2/configs/ablation_nospatial.yaml b/tinyvla2/configs/ablation_nospatial.yaml new file mode 100644 index 0000000000000000000000000000000000000000..e2bcc0ea6123d64cac7ec1c555d953e9b54b0fc1 --- /dev/null +++ b/tinyvla2/configs/ablation_nospatial.yaml @@ -0,0 +1,25 @@ +# Stage-2 gate baseline: same mixture, 5k steps, nospatial ablation +output_dir: /home/alexw/tinyvla/outputs/ablation_nospatial +wandb: null +datasets: + - root_glob: /home/alexw/tinyvla_data/so101_v3/* + weight: 0.65 + - repo_id: nvidia/BridgeData2_LeRobot_v3 + root: /home/alexw/tinyvla_data/bridge_v3_10k + revision: main + episodes: 10000 + weight: 0.35 +policy: + chunk_size: 50 + n_action_steps: 50 + image_size: 256 + use_spatial_tokens: false +batch_size: 64 +grad_accum: 2 +num_workers: 12 +lr: 1.0e-4 +warmup_steps: 500 +steps: 5000 +log_freq: 50 +save_freq: 2500 +mixed_precision: bf16 diff --git a/tinyvla2/configs/libero_ft.yaml b/tinyvla2/configs/libero_ft.yaml new file mode 100644 index 0000000000000000000000000000000000000000..da4208f1273a4ad9c3388ece78ea7438472fef08 --- /dev/null +++ b/tinyvla2/configs/libero_ft.yaml @@ -0,0 +1,31 @@ +# LIBERO fine-tune from the Stage-2 pretrained checkpoint (XS-VLA/SmolVLA protocol). +# LIBERO = new embodiment (id 2): Franka 7-DoF EE @10fps, two 256^2 cameras. + +output_dir: /home/alexw/tinyvla/outputs/libero_ft +wandb: null + +datasets: + - repo_id: HuggingFaceVLA/libero + weight: 1.0 + embodiment_id: 2 + +policy: + chunk_size: 50 + n_action_steps: 50 + image_size: 256 + freeze_lm: true + freeze_vision_encoder: true + +resume_from: /home/alexw/tinyvla/outputs/stage2_mixture/final +resume_step: 0 + +batch_size: 64 +grad_accum: 1 +num_workers: 12 +lr: 5.0e-5 +warmup_steps: 500 +steps: 15000 +grad_clip: 10.0 +log_freq: 50 +save_freq: 2500 +mixed_precision: bf16 diff --git a/tinyvla2/configs/libero_ft2.yaml b/tinyvla2/configs/libero_ft2.yaml new file mode 100644 index 0000000000000000000000000000000000000000..6eab4bdded3eb6fd9e90a984561afa326f16d26b --- /dev/null +++ b/tinyvla2/configs/libero_ft2.yaml @@ -0,0 +1,33 @@ +# LIBERO fine-tune v2: unfreeze the vision tower (sim renders are OOD for the +# real-world-pretrained encoder), longer schedule. + +output_dir: /home/alexw/tinyvla/outputs/libero_ft2 +wandb: null + +datasets: + - repo_id: HuggingFaceVLA/libero + weight: 1.0 + embodiment_id: 2 + +policy: + chunk_size: 50 + n_action_steps: 50 + image_size: 256 + freeze_lm: true + freeze_vision_encoder: false + +backbone_lr_mult: 0.1 + +resume_from: /home/alexw/tinyvla/outputs/stage2_mixture/final +resume_step: 0 + +batch_size: 48 +grad_accum: 1 +num_workers: 12 +lr: 1.0e-4 +warmup_steps: 1000 +steps: 30000 +grad_clip: 10.0 +log_freq: 50 +save_freq: 5000 +mixed_precision: bf16 diff --git a/tinyvla2/configs/libero_ft3.yaml b/tinyvla2/configs/libero_ft3.yaml new file mode 100644 index 0000000000000000000000000000000000000000..d86dcf85332e0e0cb2b0e3d4434fa15fb978a5f5 --- /dev/null +++ b/tinyvla2/configs/libero_ft3.yaml @@ -0,0 +1,33 @@ +# LIBERO fine-tune v2: unfreeze the vision tower (sim renders are OOD for the +# real-world-pretrained encoder), longer schedule. + +output_dir: /home/alexw/tinyvla/outputs/libero_ft3 +wandb: null + +datasets: + - repo_id: HuggingFaceVLA/libero + weight: 1.0 + embodiment_id: 2 + +policy: + chunk_size: 50 + n_action_steps: 50 + image_size: 256 + freeze_lm: true + freeze_vision_encoder: false + +backbone_lr_mult: 0.1 + +resume_from: /home/alexw/tinyvla/outputs/stage3_distill/final +resume_step: 0 + +batch_size: 48 +grad_accum: 1 +num_workers: 12 +lr: 1.0e-4 +warmup_steps: 1000 +steps: 30000 +grad_clip: 10.0 +log_freq: 50 +save_freq: 5000 +mixed_precision: bf16 diff --git a/tinyvla2/configs/libero_ft4.yaml b/tinyvla2/configs/libero_ft4.yaml new file mode 100644 index 0000000000000000000000000000000000000000..a719f9ef4c5abec2f9d3d699241c54f4064d0fe0 --- /dev/null +++ b/tinyvla2/configs/libero_ft4.yaml @@ -0,0 +1,33 @@ +# LIBERO fine-tune v4: unfreeze BOTH vision tower and language stack +# (12 truncated layers; vocab embeddings stay frozen) at 0.1x lr, longer schedule. + +output_dir: /home/alexw/tinyvla/outputs/libero_ft4 +wandb: null + +datasets: + - repo_id: HuggingFaceVLA/libero + weight: 1.0 + embodiment_id: 2 + +policy: + chunk_size: 50 + n_action_steps: 50 + image_size: 256 + freeze_lm: false + freeze_vision_encoder: false + +backbone_lr_mult: 0.1 + +resume_from: /home/alexw/tinyvla/outputs/stage3_distill/final +resume_step: 0 + +batch_size: 48 +grad_accum: 1 +num_workers: 12 +lr: 1.0e-4 +warmup_steps: 1000 +steps: 40000 +grad_clip: 10.0 +log_freq: 50 +save_freq: 5000 +mixed_precision: bf16 diff --git a/tinyvla2/configs/morphology/descriptors.yaml b/tinyvla2/configs/morphology/descriptors.yaml new file mode 100644 index 0000000000000000000000000000000000000000..52a2a230ec935c31cbe32547d190746e9afca536 --- /dev/null +++ b/tinyvla2/configs/morphology/descriptors.yaml @@ -0,0 +1,277 @@ +# Per-embodiment morphology descriptors (variant C). +# Fields hand-normalized to O(1); order MUST match MORPH_FIELDS in +# src/tinyvla/modules/embodiment.py. Values: raw → normalized applied at load. +# arm_dof/10, reach_m*2, gripper_width_m*10, num_cameras/3, is_mobile ±1, +# control_hz/30, joint_lo_mean/pi, joint_hi_mean/pi, workspace_xyz*2, +# payload_kg/5, ee_type_parallel ±1, ee_type_multi ±1, base_holonomic ±1, reserved. + +so101: # SO-100/101 5-DOF + gripper + arm_dof: 5 + reach_m: 0.35 + gripper_width_m: 0.03 + num_cameras: 2 + is_mobile: -1 + control_hz: 30 + joint_lo_mean: -1.9 + joint_hi_mean: 1.9 + workspace_x: 0.30 + workspace_y: 0.30 + workspace_z: 0.25 + payload_kg: 0.25 + ee_type_parallel: 1 + ee_type_multi: -1 + base_holonomic: -1 + reserved: 0 + +bridge: # WidowX 250 6-DOF + arm_dof: 6 + reach_m: 0.60 + gripper_width_m: 0.06 + num_cameras: 1 + is_mobile: -1 + control_hz: 5 + joint_lo_mean: -2.6 + joint_hi_mean: 2.6 + workspace_x: 0.40 + workspace_y: 0.40 + workspace_z: 0.35 + payload_kg: 0.25 + ee_type_parallel: 1 + ee_type_multi: -1 + base_holonomic: -1 + reserved: 0 + +rt1: # Google robot (mobile manipulator, 7-DOF arm) + arm_dof: 7 + reach_m: 0.90 + gripper_width_m: 0.08 + num_cameras: 1 + is_mobile: 1 + control_hz: 3 + joint_lo_mean: -2.9 + joint_hi_mean: 2.9 + workspace_x: 0.70 + workspace_y: 0.70 + workspace_z: 0.60 + payload_kg: 1.0 + ee_type_parallel: 1 + ee_type_multi: -1 + base_holonomic: 1 + reserved: 0 + +droid: # Franka Panda 7-DOF + arm_dof: 7 + reach_m: 0.85 + gripper_width_m: 0.08 + num_cameras: 2 + is_mobile: -1 + control_hz: 15 + joint_lo_mean: -2.8 + joint_hi_mean: 2.8 + workspace_x: 0.60 + workspace_y: 0.60 + workspace_z: 0.55 + payload_kg: 3.0 + ee_type_parallel: 1 + ee_type_multi: -1 + base_holonomic: -1 + reserved: 0 + +libero: # Franka Panda (sim) — same morphology as droid + arm_dof: 7 + reach_m: 0.85 + gripper_width_m: 0.08 + num_cameras: 2 + is_mobile: -1 + control_hz: 20 + joint_lo_mean: -2.8 + joint_hi_mean: 2.8 + workspace_x: 0.60 + workspace_y: 0.60 + workspace_z: 0.55 + payload_kg: 3.0 + ee_type_parallel: 1 + ee_type_multi: -1 + base_holonomic: -1 + reserved: 0 + +# --- held-out (never in pretraining; descriptors ready for zero-shot variant C) --- +jaco: # Kinova Jaco 6-DOF, 3-finger gripper + arm_dof: 6 + reach_m: 0.90 + gripper_width_m: 0.10 + num_cameras: 1 + is_mobile: -1 + control_hz: 10 + joint_lo_mean: -3.1 + joint_hi_mean: 3.1 + workspace_x: 0.60 + workspace_y: 0.60 + workspace_z: 0.55 + payload_kg: 1.6 + ee_type_parallel: -1 + ee_type_multi: 1 + base_holonomic: -1 + reserved: 0 + +fanuc: # Fanuc Mate 6-DOF industrial + arm_dof: 6 + reach_m: 0.70 + gripper_width_m: 0.05 + num_cameras: 1 + is_mobile: -1 + control_hz: 10 + joint_lo_mean: -3.1 + joint_hi_mean: 3.1 + workspace_x: 0.50 + workspace_y: 0.50 + workspace_z: 0.45 + payload_kg: 4.0 + ee_type_parallel: 1 + ee_type_multi: -1 + base_holonomic: -1 + reserved: 0 + +lekiwi: # LeKiwi — SO-100 (5-DOF) arm on a 3-wheel holonomic kiwi base (mobile manipulator) + arm_dof: 5 + reach_m: 0.35 + gripper_width_m: 0.03 + num_cameras: 2 + is_mobile: 1 + control_hz: 30 + joint_lo_mean: -1.9 + joint_hi_mean: 1.9 + workspace_x: 1.5 + workspace_y: 1.5 + workspace_z: 0.30 + payload_kg: 0.25 + ee_type_parallel: 1 + ee_type_multi: -1 + base_holonomic: 1 + reserved: 0 + +ur5: # Universal Robots UR5e — 6-DOF industrial arm, fixed base + arm_dof: 6 + reach_m: 0.85 + gripper_width_m: 0.085 + num_cameras: 1 + is_mobile: -1 + control_hz: 5 + joint_lo_mean: -3.1 + joint_hi_mean: 3.1 + workspace_x: 0.70 + workspace_y: 0.70 + workspace_z: 0.60 + payload_kg: 5.0 + ee_type_parallel: 1 + ee_type_multi: -1 + base_holonomic: -1 + reserved: 0 + +xarm: # UFactory xArm — 6-DOF arm, fixed base + arm_dof: 6 + reach_m: 0.70 + gripper_width_m: 0.08 + num_cameras: 1 + is_mobile: -1 + control_hz: 5 + joint_lo_mean: -3.1 + joint_hi_mean: 3.1 + workspace_x: 0.55 + workspace_y: 0.55 + workspace_z: 0.50 + payload_kg: 3.5 + ee_type_parallel: 1 + ee_type_multi: -1 + base_holonomic: -1 + reserved: 0 + +dlr_edan: # DLR EDAN — 7-DOF assistive arm on a wheelchair (mobile assistive) + arm_dof: 7 + reach_m: 0.80 + gripper_width_m: 0.07 + num_cameras: 1 + is_mobile: 1 + control_hz: 5 + joint_lo_mean: -2.9 + joint_hi_mean: 2.9 + workspace_x: 0.65 + workspace_y: 0.65 + workspace_z: 0.55 + payload_kg: 2.0 + ee_type_parallel: 1 + ee_type_multi: -1 + base_holonomic: -1 + reserved: 0 + +tidybot: # holonomic wheeled mobile manipulator (Kinova arm on omni base), whole-body EE control + arm_dof: 7 + reach_m: 0.90 + gripper_width_m: 0.08 + num_cameras: 2 + is_mobile: 1 + control_hz: 10 + joint_lo_mean: -2.9 + joint_hi_mean: 2.9 + workspace_x: 2.0 + workspace_y: 2.0 + workspace_z: 1.0 + payload_kg: 1.0 + ee_type_parallel: 1 + ee_type_multi: -1 + base_holonomic: 1 + reserved: 0 + +uav: # quadrotor drone — free-flying 6-DOF body, no arm, no gripper (max morphology gap) + arm_dof: 0 + reach_m: 0.0 + gripper_width_m: 0.0 + num_cameras: 6 + is_mobile: 1 + control_hz: 30 + joint_lo_mean: 0.0 + joint_hi_mean: 0.0 + workspace_x: 5.0 + workspace_y: 5.0 + workspace_z: 5.0 + payload_kg: 0.5 + ee_type_parallel: -1 + ee_type_multi: -1 + base_holonomic: 1 + reserved: 1 + +hello_stretch: # Hello Robot Stretch — mobile base + prismatic lift + arm_dof: 4 + reach_m: 0.52 + gripper_width_m: 0.08 + num_cameras: 1 + is_mobile: 1 + control_hz: 10 + joint_lo_mean: -1.5 + joint_hi_mean: 1.5 + workspace_x: 0.50 + workspace_y: 0.30 + workspace_z: 1.10 + payload_kg: 1.5 + ee_type_parallel: 1 + ee_type_multi: -1 + base_holonomic: -1 + reserved: 0 + +panda_omron: # RoboCasa PandaOmron — Franka 7-DOF arm on a holonomic omni mobile base (kitchen scenes) + arm_dof: 7 + reach_m: 0.85 + gripper_width_m: 0.08 + num_cameras: 3 + is_mobile: 1 + control_hz: 20 + joint_lo_mean: -2.8 + joint_hi_mean: 2.8 + workspace_x: 2.0 + workspace_y: 2.0 + workspace_z: 1.2 + payload_kg: 3.0 + ee_type_parallel: 1 + ee_type_multi: -1 + base_holonomic: 1 + reserved: 0 diff --git a/tinyvla2/configs/morphology/robot_prompts.yaml b/tinyvla2/configs/morphology/robot_prompts.yaml new file mode 100644 index 0000000000000000000000000000000000000000..2094f6128dace5cbb5a772fb4623cf7e3054bf2a --- /dev/null +++ b/tinyvla2/configs/morphology/robot_prompts.yaml @@ -0,0 +1,13 @@ +# Natural-language robot descriptions prepended to the task instruction (Qwen slow path). +# Complements the numeric morphology descriptor (C scheme) with a text channel. +so101: "Robot: SO-101, a 5-DOF tabletop arm with a parallel-jaw gripper on a fixed base." +bridge: "Robot: WidowX 250, a 6-DOF tabletop arm with a parallel gripper on a fixed base." +rt1: "Robot: Google Robot, a 7-DOF arm mounted on a wheeled mobile base." +ur5: "Robot: UR5e, a 6-DOF industrial arm with a parallel gripper on a fixed base." +xarm: "Robot: xArm, a 6-DOF arm with a parallel gripper on a fixed base." +dlr_edan: "Robot: DLR EDAN, a 7-DOF assistive arm mounted on a wheelchair mobile base." +tidybot: "Robot: TidyBot, a 7-DOF arm on a holonomic omnidirectional wheeled base; whole-body end-effector control." +hello_stretch: "Robot: Hello Stretch, an arm on a telescoping prismatic lift atop a mobile base." +jaco: "Robot: Kinova Jaco, a 6-DOF arm with a three-finger gripper on a fixed base." +lekiwi: "Robot: LeKiwi, a 5-DOF SO-100 arm mounted on a three-wheel holonomic mobile base." +panda_omron: "Robot: RoboCasa PandaOmron, a 7-DOF Franka arm on a holonomic omnidirectional mobile base, operating in kitchen scenes." diff --git a/tinyvla2/configs/stage2_mixture.yaml b/tinyvla2/configs/stage2_mixture.yaml new file mode 100644 index 0000000000000000000000000000000000000000..7f0498bee6ba0cb7ad7b15aa825dc903a3885a4b --- /dev/null +++ b/tinyvla2/configs/stage2_mixture.yaml @@ -0,0 +1,37 @@ +# Stage 2: small-scale mixture pretraining on one A6000. +# SO-100/101 community slice (converted v3.0) 0.65 + Bridge V2 0.35. +# Fractal dropped for MVP (v2.0, unconvertible in lerobot 0.6.1). + +output_dir: /home/alexw/tinyvla/outputs/stage2_mixture +wandb: null # set a project name to enable + +datasets: + - root_glob: /home/alexw/tinyvla_data/so101_v3/* + weight: 0.65 + - repo_id: nvidia/BridgeData2_LeRobot_v3 + root: /home/alexw/tinyvla_data/bridge_v3_10k + revision: main + episodes: 10000 + weight: 0.35 + +policy: + chunk_size: 50 + n_action_steps: 50 + image_size: 256 + freeze_lm: true + freeze_vision_encoder: true + +batch_size: 64 +grad_accum: 2 +num_workers: 14 +resume_from: /home/alexw/tinyvla/outputs/stage2_mixture/step_2000 +resume_step: 2000 +lr: 1.0e-4 +warmup_steps: 1000 +steps: 40000 +grad_clip: 10.0 +log_freq: 50 +save_freq: 2000 +mixed_precision: bf16 +staleness_start_step: 20000 +staleness_prob: 0.5 diff --git a/tinyvla2/configs/stage3_distill.yaml b/tinyvla2/configs/stage3_distill.yaml new file mode 100644 index 0000000000000000000000000000000000000000..362662289a994c0c6d1d948985af39c24d8b64f8 --- /dev/null +++ b/tinyvla2/configs/stage3_distill.yaml @@ -0,0 +1,40 @@ +# Stage 3: continue SO101+Bridge pretraining with the spatial-distillation +# aux CE loss (teacher labels from Qwen3.5-4B) and staleness aug from the start. + +output_dir: /home/alexw/tinyvla/outputs/stage3_distill +wandb: null + +spatial_labels_dir: /home/alexw/tinyvla_data/spatial_labels + +datasets: + - root_glob: /home/alexw/tinyvla_data/so101_v3/* + weight: 0.65 + - repo_id: nvidia/BridgeData2_LeRobot_v3 + root: /home/alexw/tinyvla_data/bridge_v3_10k + revision: main + episodes: 10000 + weight: 0.35 + +policy: + chunk_size: 50 + n_action_steps: 50 + image_size: 256 + freeze_lm: true + freeze_vision_encoder: true + spatial_loss_weight: 0.15 + +resume_from: /home/alexw/tinyvla/outputs/stage2_mixture/final +resume_step: 0 + +batch_size: 64 +grad_accum: 2 +num_workers: 14 +lr: 5.0e-5 +warmup_steps: 500 +steps: 20000 +grad_clip: 10.0 +log_freq: 50 +save_freq: 2500 +mixed_precision: bf16 +staleness_start_step: 0 +staleness_prob: 0.5 diff --git a/tinyvla2/configs/tv2_A.yaml b/tinyvla2/configs/tv2_A.yaml new file mode 100644 index 0000000000000000000000000000000000000000..9ca4710a7501ccaed9a87ca64d68bada12bc880e --- /dev/null +++ b/tinyvla2/configs/tv2_A.yaml @@ -0,0 +1,42 @@ +# TinyVLA-2 variant A: conditioning=id, action_space=native +# Identical data/weights/seed across A/B/C (matched-pairs design). +output_dir: /home/alexw/tinyvla/outputs/tv2_A_id_native +wandb: null +seed: 42 +morphology_descriptors: /home/alexw/tinyvla/configs/morphology/descriptors.yaml + +datasets: + - root_glob: /home/alexw/tinyvla_data/so101_v3/* + weight: 0.40 + morph_key: so101 + - repo_id: MolmoAct2-Bridge-Dataset + root: /home/alexw/tinyvla_data/MolmoAct2-Bridge-Dataset + episodes: 20000 + weight: 0.30 + morph_key: bridge + - repo_id: MolmoAct2-RT-1-Dataset + root: /home/alexw/tinyvla_data/MolmoAct2-RT-1-Dataset + episodes: 20000 + weight: 0.30 + morph_key: rt1 + +policy: + chunk_size: 50 + n_action_steps: 50 + image_size: 256 + freeze_lm: true + freeze_vision_encoder: true + num_embodiments: 16 + conditioning: id + action_space: native + +batch_size: 64 +grad_accum: 2 +num_workers: 12 +lr: 1.0e-4 +warmup_steps: 1000 +steps: 40000 +grad_clip: 10.0 +log_freq: 50 +save_freq: 5000 +mixed_precision: bf16 diff --git a/tinyvla2/configs/tv2_B.yaml b/tinyvla2/configs/tv2_B.yaml new file mode 100644 index 0000000000000000000000000000000000000000..d7ad50d35ce852f16ce67017ea4f1c26505a929c --- /dev/null +++ b/tinyvla2/configs/tv2_B.yaml @@ -0,0 +1,42 @@ +# TinyVLA-2 variant B: conditioning=id, action_space=canonical +# Identical data/weights/seed across A/B/C (matched-pairs design). +output_dir: /home/alexw/tinyvla/outputs/tv2_B_id_canon +wandb: null +seed: 42 +morphology_descriptors: /home/alexw/tinyvla/configs/morphology/descriptors.yaml + +datasets: + - root_glob: /home/alexw/tinyvla_data/so101_v3/* + weight: 0.40 + morph_key: so101 + - repo_id: MolmoAct2-Bridge-Dataset + root: /home/alexw/tinyvla_data/MolmoAct2-Bridge-Dataset + episodes: 20000 + weight: 0.30 + morph_key: bridge + - repo_id: MolmoAct2-RT-1-Dataset + root: /home/alexw/tinyvla_data/MolmoAct2-RT-1-Dataset + episodes: 20000 + weight: 0.30 + morph_key: rt1 + +policy: + chunk_size: 50 + n_action_steps: 50 + image_size: 256 + freeze_lm: true + freeze_vision_encoder: true + num_embodiments: 16 + conditioning: id + action_space: canonical + +batch_size: 64 +grad_accum: 2 +num_workers: 12 +lr: 1.0e-4 +warmup_steps: 1000 +steps: 40000 +grad_clip: 10.0 +log_freq: 50 +save_freq: 5000 +mixed_precision: bf16 diff --git a/tinyvla2/configs/tv2_C.yaml b/tinyvla2/configs/tv2_C.yaml new file mode 100644 index 0000000000000000000000000000000000000000..bf6887877c937c5aa839e946db3f84be15acff65 --- /dev/null +++ b/tinyvla2/configs/tv2_C.yaml @@ -0,0 +1,42 @@ +# TinyVLA-2 variant C: conditioning=morph, action_space=canonical +# Identical data/weights/seed across A/B/C (matched-pairs design). +output_dir: /home/alexw/tinyvla/outputs/tv2_C_morph_canon +wandb: null +seed: 42 +morphology_descriptors: /home/alexw/tinyvla/configs/morphology/descriptors.yaml + +datasets: + - root_glob: /home/alexw/tinyvla_data/so101_v3/* + weight: 0.40 + morph_key: so101 + - repo_id: MolmoAct2-Bridge-Dataset + root: /home/alexw/tinyvla_data/MolmoAct2-Bridge-Dataset + episodes: 20000 + weight: 0.30 + morph_key: bridge + - repo_id: MolmoAct2-RT-1-Dataset + root: /home/alexw/tinyvla_data/MolmoAct2-RT-1-Dataset + episodes: 20000 + weight: 0.30 + morph_key: rt1 + +policy: + chunk_size: 50 + n_action_steps: 50 + image_size: 256 + freeze_lm: true + freeze_vision_encoder: true + num_embodiments: 16 + conditioning: morph + action_space: canonical + +batch_size: 64 +grad_accum: 2 +num_workers: 12 +lr: 1.0e-4 +warmup_steps: 1000 +steps: 40000 +grad_clip: 10.0 +log_freq: 50 +save_freq: 5000 +mixed_precision: bf16 diff --git a/tinyvla2/configs/tv2_C_diverse.yaml b/tinyvla2/configs/tv2_C_diverse.yaml new file mode 100644 index 0000000000000000000000000000000000000000..1a9f07411a62321eb7eded69b7de0a6906175f71 --- /dev/null +++ b/tinyvla2/configs/tv2_C_diverse.yaml @@ -0,0 +1,67 @@ +# Continue-train variant C on an EXPANDED, morphology-diverse mixture to fix the +# OOD descriptor extrapolation (original C's morph-encoder saw only 3 arm-ish +# embodiments → failed zero-shot on mobile TidyBot). Now 8 embodiments spanning +# fixed arms (SO101, WidowX, UR5, xArm) + mobile (Google-robot, dlr_edan wheelchair, +# TidyBot holonomic, Stretch). jaco held out entirely for the transfer test. +output_dir: /home/alexw/tinyvla/outputs/tv2_C_diverse +wandb: null +seed: 42 +morphology_descriptors: /home/alexw/tinyvla/configs/morphology/descriptors.yaml +resume_from: /home/alexw/tinyvla/outputs/tv2_C_morph_canon/final +resume_step: 0 + +datasets: + - root_glob: /home/alexw/tinyvla_data/so101_v3/* + weight: 0.25 + morph_key: so101 + - repo_id: MolmoAct2-Bridge-Dataset + root: /home/alexw/tinyvla_data/MolmoAct2-Bridge-Dataset + episodes: 15000 + weight: 0.18 + morph_key: bridge + - repo_id: MolmoAct2-RT-1-Dataset + root: /home/alexw/tinyvla_data/MolmoAct2-RT-1-Dataset + episodes: 15000 + weight: 0.18 + morph_key: rt1 + - repo_id: div_ur5 + root: /home/alexw/tinyvla_data/diverse/div_ur5 + weight: 0.10 + morph_key: ur5 + - repo_id: div_xarm + root: /home/alexw/tinyvla_data/diverse/div_xarm + weight: 0.09 + morph_key: xarm + - repo_id: div_dlr_edan + root: /home/alexw/tinyvla_data/diverse/div_dlr_edan + weight: 0.05 + morph_key: dlr_edan + - repo_id: heldout_tidybot + root: /home/alexw/tinyvla_data/heldout/heldout_tidybot + weight: 0.08 + morph_key: tidybot + - repo_id: heldout_cmu_stretch + root: /home/alexw/tinyvla_data/heldout/heldout_cmu_stretch + weight: 0.07 + morph_key: hello_stretch + +policy: + chunk_size: 50 + n_action_steps: 50 + image_size: 256 + freeze_lm: true + freeze_vision_encoder: true + num_embodiments: 16 + conditioning: morph + action_space: canonical + +batch_size: 64 +grad_accum: 2 +num_workers: 12 +lr: 7.0e-5 +warmup_steps: 500 +steps: 25000 +grad_clip: 10.0 +log_freq: 50 +save_freq: 5000 +mixed_precision: bf16 diff --git a/tinyvla2/configs/tv2_C_diverse_slow.yaml b/tinyvla2/configs/tv2_C_diverse_slow.yaml new file mode 100644 index 0000000000000000000000000000000000000000..b144630a7e980e40924ee8cfe7ef882979337a28 --- /dev/null +++ b/tinyvla2/configs/tv2_C_diverse_slow.yaml @@ -0,0 +1,68 @@ +# Continue-train variant C on an EXPANDED, morphology-diverse mixture to fix the +# OOD descriptor extrapolation (original C's morph-encoder saw only 3 arm-ish +# embodiments → failed zero-shot on mobile TidyBot). Now 8 embodiments spanning +# fixed arms (SO101, WidowX, UR5, xArm) + mobile (Google-robot, dlr_edan wheelchair, +# TidyBot holonomic, Stretch). jaco held out entirely for the transfer test. +output_dir: /home/alexw/tinyvla/outputs/tv2_C_diverse_slow +wandb: null +seed: 42 +morphology_descriptors: /home/alexw/tinyvla/configs/morphology/descriptors.yaml +resume_from: /home/alexw/tinyvla/outputs/tv2_C_diverse/final +resume_step: 0 + +datasets: + - root_glob: /home/alexw/tinyvla_data/so101_v3/* + weight: 0.25 + morph_key: so101 + - repo_id: MolmoAct2-Bridge-Dataset + root: /home/alexw/tinyvla_data/MolmoAct2-Bridge-Dataset + episodes: 15000 + weight: 0.18 + morph_key: bridge + - repo_id: MolmoAct2-RT-1-Dataset + root: /home/alexw/tinyvla_data/MolmoAct2-RT-1-Dataset + episodes: 15000 + weight: 0.18 + morph_key: rt1 + - repo_id: div_ur5 + root: /home/alexw/tinyvla_data/diverse/div_ur5 + weight: 0.10 + morph_key: ur5 + - repo_id: div_xarm + root: /home/alexw/tinyvla_data/diverse/div_xarm + weight: 0.09 + morph_key: xarm + - repo_id: div_dlr_edan + root: /home/alexw/tinyvla_data/diverse/div_dlr_edan + weight: 0.05 + morph_key: dlr_edan + - repo_id: heldout_tidybot + root: /home/alexw/tinyvla_data/heldout/heldout_tidybot + weight: 0.08 + morph_key: tidybot + - repo_id: heldout_cmu_stretch + root: /home/alexw/tinyvla_data/heldout/heldout_cmu_stretch + weight: 0.07 + morph_key: hello_stretch + +policy: + chunk_size: 50 + n_action_steps: 50 + image_size: 256 + freeze_lm: true + freeze_vision_encoder: true + num_embodiments: 16 + conditioning: morph + morph_to_slow: true + action_space: canonical + +batch_size: 64 +grad_accum: 2 +num_workers: 12 +lr: 7.0e-5 +warmup_steps: 500 +steps: 25000 +grad_clip: 10.0 +log_freq: 50 +save_freq: 5000 +mixed_precision: bf16 diff --git a/tinyvla2/configs/tv2_C_incontext.yaml b/tinyvla2/configs/tv2_C_incontext.yaml new file mode 100644 index 0000000000000000000000000000000000000000..2b33fcaa5cb2254b52ce6656969fc9a34e9d1ccc --- /dev/null +++ b/tinyvla2/configs/tv2_C_incontext.yaml @@ -0,0 +1,72 @@ +# Variant: in-context few-shot DEMO conditioning — a few (obs, action) exemplar +# pairs from the target robot's own data, encoded via the shared spatial-aware +# vision encoder + a small MLP, WITHOUT any gradient update to adapt to that robot. +# Composed with the text-via-Qwen morphology channel (both active). Resumes from +# tv2_C_qwen_morph/final (LM already unfrozen, morph_out_proj already trained). +output_dir: /home/alexw/tinyvla/outputs/tv2_C_incontext +wandb: null +seed: 42 +robot_prompts: /home/alexw/tinyvla/configs/morphology/robot_prompts.yaml +resume_from: /home/alexw/tinyvla/outputs/tv2_C_qwen_morph/final +resume_step: 0 + +datasets: + - root_glob: /home/alexw/tinyvla_data/so101_v3/* + weight: 0.25 + morph_key: so101 + - repo_id: MolmoAct2-Bridge-Dataset + root: /home/alexw/tinyvla_data/MolmoAct2-Bridge-Dataset + episodes: 15000 + weight: 0.18 + morph_key: bridge + - repo_id: MolmoAct2-RT-1-Dataset + root: /home/alexw/tinyvla_data/MolmoAct2-RT-1-Dataset + episodes: 15000 + weight: 0.18 + morph_key: rt1 + - repo_id: div_ur5 + root: /home/alexw/tinyvla_data/diverse/div_ur5 + weight: 0.10 + morph_key: ur5 + - repo_id: div_xarm + root: /home/alexw/tinyvla_data/diverse/div_xarm + weight: 0.09 + morph_key: xarm + - repo_id: div_dlr_edan + root: /home/alexw/tinyvla_data/diverse/div_dlr_edan + weight: 0.05 + morph_key: dlr_edan + - repo_id: heldout_tidybot + root: /home/alexw/tinyvla_data/heldout/heldout_tidybot + weight: 0.08 + morph_key: tidybot + - repo_id: heldout_cmu_stretch + root: /home/alexw/tinyvla_data/heldout/heldout_cmu_stretch + weight: 0.07 + morph_key: hello_stretch + +policy: + chunk_size: 50 + n_action_steps: 50 + image_size: 256 + freeze_lm: false + freeze_vision_encoder: true + num_embodiments: 16 + conditioning: morph_qwen + num_morph_readout: 4 + morph_text_max_len: 32 + action_space: canonical + use_demo_conditioning: true + n_support: 3 + +backbone_lr_mult: 0.1 +batch_size: 40 +grad_accum: 2 +num_workers: 12 +lr: 7.0e-5 +warmup_steps: 500 +steps: 20000 +grad_clip: 10.0 +log_freq: 50 +save_freq: 5000 +mixed_precision: bf16 diff --git a/tinyvla2/configs/tv2_C_incontext_v2.yaml b/tinyvla2/configs/tv2_C_incontext_v2.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c99bf3d2cfa936e07e022a616a6d47f042ae116a --- /dev/null +++ b/tinyvla2/configs/tv2_C_incontext_v2.yaml @@ -0,0 +1,73 @@ +# Variant: in-context few-shot DEMO conditioning — a few (obs, action) exemplar +# pairs from the target robot's own data, encoded via the shared spatial-aware +# vision encoder + a small MLP, WITHOUT any gradient update to adapt to that robot. +# Composed with the text-via-Qwen morphology channel (both active). Resumes from +# tv2_C_qwen_morph/final (LM already unfrozen, morph_out_proj already trained). +output_dir: /home/alexw/tinyvla/outputs/tv2_C_incontext_v2 +wandb: null +seed: 42 +robot_prompts: /home/alexw/tinyvla/configs/morphology/robot_prompts.yaml +resume_from: /home/alexw/tinyvla/outputs/tv2_C_qwen_morph/final +resume_step: 0 + +datasets: + - root_glob: /home/alexw/tinyvla_data/so101_v3/* + weight: 0.25 + morph_key: so101 + - repo_id: MolmoAct2-Bridge-Dataset + root: /home/alexw/tinyvla_data/MolmoAct2-Bridge-Dataset + episodes: 15000 + weight: 0.18 + morph_key: bridge + - repo_id: MolmoAct2-RT-1-Dataset + root: /home/alexw/tinyvla_data/MolmoAct2-RT-1-Dataset + episodes: 15000 + weight: 0.18 + morph_key: rt1 + - repo_id: div_ur5 + root: /home/alexw/tinyvla_data/diverse/div_ur5 + weight: 0.10 + morph_key: ur5 + - repo_id: div_xarm + root: /home/alexw/tinyvla_data/diverse/div_xarm + weight: 0.09 + morph_key: xarm + - repo_id: div_dlr_edan + root: /home/alexw/tinyvla_data/diverse/div_dlr_edan + weight: 0.05 + morph_key: dlr_edan + - repo_id: heldout_tidybot + root: /home/alexw/tinyvla_data/heldout/heldout_tidybot + weight: 0.08 + morph_key: tidybot + - repo_id: heldout_cmu_stretch + root: /home/alexw/tinyvla_data/heldout/heldout_cmu_stretch + weight: 0.07 + morph_key: hello_stretch + +policy: + chunk_size: 50 + n_action_steps: 50 + image_size: 256 + freeze_lm: false + freeze_vision_encoder: true + num_embodiments: 16 + conditioning: morph_qwen + num_morph_readout: 4 + morph_text_max_len: 32 + action_space: canonical + use_demo_conditioning: true + n_support: 3 + support_other_task: true + +backbone_lr_mult: 0.1 +batch_size: 40 +grad_accum: 2 +num_workers: 12 +lr: 7.0e-5 +warmup_steps: 500 +steps: 20000 +grad_clip: 10.0 +log_freq: 50 +save_freq: 5000 +mixed_precision: bf16 diff --git a/tinyvla2/configs/tv2_C_incontext_v3.yaml b/tinyvla2/configs/tv2_C_incontext_v3.yaml new file mode 100644 index 0000000000000000000000000000000000000000..89941d922461d2770bf1186ce5572ba744e8cb17 --- /dev/null +++ b/tinyvla2/configs/tv2_C_incontext_v3.yaml @@ -0,0 +1,77 @@ +# Variant: in-context few-shot DEMO conditioning — a few (obs, action) exemplar +# pairs from the target robot's own data, encoded via the shared spatial-aware +# vision encoder + a small MLP, WITHOUT any gradient update to adapt to that robot. +# Composed with the text-via-Qwen morphology channel (both active). Resumes from +# tv2_C_qwen_morph/final (LM already unfrozen, morph_out_proj already trained). +output_dir: /home/alexw/tinyvla/outputs/tv2_C_incontext_v3 +wandb: null +seed: 42 +robot_prompts: /home/alexw/tinyvla/configs/morphology/robot_prompts.yaml +resume_from: /home/alexw/tinyvla/outputs/tv2_C_qwen_morph/final +resume_step: 0 + +datasets: + - root_glob: /home/alexw/tinyvla_data/so101_v3/* + weight: 0.25 + morph_key: so101 + - repo_id: MolmoAct2-Bridge-Dataset + root: /home/alexw/tinyvla_data/MolmoAct2-Bridge-Dataset + episodes: 15000 + weight: 0.18 + morph_key: bridge + - repo_id: MolmoAct2-RT-1-Dataset + root: /home/alexw/tinyvla_data/MolmoAct2-RT-1-Dataset + episodes: 15000 + weight: 0.18 + morph_key: rt1 + - repo_id: div_ur5 + root: /home/alexw/tinyvla_data/diverse/div_ur5 + weight: 0.10 + morph_key: ur5 + - repo_id: div_xarm + root: /home/alexw/tinyvla_data/diverse/div_xarm + weight: 0.09 + morph_key: xarm + - repo_id: div_dlr_edan + root: /home/alexw/tinyvla_data/diverse/div_dlr_edan + weight: 0.05 + morph_key: dlr_edan + - repo_id: heldout_tidybot + root: /home/alexw/tinyvla_data/heldout/heldout_tidybot + weight: 0.08 + morph_key: tidybot + - repo_id: heldout_cmu_stretch + root: /home/alexw/tinyvla_data/heldout/heldout_cmu_stretch + weight: 0.07 + morph_key: hello_stretch + +policy: + chunk_size: 50 + n_action_steps: 50 + image_size: 256 + freeze_lm: false + freeze_vision_encoder: true + num_embodiments: 16 + conditioning: morph_qwen + num_morph_readout: 4 + morph_text_max_len: 32 + action_space: canonical + use_demo_conditioning: true + n_support: 3 + support_other_task: true + demo_tokens_per_example: 4 + demo_visual_groups: 4 + demo_action_keys: 8 + demo_hidden_mult: 4 + +backbone_lr_mult: 0.1 +batch_size: 40 +grad_accum: 2 +num_workers: 12 +lr: 7.0e-5 +warmup_steps: 500 +steps: 20000 +grad_clip: 10.0 +log_freq: 50 +save_freq: 5000 +mixed_precision: bf16 diff --git a/tinyvla2/configs/tv2_C_mega.yaml b/tinyvla2/configs/tv2_C_mega.yaml new file mode 100644 index 0000000000000000000000000000000000000000..740047f16175e1d9f193a36d187201014cf3e2f3 --- /dev/null +++ b/tinyvla2/configs/tv2_C_mega.yaml @@ -0,0 +1,71 @@ +# Continue-train C on the MAXIMALLY diverse robot set (9 embodiments) with BOTH +# channels describing the robot: numeric morphology descriptor (C scheme) AND a +# natural-language robot description in the Qwen prompt (slow path). Held out: LeKiwi +# (mobile SO-100 + kiwi base) — we have its zero-shot baseline (295mm) to compare. +output_dir: /home/alexw/tinyvla/outputs/tv2_C_mega +wandb: null +seed: 42 +morphology_descriptors: /home/alexw/tinyvla/configs/morphology/descriptors.yaml +robot_prompts: /home/alexw/tinyvla/configs/morphology/robot_prompts.yaml +resume_from: /home/alexw/tinyvla/outputs/tv2_C_diverse/final +resume_step: 0 + +datasets: + - root_glob: /home/alexw/tinyvla_data/so101_v3/* + weight: 0.20 + morph_key: so101 + - repo_id: MolmoAct2-Bridge-Dataset + root: /home/alexw/tinyvla_data/MolmoAct2-Bridge-Dataset + episodes: 12000 + weight: 0.15 + morph_key: bridge + - repo_id: MolmoAct2-RT-1-Dataset + root: /home/alexw/tinyvla_data/MolmoAct2-RT-1-Dataset + episodes: 12000 + weight: 0.15 + morph_key: rt1 + - repo_id: div_ur5 + root: /home/alexw/tinyvla_data/diverse/div_ur5 + weight: 0.10 + morph_key: ur5 + - repo_id: div_xarm + root: /home/alexw/tinyvla_data/diverse/div_xarm + weight: 0.09 + morph_key: xarm + - repo_id: div_dlr_edan + root: /home/alexw/tinyvla_data/diverse/div_dlr_edan + weight: 0.06 + morph_key: dlr_edan + - repo_id: heldout_tidybot + root: /home/alexw/tinyvla_data/heldout/heldout_tidybot + weight: 0.09 + morph_key: tidybot + - repo_id: heldout_cmu_stretch + root: /home/alexw/tinyvla_data/heldout/heldout_cmu_stretch + weight: 0.08 + morph_key: hello_stretch + - repo_id: heldout_jaco_play + root: /home/alexw/tinyvla_data/heldout/heldout_jaco_play + weight: 0.08 + morph_key: jaco + +policy: + chunk_size: 50 + n_action_steps: 50 + image_size: 256 + freeze_lm: true + freeze_vision_encoder: true + num_embodiments: 16 + conditioning: morph + action_space: canonical + +batch_size: 64 +grad_accum: 2 +num_workers: 12 +lr: 7.0e-5 +warmup_steps: 500 +steps: 20000 +grad_clip: 10.0 +log_freq: 50 +save_freq: 5000 +mixed_precision: bf16 diff --git a/tinyvla2/configs/tv2_C_qwen_morph.yaml b/tinyvla2/configs/tv2_C_qwen_morph.yaml new file mode 100644 index 0000000000000000000000000000000000000000..9a41ac91c35f781b82e43b1528ace3459b820e7b --- /dev/null +++ b/tinyvla2/configs/tv2_C_qwen_morph.yaml @@ -0,0 +1,73 @@ +# Variant: morphology conditioning via the SAME shared Qwen (text-only pass over a +# robot description, e.g. "Robot: LeKiwi, a 5-DOF SO-100 arm on a three-wheel +# holonomic mobile base."), replacing the from-scratch MLP-over-numbers encoder. +# ALSO unfreezes the LM (freeze_lm: false) — per user request, since the frozen LM +# was the likely reason the earlier text-prompt-in-instruction experiment failed to +# learn to use new text. Resumes from C-diverse (8 robots, our best checkpoint: +# 150mm on held-out jaco). LeKiwi held out again for direct comparison (prior +# baselines: C-diverse numeric-only zero-shot 295mm; C-mega text-in-prompt 522mm). +output_dir: /home/alexw/tinyvla/outputs/tv2_C_qwen_morph +wandb: null +seed: 42 +robot_prompts: /home/alexw/tinyvla/configs/morphology/robot_prompts.yaml +resume_from: /home/alexw/tinyvla/outputs/tv2_C_diverse/final +resume_step: 0 + +datasets: + - root_glob: /home/alexw/tinyvla_data/so101_v3/* + weight: 0.25 + morph_key: so101 + - repo_id: MolmoAct2-Bridge-Dataset + root: /home/alexw/tinyvla_data/MolmoAct2-Bridge-Dataset + episodes: 15000 + weight: 0.18 + morph_key: bridge + - repo_id: MolmoAct2-RT-1-Dataset + root: /home/alexw/tinyvla_data/MolmoAct2-RT-1-Dataset + episodes: 15000 + weight: 0.18 + morph_key: rt1 + - repo_id: div_ur5 + root: /home/alexw/tinyvla_data/diverse/div_ur5 + weight: 0.10 + morph_key: ur5 + - repo_id: div_xarm + root: /home/alexw/tinyvla_data/diverse/div_xarm + weight: 0.09 + morph_key: xarm + - repo_id: div_dlr_edan + root: /home/alexw/tinyvla_data/diverse/div_dlr_edan + weight: 0.05 + morph_key: dlr_edan + - repo_id: heldout_tidybot + root: /home/alexw/tinyvla_data/heldout/heldout_tidybot + weight: 0.08 + morph_key: tidybot + - repo_id: heldout_cmu_stretch + root: /home/alexw/tinyvla_data/heldout/heldout_cmu_stretch + weight: 0.07 + morph_key: hello_stretch + +policy: + chunk_size: 50 + n_action_steps: 50 + image_size: 256 + freeze_lm: false + freeze_vision_encoder: true + num_embodiments: 16 + conditioning: morph_qwen + num_morph_readout: 4 + morph_text_max_len: 32 + action_space: canonical + +backbone_lr_mult: 0.1 +batch_size: 48 +grad_accum: 2 +num_workers: 12 +lr: 7.0e-5 +warmup_steps: 500 +steps: 20000 +grad_clip: 10.0 +log_freq: 50 +save_freq: 5000 +mixed_precision: bf16 diff --git a/tinyvla2/configs/tv2_C_scaled.yaml b/tinyvla2/configs/tv2_C_scaled.yaml new file mode 100644 index 0000000000000000000000000000000000000000..f437d91198c5fb7b3483cebdbb31fdbdbc8d7f98 --- /dev/null +++ b/tinyvla2/configs/tv2_C_scaled.yaml @@ -0,0 +1,76 @@ +# SCALED C: the setup that works (canonical actions + numeric morphology +# descriptor), given a real budget. The reality check showed the model barely beat +# the trivial baseline (mean ratio 0.99 in-training) because 40k+20k steps were +# spread across 8 morphologies with a frozen backbone. This run scales: +# data : full Bridge (53K eps) + full RT-1 (87K eps) instead of 12-15K subsets +# -> pool 4.4M -> ~9.1M frames +# steps: 60k (was 20k) at eff batch 128 -> 7.7M samples (~0.85 epoch) +# model: LM UNFROZEN at 0.1x lr (the single biggest lever in TinyVLA v1: +# frozen 60% -> unfrozen 86% on LIBERO) +# Held out: LeKiwi (mobile) + jaco stays in for training here, so held-out +# comparison uses LeKiwi; jaco numbers become in-distribution reference. +output_dir: /home/alexw/tinyvla/outputs/tv2_C_scaled +wandb: null +seed: 42 +morphology_descriptors: /home/alexw/tinyvla/configs/morphology/descriptors.yaml +resume_from: /home/alexw/tinyvla/outputs/tv2_C_diverse/final +resume_step: 0 + +datasets: + - root_glob: /home/alexw/tinyvla_data/so101_v3/* + weight: 0.22 + morph_key: so101 + - repo_id: MolmoAct2-Bridge-Dataset # FULL 53,192 eps (1.89M frames) + root: /home/alexw/tinyvla_data/MolmoAct2-Bridge-Dataset + weight: 0.22 + morph_key: bridge + - repo_id: MolmoAct2-RT-1-Dataset # FULL 87,212 eps (3.79M frames) + root: /home/alexw/tinyvla_data/MolmoAct2-RT-1-Dataset + weight: 0.22 + morph_key: rt1 + - repo_id: div_ur5 + root: /home/alexw/tinyvla_data/diverse/div_ur5 + weight: 0.09 + morph_key: ur5 + - repo_id: div_xarm + root: /home/alexw/tinyvla_data/diverse/div_xarm + weight: 0.08 + morph_key: xarm + - repo_id: div_dlr_edan + root: /home/alexw/tinyvla_data/diverse/div_dlr_edan + weight: 0.04 + morph_key: dlr_edan + - repo_id: heldout_tidybot + root: /home/alexw/tinyvla_data/heldout/heldout_tidybot + weight: 0.05 + morph_key: tidybot + - repo_id: heldout_cmu_stretch + root: /home/alexw/tinyvla_data/heldout/heldout_cmu_stretch + weight: 0.04 + morph_key: hello_stretch + - repo_id: heldout_jaco_play + root: /home/alexw/tinyvla_data/heldout/heldout_jaco_play + weight: 0.04 + morph_key: jaco + +policy: + chunk_size: 50 + n_action_steps: 50 + image_size: 256 + freeze_lm: false # the biggest known lever + freeze_vision_encoder: true + num_embodiments: 16 + conditioning: morph # numeric descriptor: the channel that actually works + action_space: canonical + +backbone_lr_mult: 0.1 +batch_size: 48 +grad_accum: 3 # eff batch 144 +num_workers: 12 +lr: 1.0e-4 +warmup_steps: 1000 +steps: 60000 +grad_clip: 10.0 +log_freq: 100 +save_freq: 10000 +mixed_precision: bf16 diff --git a/tinyvla2/configs/tv2_demoonly.yaml b/tinyvla2/configs/tv2_demoonly.yaml new file mode 100644 index 0000000000000000000000000000000000000000..93b3cfea1571357ee801a03d4519164a6a3880df --- /dev/null +++ b/tinyvla2/configs/tv2_demoonly.yaml @@ -0,0 +1,89 @@ +# DEMO-ONLY: every robot-identity channel removed (no embodiment ID, no +# morphology descriptor, no robot text) — the K demo exemplars in the LM +# attention stream are the ONLY cue for which robot this is. Per-source action +# normalization creates the pressure: to predict correctly the model MUST infer +# the robot's action scale, and only the demos carry it. +# (was: RICH SLOW PATH) demos + robot text + instruction + image all inside the LM +# sequence (demos in the ATTENTION STREAM), LM emits latent tokens conditioning +# the fast path. Dual-rate PRESERVED (fast path unchanged, latent cacheable). +# (was: in-context few-shot DEMO conditioning) — a few (obs, action) exemplar +# pairs from the target robot's own data, encoded via the shared spatial-aware +# vision encoder + a small MLP, WITHOUT any gradient update to adapt to that robot. +# Composed with the text-via-Qwen morphology channel (both active). Resumes from +# tv2_C_qwen_morph/final (LM already unfrozen, morph_out_proj already trained). +output_dir: /home/alexw/tinyvla/outputs/tv2_demoonly +wandb: null +seed: 42 +robot_prompts: /home/alexw/tinyvla/configs/morphology/robot_prompts.yaml +resume_from: /home/alexw/tinyvla/outputs/tv2_C_qwen_morph/final +resume_step: 0 + +datasets: + - root_glob: /home/alexw/tinyvla_data/so101_v3/* + weight: 0.25 + morph_key: so101 + - repo_id: MolmoAct2-Bridge-Dataset + root: /home/alexw/tinyvla_data/MolmoAct2-Bridge-Dataset + episodes: 15000 + weight: 0.18 + morph_key: bridge + - repo_id: MolmoAct2-RT-1-Dataset + root: /home/alexw/tinyvla_data/MolmoAct2-RT-1-Dataset + episodes: 15000 + weight: 0.18 + morph_key: rt1 + - repo_id: div_ur5 + root: /home/alexw/tinyvla_data/diverse/div_ur5 + weight: 0.10 + morph_key: ur5 + - repo_id: div_xarm + root: /home/alexw/tinyvla_data/diverse/div_xarm + weight: 0.09 + morph_key: xarm + - repo_id: div_dlr_edan + root: /home/alexw/tinyvla_data/diverse/div_dlr_edan + weight: 0.05 + morph_key: dlr_edan + - repo_id: heldout_tidybot + root: /home/alexw/tinyvla_data/heldout/heldout_tidybot + weight: 0.08 + morph_key: tidybot + - repo_id: heldout_cmu_stretch + root: /home/alexw/tinyvla_data/heldout/heldout_cmu_stretch + weight: 0.07 + morph_key: hello_stretch + +policy: + chunk_size: 50 + n_action_steps: 50 + image_size: 256 + freeze_lm: false + freeze_vision_encoder: true + num_embodiments: 16 + conditioning: morph_qwen + num_morph_readout: 4 + morph_text_max_len: 32 + action_space: canonical + use_demo_conditioning: false # demos now go INTO the LM sequence, not the MLP encoder + n_support: 4 + support_other_task: true + demo_tokens_per_example: 4 + demo_visual_groups: 4 + demo_action_keys: 8 + demo_hidden_mult: 4 + vlm_native: true + num_action_readout: 16 + demo_action_keys_lm: 4 + demo_only: true + +backbone_lr_mult: 0.1 +batch_size: 24 +grad_accum: 3 +num_workers: 12 +lr: 7.0e-5 +warmup_steps: 500 +steps: 20000 +grad_clip: 10.0 +log_freq: 50 +save_freq: 5000 +mixed_precision: bf16 diff --git a/tinyvla2/configs/tv2_richslow.yaml b/tinyvla2/configs/tv2_richslow.yaml new file mode 100644 index 0000000000000000000000000000000000000000..2adee5541cbd75f28a6d96069bb7917b01af7032 --- /dev/null +++ b/tinyvla2/configs/tv2_richslow.yaml @@ -0,0 +1,83 @@ +# RICH SLOW PATH: demos + robot text + instruction + image all inside the LM +# sequence (demos in the ATTENTION STREAM), LM emits latent tokens conditioning +# the fast path. Dual-rate PRESERVED (fast path unchanged, latent cacheable). +# (was: in-context few-shot DEMO conditioning) — a few (obs, action) exemplar +# pairs from the target robot's own data, encoded via the shared spatial-aware +# vision encoder + a small MLP, WITHOUT any gradient update to adapt to that robot. +# Composed with the text-via-Qwen morphology channel (both active). Resumes from +# tv2_C_qwen_morph/final (LM already unfrozen, morph_out_proj already trained). +output_dir: /home/alexw/tinyvla/outputs/tv2_richslow +wandb: null +seed: 42 +robot_prompts: /home/alexw/tinyvla/configs/morphology/robot_prompts.yaml +resume_from: /home/alexw/tinyvla/outputs/tv2_C_qwen_morph/final +resume_step: 0 + +datasets: + - root_glob: /home/alexw/tinyvla_data/so101_v3/* + weight: 0.25 + morph_key: so101 + - repo_id: MolmoAct2-Bridge-Dataset + root: /home/alexw/tinyvla_data/MolmoAct2-Bridge-Dataset + episodes: 15000 + weight: 0.18 + morph_key: bridge + - repo_id: MolmoAct2-RT-1-Dataset + root: /home/alexw/tinyvla_data/MolmoAct2-RT-1-Dataset + episodes: 15000 + weight: 0.18 + morph_key: rt1 + - repo_id: div_ur5 + root: /home/alexw/tinyvla_data/diverse/div_ur5 + weight: 0.10 + morph_key: ur5 + - repo_id: div_xarm + root: /home/alexw/tinyvla_data/diverse/div_xarm + weight: 0.09 + morph_key: xarm + - repo_id: div_dlr_edan + root: /home/alexw/tinyvla_data/diverse/div_dlr_edan + weight: 0.05 + morph_key: dlr_edan + - repo_id: heldout_tidybot + root: /home/alexw/tinyvla_data/heldout/heldout_tidybot + weight: 0.08 + morph_key: tidybot + - repo_id: heldout_cmu_stretch + root: /home/alexw/tinyvla_data/heldout/heldout_cmu_stretch + weight: 0.07 + morph_key: hello_stretch + +policy: + chunk_size: 50 + n_action_steps: 50 + image_size: 256 + freeze_lm: false + freeze_vision_encoder: true + num_embodiments: 16 + conditioning: morph_qwen + num_morph_readout: 4 + morph_text_max_len: 32 + action_space: canonical + use_demo_conditioning: false # demos now go INTO the LM sequence, not the MLP encoder + n_support: 3 + support_other_task: true + demo_tokens_per_example: 4 + demo_visual_groups: 4 + demo_action_keys: 8 + demo_hidden_mult: 4 + vlm_native: true + num_action_readout: 16 + demo_action_keys_lm: 4 + +backbone_lr_mult: 0.1 +batch_size: 24 +grad_accum: 3 +num_workers: 12 +lr: 7.0e-5 +warmup_steps: 500 +steps: 20000 +grad_clip: 10.0 +log_freq: 50 +save_freq: 5000 +mixed_precision: bf16 diff --git a/tinyvla2/configs/tv2_robocasa.yaml b/tinyvla2/configs/tv2_robocasa.yaml new file mode 100644 index 0000000000000000000000000000000000000000..481440eb1a1ff13c113ab687f7e0dc7d0975ed97 --- /dev/null +++ b/tinyvla2/configs/tv2_robocasa.yaml @@ -0,0 +1,77 @@ +# Continue-train the best checkpoint (C-scaled, 60k) on RoboCasa365 — the first +# genuinely NEW domain in this project: simulated kitchen scenes (everything before +# was real-world tabletops) with PandaOmron, a Franka arm on a holonomic mobile base. +# 7356 eps / 1.5M frames / 20fps / 3 cameras. +# +# RoboCasa gets half the sampling weight; the other half replays the previous robots +# so we can tell integration from catastrophic forgetting (all old robots stay +# measurable with the same eval). LeKiwi remains fully held out. +output_dir: /home/alexw/tinyvla/outputs/tv2_robocasa +wandb: null +seed: 42 +morphology_descriptors: /home/alexw/tinyvla/configs/morphology/descriptors.yaml +resume_from: /home/alexw/tinyvla/outputs/tv2_C_scaled/final +resume_step: 0 + +datasets: + - repo_id: robocasa365 # NEW domain: sim kitchens, PandaOmron + root: /home/alexw/tinyvla_data/robocasa365 + weight: 0.50 + morph_key: panda_omron + - root_glob: /home/alexw/tinyvla_data/so101_v3/* + weight: 0.12 + morph_key: so101 + - repo_id: MolmoAct2-Bridge-Dataset + root: /home/alexw/tinyvla_data/MolmoAct2-Bridge-Dataset + weight: 0.11 + morph_key: bridge + - repo_id: MolmoAct2-RT-1-Dataset + root: /home/alexw/tinyvla_data/MolmoAct2-RT-1-Dataset + weight: 0.11 + morph_key: rt1 + - repo_id: div_ur5 + root: /home/alexw/tinyvla_data/diverse/div_ur5 + weight: 0.05 + morph_key: ur5 + - repo_id: div_xarm + root: /home/alexw/tinyvla_data/diverse/div_xarm + weight: 0.04 + morph_key: xarm + - repo_id: div_dlr_edan + root: /home/alexw/tinyvla_data/diverse/div_dlr_edan + weight: 0.02 + morph_key: dlr_edan + - repo_id: heldout_tidybot + root: /home/alexw/tinyvla_data/heldout/heldout_tidybot + weight: 0.02 + morph_key: tidybot + - repo_id: heldout_cmu_stretch + root: /home/alexw/tinyvla_data/heldout/heldout_cmu_stretch + weight: 0.02 + morph_key: hello_stretch + - repo_id: heldout_jaco_play + root: /home/alexw/tinyvla_data/heldout/heldout_jaco_play + weight: 0.01 + morph_key: jaco + +policy: + chunk_size: 50 + n_action_steps: 50 + image_size: 256 + freeze_lm: false + freeze_vision_encoder: true + num_embodiments: 16 + conditioning: morph + action_space: canonical + +backbone_lr_mult: 0.1 +batch_size: 48 +grad_accum: 3 +num_workers: 12 +lr: 7.0e-5 +warmup_steps: 500 +steps: 30000 +grad_clip: 10.0 +log_freq: 100 +save_freq: 10000 +mixed_precision: bf16 diff --git a/tinyvla2/configs/tv2_vlm_native.yaml b/tinyvla2/configs/tv2_vlm_native.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c1cb77045da2ac3665c59fcb3f058a07842ac78b --- /dev/null +++ b/tinyvla2/configs/tv2_vlm_native.yaml @@ -0,0 +1,83 @@ +# LLaVA-style: ALL modalities as tokens in one LM sequence, policy conditions on +# latent readout tokens. Demos live IN the attention stream (fix for v1-v3 failure). +# Trades away the dual-rate fast path. +# (was: in-context few-shot DEMO conditioning) — a few (obs, action) exemplar +# pairs from the target robot's own data, encoded via the shared spatial-aware +# vision encoder + a small MLP, WITHOUT any gradient update to adapt to that robot. +# Composed with the text-via-Qwen morphology channel (both active). Resumes from +# tv2_C_qwen_morph/final (LM already unfrozen, morph_out_proj already trained). +output_dir: /home/alexw/tinyvla/outputs/tv2_vlm_native +wandb: null +seed: 42 +robot_prompts: /home/alexw/tinyvla/configs/morphology/robot_prompts.yaml +resume_from: /home/alexw/tinyvla/outputs/tv2_C_qwen_morph/final +resume_step: 0 + +datasets: + - root_glob: /home/alexw/tinyvla_data/so101_v3/* + weight: 0.25 + morph_key: so101 + - repo_id: MolmoAct2-Bridge-Dataset + root: /home/alexw/tinyvla_data/MolmoAct2-Bridge-Dataset + episodes: 15000 + weight: 0.18 + morph_key: bridge + - repo_id: MolmoAct2-RT-1-Dataset + root: /home/alexw/tinyvla_data/MolmoAct2-RT-1-Dataset + episodes: 15000 + weight: 0.18 + morph_key: rt1 + - repo_id: div_ur5 + root: /home/alexw/tinyvla_data/diverse/div_ur5 + weight: 0.10 + morph_key: ur5 + - repo_id: div_xarm + root: /home/alexw/tinyvla_data/diverse/div_xarm + weight: 0.09 + morph_key: xarm + - repo_id: div_dlr_edan + root: /home/alexw/tinyvla_data/diverse/div_dlr_edan + weight: 0.05 + morph_key: dlr_edan + - repo_id: heldout_tidybot + root: /home/alexw/tinyvla_data/heldout/heldout_tidybot + weight: 0.08 + morph_key: tidybot + - repo_id: heldout_cmu_stretch + root: /home/alexw/tinyvla_data/heldout/heldout_cmu_stretch + weight: 0.07 + morph_key: hello_stretch + +policy: + chunk_size: 50 + n_action_steps: 50 + image_size: 256 + freeze_lm: false + freeze_vision_encoder: true + num_embodiments: 16 + conditioning: id # unused in vlm_native; robot text goes into the LM sequence + num_morph_readout: 4 + morph_text_max_len: 32 + action_space: canonical + use_demo_conditioning: true + n_support: 3 + support_other_task: true + demo_tokens_per_example: 4 + demo_visual_groups: 4 + demo_action_keys: 8 + demo_hidden_mult: 4 + vlm_native: true + num_action_readout: 16 + demo_action_keys_lm: 4 + +backbone_lr_mult: 0.1 +batch_size: 20 +grad_accum: 4 +num_workers: 12 +lr: 7.0e-5 +warmup_steps: 500 +steps: 20000 +grad_clip: 10.0 +log_freq: 50 +save_freq: 5000 +mixed_precision: bf16 diff --git a/tinyvla2/scripts/__pycache__/fewshot_grid.cpython-312.pyc b/tinyvla2/scripts/__pycache__/fewshot_grid.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..509d97fa4289b698d5b6afae469e9e50e3042c27 Binary files /dev/null and b/tinyvla2/scripts/__pycache__/fewshot_grid.cpython-312.pyc differ diff --git a/tinyvla2/scripts/_eval_richslow_10k.py b/tinyvla2/scripts/_eval_richslow_10k.py new file mode 100644 index 0000000000000000000000000000000000000000..d80bac2af1808b8e1bba02d6e0c61eaedc0bd484 --- /dev/null +++ b/tinyvla2/scripts/_eval_richslow_10k.py @@ -0,0 +1,141 @@ +#!/usr/bin/env python +"""Auto-eval for in-context demo conditioning on held-out LeKiwi. + +The key test: give the model K=3 REAL (obs, action) example pairs from LeKiwi's +OWN data at test time — NO gradient update, NO fine-tuning — and see if that beats +zero-shot with no examples, wrong examples, or the numeric/text descriptor alone. +""" + +from __future__ import annotations + +import numpy as np +import torch +import yaml +from lerobot.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata +from transformers import AutoTokenizer + +from tinyvla.data.canonical import CanonicalChunkStore, quantile_normalize +from tinyvla.modeling_tinyvla import TinyVLAPolicy + +CKPT = "outputs/tv2_richslow/step_10000" +DS = "lekiwi_cleanup" +ROOT = f"/home/alexw/tinyvla_data/lekiwi/{DS}" +K = 3 + + +@torch.no_grad() +def main(): + pol = TinyVLAPolicy.from_pretrained(CKPT).cuda().eval() + cfg = pol.config + tok = AutoTokenizer.from_pretrained(cfg.lm_model_name) + prompts = yaml.safe_load(open("configs/morphology/robot_prompts.yaml")) + + m = LeRobotDatasetMetadata(DS, root=ROOT) + ds = LeRobotDataset(DS, root=ROOT, delta_timestamps={"action": [t / m.fps for t in range(50)]}, + video_backend="torchcodec") + store = CanonicalChunkStore(DS, src_fps=m.fps, chunk=50) + st = store.compute_stats() + q01, q99 = np.asarray(st["q01"]), np.asarray(st["q99"]) + span = np.maximum(q99 - q01, 0.01 * np.median(np.abs(np.concatenate([q01, q99])) + 1e-6)) + mid = 0.5 * (q01 + q99) + + imk = sorted(k for k in ds.meta.features if k.startswith("observation.images")) + prim = next((k for k in imk if "front" in k or "base" in k), imk[0]) + + def integ(d): + return np.cumsum(d[:, :3], 0) + + def img_at(idx, ep, ep_start): + it = ds[idx] + im = torch.nn.functional.interpolate(it[prim][None], size=(256, 256), mode="bilinear")[0] + raw = quantile_normalize(store.chunk_for(ep, idx - ep_start), q01, q99) + act = torch.from_numpy(raw[:, :7]).float() + act = torch.nn.functional.pad(act, (0, cfg.max_action_dim - 7)) + return im, act + + # REAL LeKiwi demo pairs from early episodes (support pool), disjoint from test + support_pool_eps = range(0, min(20, m.total_episodes - 15)) + test = range(max(0, m.total_episodes - 15), m.total_episodes) + + def get_real_support(k): + imgs, acts = [], [] + for _ in range(k): + ep = int(np.random.choice(list(support_pool_eps))) + s = int(m.episodes["dataset_from_index"][ep]) + e = int(m.episodes["dataset_to_index"][ep]) + idx = int(np.random.randint(s, max(s + 1, e - 1))) + im, act = img_at(idx, ep, s) + imgs.append(im) + acts.append(act) + return torch.stack(imgs), torch.stack(acts) + + def get_wrong_support(k): + # random noise images + random actions — a garbage support set control + return torch.rand(k, 3, 256, 256), torch.randn(k, 50, cfg.max_action_dim) * 0.3 + + + # control: demos from a DIFFERENT ROBOT (jaco) — if the model truly extracts + # robot-specific dynamics from exemplars, these should be worse than LeKiwi's own + def get_otherrobot_support(k): + from lerobot.datasets.lerobot_dataset import LeRobotDataset as _LD, LeRobotDatasetMetadata as _LM + from tinyvla.data.canonical import CanonicalChunkStore as _CS + jr = "/home/alexw/tinyvla_data/heldout/heldout_jaco_play" + jm = _LM("heldout_jaco_play", root=jr) + jds = _LD("heldout_jaco_play", root=jr, delta_timestamps={"action": [t / jm.fps for t in range(50)]}, video_backend="torchcodec") + jstore = _CS("heldout_jaco_play", src_fps=jm.fps, chunk=50) + jst = jstore.compute_stats(200) + jq01, jq99 = np.asarray(jst["q01"]), np.asarray(jst["q99"]) + jimk = sorted(kk for kk in jds.meta.features if kk.startswith("observation.images"))[0] + imgs, acts = [], [] + for _ in range(k): + ep = int(np.random.randint(0, 50)) + s = int(jm.episodes["dataset_from_index"][ep]); e = int(jm.episodes["dataset_to_index"][ep]) + idx = int(np.random.randint(s, max(s + 1, e - 1))) + it = jds[idx] + im = torch.nn.functional.interpolate(it[jimk][None], size=(256, 256), mode="bilinear")[0] + raw = quantile_normalize(jstore.chunk_for(ep, idx - s), jq01, jq99) + act = torch.nn.functional.pad(torch.from_numpy(raw[:, :7]).float(), (0, cfg.max_action_dim - 7)) + imgs.append(im); acts.append(act) + return torch.stack(imgs), torch.stack(acts) + + def run(support_fn, desc_text): + ids_t = tok([desc_text] if desc_text else [""], padding="max_length", truncation=True, + max_length=cfg.morph_text_max_len, return_tensors="pt") + errs = [] + for ep in test: + s = int(m.episodes["dataset_from_index"][ep]) + e = int(m.episodes["dataset_to_index"][ep]) + for idx in range(s, e - 1, 30): + item = ds[idx] + task = item.get("task") or "" + t = tok([task], padding=True, truncation=True, max_length=48, return_tensors="pt") + img = torch.nn.functional.interpolate(item[prim][None], size=(256, 256), mode="bilinear")[0] + stt = torch.nn.functional.pad(item["observation.state"].float()[:6], (0, cfg.max_state_dim - 6)) + sup_img, sup_act = support_fn(K) + b = {"observation.images.cam0": img[None].cuda(), + "observation.images.cam1": torch.zeros_like(img)[None].cuda(), + "observation.state": stt[None].cuda(), + "observation.language.tokens": t["input_ids"].cuda(), + "observation.language.attention_mask": t["attention_mask"].bool().cuda(), + "morph_text_ids": ids_t["input_ids"].cuda(), + "morph_text_mask": ids_t["attention_mask"].bool().cuda(), + "support_images": sup_img[None].cuda(), "support_actions": sup_act[None].cuda(), + "embodiment_id": torch.tensor([0]).cuda()} + with torch.autocast("cuda", torch.bfloat16): + pr = pol.predict_action_chunk(b)[0].cpu().float().numpy() + gu = quantile_normalize(store.chunk_for(ep, idx - s), q01, q99)[:, :7] * span / 2 + mid + pu = pr[:, :7] * span / 2 + mid + errs.append(np.linalg.norm(integ(pu)[-1] - integ(gu)[-1]) * 1000) + return np.mean(errs) + + print("=== RICH-SLOW (demos in LM attention stream, no FT) on held-out LeKiwi ===") + print(f"REAL LeKiwi demos (K={K}) + desc=lekiwi endpoint {run(get_real_support, prompts['lekiwi']):.1f}mm") + print(f"REAL LeKiwi demos (K={K}) + desc=none endpoint {run(get_real_support, None):.1f}mm") + print(f"WRONG/garbage demos + desc=lekiwi endpoint {run(get_wrong_support, prompts['lekiwi']):.1f}mm") + print(f"WRONG/garbage demos + desc=none endpoint {run(get_wrong_support, None):.1f}mm") + print(f"OTHER-ROBOT (jaco) demos + desc=none endpoint {run(get_otherrobot_support, None):.1f}mm") + print("\nprior baselines: C-diverse 295mm | C-qwen-morph desc=none 294mm | v1 demo(same-task) real=290.6 garbage=288.5 (dead)") + + +if __name__ == "__main__": + main() diff --git a/tinyvla2/scripts/audit_so101_fk.py b/tinyvla2/scripts/audit_so101_fk.py new file mode 100644 index 0000000000000000000000000000000000000000..e2b0268d970bc6e39e54cd8a853c8038138941e8 --- /dev/null +++ b/tinyvla2/scripts/audit_so101_fk.py @@ -0,0 +1,103 @@ +#!/usr/bin/env python +"""Gate 1: per-dataset SO101 FK calibration audit. + +Same recorded degrees can map to different physical poses across lerobot +zero-conventions → FK garbage that would poison B/C canonical labels while +leaving A untouched (a confound AGAINST the hypothesis). Drop datasets whose +FK produces non-physical EE trajectories. + +Checks per dataset (sampled frames across episodes): + - reach in [0.02, 0.40] m (SO101 max reach ~0.35) + - z above a floor (> -0.20 m; base frame origin at arm mount) + - trajectory smoothness (median consecutive EE step < 0.05 m at 30fps) + - FK(state) vs FK(action-target) diffs correlate (both go through FK, deltas + should track since action = commanded target of the same arm) + +Writes ~/tinyvla_data/so101_fk_audit.json with per-dataset verdict. +""" + +from __future__ import annotations + +import json +from pathlib import Path + +import numpy as np + +DATA_ROOT = Path.home() / "tinyvla_data" / "so101_v3" +OUT = Path.home() / "tinyvla_data" / "so101_fk_audit.json" + + +def audit_one(name, root, fk, n_frames=200): + from lerobot.datasets.lerobot_dataset import LeRobotDataset + + ds = LeRobotDataset(name, root=root) + anames = ds.meta.features["action"]["shape"][0] + if anames != 6: + return {"verdict": "SKIP", "reason": f"action dim {anames} != 6"} + # smoothness needs CONSECUTIVE frames → sample a few contiguous windows; + # reach/corr can use sparse knots across the whole dataset + n = len(ds) + eps = ds.meta.episodes + sparse = np.linspace(0, n - 1, min(n_frames, n)).astype(int) + pos_state, pos_act = [], [] + for i in sparse: + item = ds[int(i)] + pos_state.append(fk.ee_pose(item["observation.state"].numpy())[:3, 3]) + pos_act.append(fk.ee_pose(item["action"].numpy())[:3, 3]) + ps = np.array(pos_state) + pa = np.array(pos_act) + reach = np.linalg.norm(ps, axis=1) + + # consecutive-frame EE steps within the first episode (real per-frame motion) + e0, e1 = int(eps["dataset_from_index"][0]), int(eps["dataset_to_index"][0]) + consec = [] + prev = None + for i in range(e0, min(e1, e0 + 150)): + p = fk.ee_pose(ds[i]["observation.state"].numpy())[:3, 3] + if prev is not None: + consec.append(np.linalg.norm(p - prev)) + prev = p + steps = np.array(consec) if consec else np.array([0.0]) + # correlation of state-vs-target displacement over sampled knots + d_state = np.diff(ps, axis=0).flatten() + d_act = np.diff(pa, axis=0).flatten() + corr = float(np.corrcoef(d_state, d_act)[0, 1]) if d_state.std() > 1e-9 else 0.0 + + ok_reach = bool(0.02 < reach.mean() < 0.40 and reach.max() < 0.50) + ok_z = bool(ps[:, 2].min() > -0.20) + ok_smooth = bool(np.median(steps) < 0.06) + ok_corr = bool(corr > 0.5) + verdict = "KEEP" if (ok_reach and ok_z and ok_smooth and ok_corr) else "DROP" + return { + "verdict": verdict, + "reach_mean": round(float(reach.mean()), 3), + "reach_max": round(float(reach.max()), 3), + "z_min": round(float(ps[:, 2].min()), 3), + "step_median": round(float(np.median(steps)), 4), + "corr_state_target": round(corr, 3), + "flags": {"reach": ok_reach, "z": ok_z, "smooth": ok_smooth, "corr": ok_corr}, + } + + +def main(): + from tinyvla.data.kinematics_so101 import SO101FK + + fk = SO101FK() + results = {} + roots = sorted(DATA_ROOT.iterdir()) + for r in roots: + if not (r / "meta" / "info.json").exists(): + continue + try: + res = audit_one(r.name, r, fk) + except Exception as e: + res = {"verdict": "ERROR", "reason": f"{type(e).__name__}: {str(e)[:120]}"} + results[r.name] = res + print(f"{res['verdict']:6} {r.name}: {res}") + OUT.write_text(json.dumps(results, indent=1)) + keep = sum(1 for v in results.values() if v["verdict"] == "KEEP") + print(f"\nKEEP {keep}/{len(results)} -> {OUT}") + + +if __name__ == "__main__": + main() diff --git a/tinyvla2/scripts/build_ee_sidecar.py b/tinyvla2/scripts/build_ee_sidecar.py new file mode 100644 index 0000000000000000000000000000000000000000..098616bc5bc8795d7d652b73cff06a0463d23836 --- /dev/null +++ b/tinyvla2/scripts/build_ee_sidecar.py @@ -0,0 +1,169 @@ +#!/usr/bin/env python +"""Build per-frame canonical EE-pose sidecars (offline, once). + +For each frame stores the absolute end-effector pose implied by BOTH the +observation state and the action target, in the robot base frame: + cols: episode_index, frame_index, + s_px s_py s_pz s_qx s_qy s_qz s_qw (state EE pose) + a_px a_py a_pz a_qx a_qy a_qz a_qw (action-target EE pose) + grip (raw gripper channel, for q1-q99 later) + +Canonical chunk deltas are assembled at train time from these poses (anchor, +frequency, rotation convention are all decided there — this file is convention-free). + +Source kinds: + so101 : FK(joint_deg) for state and action (both are joint vectors). + ee : state/action already EE; reconstruct absolute pose. Bridge/LIBERO/DROID + store EE pose in observation.state; action is a delta → target = state ⊕ action. + (implemented per-source as adapters are verified.) + +Usage: + python scripts/build_ee_sidecar.py --kind so101 --glob '~/tinyvla_data/so101_v3/*' +""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +import numpy as np +import pyarrow as pa +import pyarrow.parquet as pq + +OUT_DIR = Path.home() / "tinyvla_data" / "ee_sidecar" + + +def _read_cols(name, root): + """State/action/ep/frame columns WITHOUT video decode (raw parquet).""" + from lerobot.datasets.lerobot_dataset import LeRobotDataset + + ds = LeRobotDataset(name, root=root) + hf = ds.reader.hf_dataset.with_format("numpy") + return ( + np.asarray(hf["observation.state"], dtype=np.float64), + np.asarray(hf["action"], dtype=np.float64), + np.asarray(hf["episode_index"]).astype(int), + np.asarray(hf["frame_index"]).astype(int), + ) + + +def build_so101(name, root, fk): + from scipy.spatial.transform import Rotation + + state, action, ep, fr = _read_cols(name, root) + n = len(state) + cols = ["episode_index", "frame_index", + "s_px", "s_py", "s_pz", "s_qx", "s_qy", "s_qz", "s_qw", + "a_px", "a_py", "a_pz", "a_qx", "a_qy", "a_qz", "a_qw", "grip"] + buf = {k: np.empty(n, dtype=np.float64) for k in cols} + buf["episode_index"] = ep.astype(np.float64) + buf["frame_index"] = fr.astype(np.float64) + for i in range(n): + Ts = fk.ee_pose(state[i]) + Ta = fk.ee_pose(action[i]) + sp, sq = Ts[:3, 3], Rotation.from_matrix(Ts[:3, :3]).as_quat() + ap, aq = Ta[:3, 3], Rotation.from_matrix(Ta[:3, :3]).as_quat() + for j, k in enumerate(("s_px", "s_py", "s_pz")): buf[k][i] = sp[j] + for j, k in enumerate(("s_qx", "s_qy", "s_qz", "s_qw")): buf[k][i] = sq[j] + for j, k in enumerate(("a_px", "a_py", "a_pz")): buf[k][i] = ap[j] + for j, k in enumerate(("a_qx", "a_qy", "a_qz", "a_qw")): buf[k][i] = aq[j] + buf["grip"][i] = action[i, 5] + return pa.table(buf) + + +def build_ee_from_state(name, root): + """EE-native sources (Bridge, RT-1): state carries absolute EE pose. + + Bridge state = [x y z roll pitch yaw pad gripper]; RT-1 state = + [x y z rx ry rz rw gripper] (quat). We store the STATE pose as canonical + absolute pose (both s_* and a_* set to state pose; targets reconstructed at + train time as state ⊕ action-delta if needed, but state-derived deltas are + the uniform choice per plan → a_* == next-frame not needed here). + """ + from lerobot.datasets.lerobot_dataset import LeRobotDataset + from scipy.spatial.transform import Rotation + + ds = LeRobotDataset(name, root=root) + snames = ds.meta.features["observation.state"].get("names", {}) + flat = snames.get("motors", snames) if isinstance(snames, dict) else snames + state, _, ep, fr = _read_cols(name, root) + n = len(state) + # RoboCasa/PandaOmron: state = base_pos(3) base_quat(4) ee_pos(3) ee_quat(4) grip(2) + # -> the EE block is not at the front, slice it out first + if state.shape[1] == 16: + state = np.concatenate([state[:, 7:14], state[:, 14:15]], axis=1) # pos,quat,grip + # detect quaternion: named "rw", or 8-dim state whose dims 3:7 are unit-norm + is_quat = ("rw" in flat) or ( + state.shape[1] >= 8 and abs(np.linalg.norm(state[:200, 3:7], axis=1).mean() - 1.0) < 1e-2 + ) + cols = ["episode_index", "frame_index", + "s_px", "s_py", "s_pz", "s_qx", "s_qy", "s_qz", "s_qw", + "a_px", "a_py", "a_pz", "a_qx", "a_qy", "a_qz", "a_qw", "grip"] + buf = {k: np.empty(n, dtype=np.float64) for k in cols} + buf["episode_index"] = ep.astype(np.float64) + buf["frame_index"] = fr.astype(np.float64) + pos = state[:, :3] + if is_quat: # [x y z rx ry rz rw gripper] + quat = state[:, 3:7] + grip = state[:, 7] + else: # [x y z roll pitch yaw (pad) (gripper)] + quat = Rotation.from_euler("xyz", state[:, 3:6]).as_quat() + grip = state[:, 7] if state.shape[1] > 7 else np.zeros(n) # free-body (UAV): no gripper + for j, k in enumerate(("s_px", "s_py", "s_pz")): buf[k] = pos[:, j] + for j, k in enumerate(("s_qx", "s_qy", "s_qz", "s_qw")): buf[k] = quat[:, j] + for j, k in enumerate(("a_px", "a_py", "a_pz")): buf[k] = pos[:, j] + for j, k in enumerate(("a_qx", "a_qy", "a_qz", "a_qw")): buf[k] = quat[:, j] + buf["grip"] = grip + return pa.table(buf) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--kind", choices=["so101", "ee"], required=True) + parser.add_argument("--glob", required=True) + parser.add_argument("--audit", default=str(Path.home() / "tinyvla_data/so101_fk_audit.json"), + help="skip datasets not marked KEEP") + args = parser.parse_args() + + import glob as _glob + import json + + OUT_DIR.mkdir(parents=True, exist_ok=True) + keep = None + audit_path = Path(args.audit).expanduser() + if args.kind == "so101" and audit_path.exists(): + audit = json.loads(audit_path.read_text()) + keep = {k for k, v in audit.items() if v.get("verdict") == "KEEP"} + + fk = None + if args.kind == "so101": + from tinyvla.data.kinematics_so101 import SO101FK + fk = SO101FK() + + roots = sorted(Path(p) for p in _glob.glob(str(Path(args.glob).expanduser()))) + done = 0 + for r in roots: + if not (r / "meta" / "info.json").exists(): + continue + if keep is not None and r.name not in keep: + print(f"skip {r.name} (not KEEP)") + continue + out = OUT_DIR / f"{r.name}.parquet" + if out.exists(): + done += 1 + continue + try: + if args.kind == "so101": + tbl = build_so101(r.name, r, fk) + else: + tbl = build_ee_from_state(r.name, r) + pq.write_table(tbl, out) + done += 1 + print(f"[{done}] {r.name}: {tbl.num_rows} frames -> {out}") + except Exception as e: + print(f"FAIL {r.name}: {type(e).__name__}: {str(e)[:150]}") + print(f"done {done}") + + +if __name__ == "__main__": + main() diff --git a/tinyvla2/scripts/diag_obs_gap.py b/tinyvla2/scripts/diag_obs_gap.py new file mode 100644 index 0000000000000000000000000000000000000000..7d6b06096dbd602123412eee029aa3353bc02151 --- /dev/null +++ b/tinyvla2/scripts/diag_obs_gap.py @@ -0,0 +1,131 @@ +#!/usr/bin/env python +"""Diagnose the env-obs adapter: compare policy predictions from env-rendered +observations vs dataset observations at the SAME init state. + +If pred(dataset obs) is close to GT but pred(env obs) differs, the observation +adapter (image orientation/cameras/state) is the remaining gap. +""" + +from __future__ import annotations + +import argparse + +import numpy as np +import torch +from scipy.spatial.transform import Rotation + + +@torch.no_grad() +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--checkpoint", default="outputs/libero_ft2/step_15000") + parser.add_argument("--suite", default="libero_spatial") + parser.add_argument("--tasks", type=int, default=3) + args = parser.parse_args() + + from lerobot.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata + from lerobot.envs.factory import make_env, make_env_config + from transformers import AutoTokenizer + from tinyvla.data.mixture import CanonicalSource + from tinyvla.modeling_tinyvla import TinyVLAPolicy + + policy = TinyVLAPolicy.from_pretrained(args.checkpoint).cuda().eval() + cfg = policy.config + tok = AutoTokenizer.from_pretrained(cfg.lm_model_name) + + meta = LeRobotDatasetMetadata("HuggingFaceVLA/libero") + ds = LeRobotDataset( + "HuggingFaceVLA/libero", + delta_timestamps={"action": [t / meta.fps for t in range(cfg.chunk_size)]}, + video_backend="torchcodec", + ) + src = CanonicalSource(ds, 2, cfg.image_size, cfg.max_state_dim, cfg.max_action_dim) + s_stats = meta.stats["observation.state"] + s_mean = torch.as_tensor(s_stats["mean"]).flatten().float() + s_std = torch.as_tensor(s_stats["std"]).flatten().float().clamp(min=1e-6) + + env_cfg = make_env_config("libero", task=args.suite) + task_envs = make_env(env_cfg, n_envs=1)[args.suite] + env_by_task = {} + for tid, env in task_envs.items(): + desc = env.get_attr("task_description")[0] + env_by_task[desc.strip().lower()] = (tid, env) + + eps_meta = ds.meta.episodes + first_ep_by_task = {} + for ep in range(ds.num_episodes): + start = int(eps_meta["dataset_from_index"][ep]) + task = ds[start]["task"].strip().lower() + if task in env_by_task and task not in first_ep_by_task: + first_ep_by_task[task] = ep + + def tok_batch(task_text): + t = tok([task_text], padding=True, truncation=True, + max_length=cfg.tokenizer_max_length, return_tensors="pt") + return t["input_ids"].cuda(), t["attention_mask"].bool().cuda() + + def env_to_batch(obs, task_text): + imgs = {} + for slot, key in (("cam0", "image"), ("cam1", "image2")): + x = torch.as_tensor(np.asarray(obs["pixels"][key]))[0].flip(0).flip(1) + x = x.permute(2, 0, 1).float() / 255.0 + x = torch.nn.functional.interpolate(x[None], size=(cfg.image_size, cfg.image_size), + mode="bilinear", align_corners=False)[0] + imgs[slot] = x + rs = obs["robot_state"] + pos = np.asarray(rs["eef"]["pos"]).flatten() + quat = np.asarray(rs["eef"]["quat"]).flatten() + rotvec = Rotation.from_quat(quat).as_rotvec() + if rotvec[0] < 0: + th = np.linalg.norm(rotvec) + rotvec = rotvec * (th - 2 * np.pi) / th + grip = np.asarray(rs["gripper"]["qpos"]).flatten() + state = torch.tensor(np.concatenate([pos, rotvec, grip]), dtype=torch.float32) + state = (state - s_mean) / s_std + state = torch.nn.functional.pad(state, (0, cfg.max_state_dim - state.shape[-1])) + ids, mask = tok_batch(task_text) + return { + "observation.images.cam0": imgs["cam0"][None].cuda(), + "observation.images.cam1": imgs["cam1"][None].cuda(), + "observation.state": state[None].cuda(), + "observation.language.tokens": ids, + "observation.language.attention_mask": mask, + "embodiment_id": torch.tensor([2], device="cuda"), + } + + def ds_to_batch(item): + ids, mask = tok_batch(item.pop("task")) + b = {k: v[None].cuda() for k, v in item.items() if torch.is_tensor(v)} + b["observation.language.tokens"] = ids + b["observation.language.attention_mask"] = mask + return b + + for task, ep in list(first_ep_by_task.items())[: args.tasks]: + tid, env = env_by_task[task] + obs, _ = env.reset(seed=0) + start = int(eps_meta["dataset_from_index"][ep]) + item = src[start] + gt = item["action"].clone()[None].cuda() + + env_b = env_to_batch(obs, task) + ds_b = ds_to_batch(dict(item)) + + torch.manual_seed(0) + pred_env = policy.predict_action_chunk(env_b) + torch.manual_seed(0) + pred_ds = policy.predict_action_chunk(ds_b) + + m = item["action_dim_mask"] + d_env_gt = ((pred_env[0, :, m].cpu() - gt[0, :, m].cpu()) ** 2).mean().sqrt().item() + d_ds_gt = ((pred_ds[0, :, m].cpu() - gt[0, :, m].cpu()) ** 2).mean().sqrt().item() + d_env_ds = ((pred_env[0, :, m] - pred_ds[0, :, m]) ** 2).mean().sqrt().item() + + # also compare the raw inputs + img_gap = (env_b["observation.images.cam0"] - ds_b["observation.images.cam0"]).abs().mean().item() + st_gap = (env_b["observation.state"] - ds_b["observation.state"]).abs().max().item() + print(f"task {tid}: RMSE(pred_env, GT)={d_env_gt:.3f} | RMSE(pred_ds, GT)={d_ds_gt:.3f} " + f"| RMSE(pred_env, pred_ds)={d_env_ds:.3f} | img L1 gap={img_gap:.4f} | state max gap={st_gap:.2f}") + + +if __name__ == "__main__": + main() diff --git a/tinyvla2/scripts/download_subsets.py b/tinyvla2/scripts/download_subsets.py new file mode 100644 index 0000000000000000000000000000000000000000..d013e0fbdac1e92850f0990b99e3082a4f46a8d0 --- /dev/null +++ b/tinyvla2/scripts/download_subsets.py @@ -0,0 +1,77 @@ +#!/usr/bin/env python +"""Download Stage-2 dataset subsets under a hard disk budget. + +Usage: + python scripts/download_subsets.py --root ~/tinyvla_data [--dry-run] + +Downloads metadata first, measures actual on-disk size of a small episode +sample, then extrapolates before committing to the full subset download. +""" + +from __future__ import annotations + +import argparse +import shutil +from pathlib import Path + +DISK_BUDGET_GB = 400 + +# (repo_id, num_episodes or None for all) +SUBSETS = [ + ("HuggingFaceVLA/community_dataset_v1", None), + ("nvidia/BridgeData2_LeRobot_v3", 10_000), + ("IPEC-COMMUNITY/fractal20220817_data_lerobot", 8_000), +] + + +def free_gb(path: Path) -> float: + return shutil.disk_usage(path).free / 1e9 + + +def used_gb(path: Path) -> float: + return sum(f.stat().st_size for f in path.rglob("*") if f.is_file()) / 1e9 + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--root", type=Path, default=Path.home() / "tinyvla_data") + parser.add_argument("--dry-run", action="store_true") + parser.add_argument("--only", type=str, default=None, help="substring filter on repo_id") + args = parser.parse_args() + + args.root.mkdir(parents=True, exist_ok=True) + + from lerobot.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata + + for repo_id, n_eps in SUBSETS: + if args.only and args.only not in repo_id: + continue + print(f"\n=== {repo_id} (episodes: {n_eps or 'all'}) ===") + meta = LeRobotDatasetMetadata(repo_id) + total_eps = meta.total_episodes + print(f"total episodes: {total_eps}, fps: {meta.fps}") + + episodes = list(range(min(n_eps, total_eps))) if n_eps else None + + # probe with 1% to estimate size + probe_n = max(10, (len(episodes) if episodes else total_eps) // 100) + probe_dir = args.root / "_probe" / repo_id.replace("/", "__") + if not args.dry_run and not probe_dir.exists(): + LeRobotDataset(repo_id, root=probe_dir, episodes=list(range(probe_n))) + probe_gb = used_gb(probe_dir) + est_gb = probe_gb / probe_n * (len(episodes) if episodes else total_eps) + print(f"probe: {probe_n} eps = {probe_gb:.2f}GB -> estimated full subset {est_gb:.0f}GB") + if est_gb > free_gb(args.root) - 100 or used_gb(args.root) + est_gb > DISK_BUDGET_GB: + print(f"SKIP {repo_id}: would exceed budget ({DISK_BUDGET_GB}GB) or disk") + continue + + if args.dry_run: + continue + + target = args.root / repo_id.replace("/", "__") + LeRobotDataset(repo_id, root=target, episodes=episodes) + print(f"downloaded {repo_id}: {used_gb(target):.1f}GB, total used {used_gb(args.root):.1f}GB") + + +if __name__ == "__main__": + main() diff --git a/tinyvla2/scripts/eval_canonical.py b/tinyvla2/scripts/eval_canonical.py new file mode 100644 index 0000000000000000000000000000000000000000..4a6e27bbf1301151f24b47c4aa17eb8be3e64ade --- /dev/null +++ b/tinyvla2/scripts/eval_canonical.py @@ -0,0 +1,186 @@ +#!/usr/bin/env python +"""Physical-space canonical-action eval (variants A/B/C, TinyVLA-2 headline metric). + +Compares A/B/C in ONE fair space: canonical base-frame EE deltas, unnormalized to +physical units — position error (mm), rotation error (deg), gripper error ([0,1]). + +- B/C predict canonical directly → unnormalize with dataset canonical stats. +- A predicts NATIVE actions → for SO101 map joint predictions through FK to EE deltas; + for EE-native sources A's native deltas are already comparable (convention aside). +- Held-out embodiment: A/B have no ID row → oracle over trained IDs (best), making any + C win conservative. C uses the written descriptor (zero-shot by construction). + +Usage: + python scripts/eval_canonical.py --checkpoint outputs/tv2_C_morph_canon/final \ + --dataset heldout_jaco_play --root ~/tinyvla_data/heldout/heldout_jaco_play \ + --morph-key jaco --episodes 20 +""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +import numpy as np +import torch + + +@torch.no_grad() +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--checkpoint", type=Path, required=True) + ap.add_argument("--dataset", required=True) + ap.add_argument("--root", required=True) + ap.add_argument("--morph-key", default=None, help="descriptor key for variant C zero-shot") + ap.add_argument("--oracle-ids", type=int, default=8, help="A/B: try IDs 0..N-1, report best") + ap.add_argument("--episodes", type=int, default=20) + ap.add_argument("--ep-start", type=int, default=0, help="first episode index (few-shot: eval on held-out test split after FT episodes)") + ap.add_argument("--stride", type=int, default=30) + args = ap.parse_args() + + import yaml + from lerobot.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata + from scipy.spatial.transform import Rotation + from transformers import AutoTokenizer + + from tinyvla.data.canonical import CanonicalChunkStore, quantile_normalize + from tinyvla.modeling_tinyvla import TinyVLAPolicy + from tinyvla.modules.embodiment import MORPH_FIELDS + + policy = TinyVLAPolicy.from_pretrained(args.checkpoint).cuda().eval() + cfg = policy.config + tok = AutoTokenizer.from_pretrained(cfg.lm_model_name) + chunk = cfg.chunk_size + + meta = LeRobotDatasetMetadata(args.dataset, root=args.root) + ds = LeRobotDataset(args.dataset, root=args.root, + delta_timestamps={"action": [t / meta.fps for t in range(chunk)]}, + video_backend="torchcodec") + store = CanonicalChunkStore(args.dataset, src_fps=ds.fps, chunk=chunk) + stats = store.compute_stats() + q01, q99 = np.asarray(stats["q01"]), np.asarray(stats["q99"]) + + # morphology descriptor (variant C) + morph = None + if cfg.conditioning == "morph" and args.morph_key: + raw = yaml.safe_load(open("configs/morphology/descriptors.yaml"))[args.morph_key] + _sc = {"arm_dof": 0.1, "reach_m": 2, "gripper_width_m": 10, "num_cameras": 1 / 3, + "control_hz": 1 / 30, "joint_lo_mean": 1 / 3.1416, "joint_hi_mean": 1 / 3.1416, + "workspace_x": 2, "workspace_y": 2, "workspace_z": 2, "payload_kg": 0.2} + morph = torch.tensor([raw.get(f, 0) * _sc.get(f, 1) for f in MORPH_FIELDS], + dtype=torch.float32) + + image_keys = sorted(k for k in ds.meta.features if k.startswith("observation.images")) + + def to_batch(item, emb_id, morph_vec): + import re + from tinyvla.data.mixture import PRIMARY_RE, WRIST_RE + prim = next((k for k in image_keys if PRIMARY_RE.search(k)), image_keys[0]) + wr = next((k for k in image_keys if k != prim and WRIST_RE.search(k)), None) + + def img(x): + if x.shape[-2:] != (cfg.image_size, cfg.image_size): + x = torch.nn.functional.interpolate(x[None], size=(cfg.image_size, cfg.image_size), + mode="bilinear", align_corners=False)[0] + return x + cam0 = img(item[prim]) + cam1 = img(item[wr]) if wr else torch.zeros_like(cam0) + state = item["observation.state"].float() + state = torch.nn.functional.pad(state, (0, cfg.max_state_dim - state.shape[-1])) + t = tok([item.get("task") or ""], padding=True, truncation=True, + max_length=cfg.tokenizer_max_length, return_tensors="pt") + b = {"observation.images.cam0": cam0[None].cuda(), + "observation.images.cam1": cam1[None].cuda(), + "observation.state": state[None].cuda(), + "observation.language.tokens": t["input_ids"].cuda(), + "observation.language.attention_mask": t["attention_mask"].bool().cuda(), + "embodiment_id": torch.tensor([emb_id], device="cuda")} + if morph_vec is not None: + b["morphology"] = morph_vec[None].cuda() + return b + + # native action stats (for variant A unnormalization) — physical EE deltas + native_stats = ds.meta.stats.get("action", {}) + nat_mean = np.asarray(native_stats.get("mean", np.zeros(7))) + nat_std = np.asarray(native_stats.get("std", np.ones(7))) + is_native = cfg.action_space == "native" + + def _unnorm_canon(x): + span = np.maximum(q99 - q01, 0.01 * np.median(np.abs(np.concatenate([q01, q99])) + 1e-6)) + mid = 0.5 * (q01 + q99) + return x[:, :7] * span / 2 + mid + + def _integrate(deltas): + """(T,6+) physical per-step EE deltas -> (T,3) cumulative positions, + list of cumulative rotations. Amplifies per-step differences into a + trajectory with real dynamic range (per-step motion is tiny).""" + pos = np.cumsum(deltas[:, :3], axis=0) + R = Rotation.identity() + rots = [] + for k in range(len(deltas)): + R = Rotation.from_rotvec(deltas[k, 3:6]) * R + rots.append(R) + return pos, rots + + def canonical_phys_err(pred_norm, gt_canon): + """Integrated-trajectory error (the per-step delta floor ~3.6mm has no + dynamic range — see control). Returns: + ep_mm : endpoint position error over the 5s chunk (mm) + path_mm: mean cumulative-position error along the chunk (mm) + ep_deg : endpoint cumulative-rotation error (deg) + """ + gu = _unnorm_canon(gt_canon) + if is_native: + d = pred_norm.shape[-1] + pu = pred_norm[:, : min(d, 7)] * nat_std[: min(d, 7)] + nat_mean[: min(d, 7)] + if pu.shape[-1] < 7: + pu = np.concatenate([pu, gu[:, pu.shape[-1]:7]], axis=1) + else: + pu = _unnorm_canon(pred_norm) + gp, gr = _integrate(gu) + pp, pr = _integrate(pu) + ep_mm = np.linalg.norm(pp[-1] - gp[-1]) * 1000 + path_mm = np.linalg.norm(pp - gp, axis=1).mean() * 1000 + ep_deg = np.degrees((pr[-1] * gr[-1].inv()).magnitude()) + return ep_mm, path_mm, ep_deg + pos_mm = np.linalg.norm(pu[:, :3] - gu[:, :3], axis=1).mean() * 1000 + # rotation error: geodesic between rotvec deltas + rp = Rotation.from_rotvec(pu[:, 3:6]) + rg = Rotation.from_rotvec(gu[:, 3:6]) + rot_deg = np.degrees((rp * rg.inv()).magnitude()).mean() + grip = np.abs(pu[:, 6] - gu[:, 6]).mean() + return pos_mm, rot_deg, grip + + eps = list(range(args.ep_start, min(args.ep_start + args.episodes, ds.num_episodes))) + # candidate embodiment conditionings + if cfg.conditioning == "morph": + candidates = [("morph", morph)] + else: + candidates = [(f"id{i}", i) for i in range(args.oracle_ids)] + + best = None + for label, cand in candidates: + errs = [] + for ep in eps: + start = int(ds.meta.episodes["dataset_from_index"][ep]) + end = int(ds.meta.episodes["dataset_to_index"][ep]) + for idx in range(start, end - 1, args.stride): + gt = quantile_normalize(store.chunk_for(ep, idx - start), q01, q99) + if cfg.conditioning == "morph": + b = to_batch(ds[idx], 0, cand) + else: + b = to_batch(ds[idx], cand, None) + pred = policy.predict_action_chunk(b)[0].cpu().numpy() + errs.append(canonical_phys_err(pred, gt)) + errs = np.array(errs) + m = errs.mean(0) + if best is None or m[0] < best[1][0]: + best = (label, m) + print(f" {label}: endpoint {m[0]:.1f}mm path {m[1]:.1f}mm rot {m[2]:.1f}deg") + + print(f"\n=== {args.dataset} | {args.checkpoint.name} | cond={cfg.conditioning} ===") + print(f"BEST ({best[0]}): endpoint {best[1][0]:.1f}mm path {best[1][1]:.1f}mm rot {best[1][2]:.1f}deg") + + +if __name__ == "__main__": + main() diff --git a/tinyvla2/scripts/eval_demoonly_lekiwi.py b/tinyvla2/scripts/eval_demoonly_lekiwi.py new file mode 100644 index 0000000000000000000000000000000000000000..52cf685e94ca3c5a99223696657e7c5c928e5972 --- /dev/null +++ b/tinyvla2/scripts/eval_demoonly_lekiwi.py @@ -0,0 +1,139 @@ +#!/usr/bin/env python +"""Auto-eval for in-context demo conditioning on held-out LeKiwi. + +The key test: give the model K=3 REAL (obs, action) example pairs from LeKiwi's +OWN data at test time — NO gradient update, NO fine-tuning — and see if that beats +zero-shot with no examples, wrong examples, or the numeric/text descriptor alone. +""" + +from __future__ import annotations + +import numpy as np +import torch +import yaml +from lerobot.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata +from transformers import AutoTokenizer + +from tinyvla.data.canonical import CanonicalChunkStore, quantile_normalize +from tinyvla.modeling_tinyvla import TinyVLAPolicy + +CKPT = "outputs/tv2_demoonly/final" +DS = "lekiwi_cleanup" +ROOT = f"/home/alexw/tinyvla_data/lekiwi/{DS}" +K = 4 + + +@torch.no_grad() +def main(): + pol = TinyVLAPolicy.from_pretrained(CKPT).cuda().eval() + cfg = pol.config + tok = AutoTokenizer.from_pretrained(cfg.lm_model_name) + prompts = yaml.safe_load(open("configs/morphology/robot_prompts.yaml")) + + m = LeRobotDatasetMetadata(DS, root=ROOT) + ds = LeRobotDataset(DS, root=ROOT, delta_timestamps={"action": [t / m.fps for t in range(50)]}, + video_backend="torchcodec") + store = CanonicalChunkStore(DS, src_fps=m.fps, chunk=50) + st = store.compute_stats() + q01, q99 = np.asarray(st["q01"]), np.asarray(st["q99"]) + span = np.maximum(q99 - q01, 0.01 * np.median(np.abs(np.concatenate([q01, q99])) + 1e-6)) + mid = 0.5 * (q01 + q99) + + imk = sorted(k for k in ds.meta.features if k.startswith("observation.images")) + prim = next((k for k in imk if "front" in k or "base" in k), imk[0]) + + def integ(d): + return np.cumsum(d[:, :3], 0) + + def img_at(idx, ep, ep_start): + it = ds[idx] + im = torch.nn.functional.interpolate(it[prim][None], size=(256, 256), mode="bilinear")[0] + raw = quantile_normalize(store.chunk_for(ep, idx - ep_start), q01, q99) + act = torch.from_numpy(raw[:, :7]).float() + act = torch.nn.functional.pad(act, (0, cfg.max_action_dim - 7)) + return im, act + + # REAL LeKiwi demo pairs from early episodes (support pool), disjoint from test + support_pool_eps = range(0, min(20, m.total_episodes - 15)) + test = range(max(0, m.total_episodes - 15), m.total_episodes) + + def get_real_support(k): + imgs, acts = [], [] + for _ in range(k): + ep = int(np.random.choice(list(support_pool_eps))) + s = int(m.episodes["dataset_from_index"][ep]) + e = int(m.episodes["dataset_to_index"][ep]) + idx = int(np.random.randint(s, max(s + 1, e - 1))) + im, act = img_at(idx, ep, s) + imgs.append(im) + acts.append(act) + return torch.stack(imgs), torch.stack(acts) + + def get_wrong_support(k): + # random noise images + random actions — a garbage support set control + return torch.rand(k, 3, 256, 256), torch.randn(k, 50, cfg.max_action_dim) * 0.3 + + + # control: demos from a DIFFERENT ROBOT (jaco) — if the model truly extracts + # robot-specific dynamics from exemplars, these should be worse than LeKiwi's own + def get_otherrobot_support(k): + from lerobot.datasets.lerobot_dataset import LeRobotDataset as _LD, LeRobotDatasetMetadata as _LM + from tinyvla.data.canonical import CanonicalChunkStore as _CS + jr = "/home/alexw/tinyvla_data/heldout/heldout_jaco_play" + jm = _LM("heldout_jaco_play", root=jr) + jds = _LD("heldout_jaco_play", root=jr, delta_timestamps={"action": [t / jm.fps for t in range(50)]}, video_backend="torchcodec") + jstore = _CS("heldout_jaco_play", src_fps=jm.fps, chunk=50) + jst = jstore.compute_stats(200) + jq01, jq99 = np.asarray(jst["q01"]), np.asarray(jst["q99"]) + jimk = sorted(kk for kk in jds.meta.features if kk.startswith("observation.images"))[0] + imgs, acts = [], [] + for _ in range(k): + ep = int(np.random.randint(0, 50)) + s = int(jm.episodes["dataset_from_index"][ep]); e = int(jm.episodes["dataset_to_index"][ep]) + idx = int(np.random.randint(s, max(s + 1, e - 1))) + it = jds[idx] + im = torch.nn.functional.interpolate(it[jimk][None], size=(256, 256), mode="bilinear")[0] + raw = quantile_normalize(jstore.chunk_for(ep, idx - s), jq01, jq99) + act = torch.nn.functional.pad(torch.from_numpy(raw[:, :7]).float(), (0, cfg.max_action_dim - 7)) + imgs.append(im); acts.append(act) + return torch.stack(imgs), torch.stack(acts) + + def run(support_fn, desc_text): + ids_t = tok([desc_text] if desc_text else [""], padding="max_length", truncation=True, + max_length=cfg.morph_text_max_len, return_tensors="pt") + errs = [] + for ep in test: + s = int(m.episodes["dataset_from_index"][ep]) + e = int(m.episodes["dataset_to_index"][ep]) + for idx in range(s, e - 1, 30): + item = ds[idx] + task = item.get("task") or "" + t = tok([task], padding=True, truncation=True, max_length=48, return_tensors="pt") + img = torch.nn.functional.interpolate(item[prim][None], size=(256, 256), mode="bilinear")[0] + stt = torch.nn.functional.pad(item["observation.state"].float()[:6], (0, cfg.max_state_dim - 6)) + sup_img, sup_act = support_fn(K) + b = {"observation.images.cam0": img[None].cuda(), + "observation.images.cam1": torch.zeros_like(img)[None].cuda(), + "observation.state": stt[None].cuda(), + "observation.language.tokens": t["input_ids"].cuda(), + "observation.language.attention_mask": t["attention_mask"].bool().cuda(), + "morph_text_ids": ids_t["input_ids"].cuda(), + "morph_text_mask": ids_t["attention_mask"].bool().cuda(), + "support_images": sup_img[None].cuda(), "support_actions": sup_act[None].cuda(), + "embodiment_id": torch.tensor([0]).cuda()} + with torch.autocast("cuda", torch.bfloat16): + pr = pol.predict_action_chunk(b)[0].cpu().float().numpy() + gu = quantile_normalize(store.chunk_for(ep, idx - s), q01, q99)[:, :7] * span / 2 + mid + pu = pr[:, :7] * span / 2 + mid + errs.append(np.linalg.norm(integ(pu)[-1] - integ(gu)[-1]) * 1000) + return np.mean(errs) + + print("=== DEMO-ONLY (demos are the ONLY identity cue, no FT) on held-out LeKiwi ===") + print(f"REAL LeKiwi demos (K={K}) endpoint {run(get_real_support, None):.1f}mm") + print(f"GARBAGE/noise demos endpoint {run(get_wrong_support, None):.1f}mm") + print(f"OTHER-ROBOT (jaco) demos endpoint {run(get_otherrobot_support, None):.1f}mm") + print("\nprior baselines: C-diverse 295mm | C-qwen-morph desc=none 294mm | v1 demo(same-task) real=290.6 garbage=288.5 (dead)") + + +if __name__ == "__main__": + main() diff --git a/tinyvla2/scripts/eval_incontext_lekiwi.py b/tinyvla2/scripts/eval_incontext_lekiwi.py new file mode 100644 index 0000000000000000000000000000000000000000..41fffdc7385350d008c50af54356b7bb85aca3fe --- /dev/null +++ b/tinyvla2/scripts/eval_incontext_lekiwi.py @@ -0,0 +1,115 @@ +#!/usr/bin/env python +"""Auto-eval for in-context demo conditioning on held-out LeKiwi. + +The key test: give the model K=3 REAL (obs, action) example pairs from LeKiwi's +OWN data at test time — NO gradient update, NO fine-tuning — and see if that beats +zero-shot with no examples, wrong examples, or the numeric/text descriptor alone. +""" + +from __future__ import annotations + +import numpy as np +import torch +import yaml +from lerobot.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata +from transformers import AutoTokenizer + +from tinyvla.data.canonical import CanonicalChunkStore, quantile_normalize +from tinyvla.modeling_tinyvla import TinyVLAPolicy + +CKPT = "outputs/tv2_C_incontext/final" +DS = "lekiwi_cleanup" +ROOT = f"/home/alexw/tinyvla_data/lekiwi/{DS}" +K = 3 + + +@torch.no_grad() +def main(): + pol = TinyVLAPolicy.from_pretrained(CKPT).cuda().eval() + cfg = pol.config + tok = AutoTokenizer.from_pretrained(cfg.lm_model_name) + prompts = yaml.safe_load(open("configs/morphology/robot_prompts.yaml")) + + m = LeRobotDatasetMetadata(DS, root=ROOT) + ds = LeRobotDataset(DS, root=ROOT, delta_timestamps={"action": [t / m.fps for t in range(50)]}, + video_backend="torchcodec") + store = CanonicalChunkStore(DS, src_fps=m.fps, chunk=50) + st = store.compute_stats() + q01, q99 = np.asarray(st["q01"]), np.asarray(st["q99"]) + span = np.maximum(q99 - q01, 0.01 * np.median(np.abs(np.concatenate([q01, q99])) + 1e-6)) + mid = 0.5 * (q01 + q99) + + imk = sorted(k for k in ds.meta.features if k.startswith("observation.images")) + prim = next((k for k in imk if "front" in k or "base" in k), imk[0]) + + def integ(d): + return np.cumsum(d[:, :3], 0) + + def img_at(idx, ep, ep_start): + it = ds[idx] + im = torch.nn.functional.interpolate(it[prim][None], size=(256, 256), mode="bilinear")[0] + raw = quantile_normalize(store.chunk_for(ep, idx - ep_start), q01, q99) + act = torch.from_numpy(raw[:, :7]).float() + act = torch.nn.functional.pad(act, (0, cfg.max_action_dim - 7)) + return im, act + + # REAL LeKiwi demo pairs from early episodes (support pool), disjoint from test + support_pool_eps = range(0, min(20, m.total_episodes - 15)) + test = range(max(0, m.total_episodes - 15), m.total_episodes) + + def get_real_support(k): + imgs, acts = [], [] + for _ in range(k): + ep = int(np.random.choice(list(support_pool_eps))) + s = int(m.episodes["dataset_from_index"][ep]) + e = int(m.episodes["dataset_to_index"][ep]) + idx = int(np.random.randint(s, max(s + 1, e - 1))) + im, act = img_at(idx, ep, s) + imgs.append(im) + acts.append(act) + return torch.stack(imgs), torch.stack(acts) + + def get_wrong_support(k): + # random noise images + random actions — a garbage support set control + return torch.rand(k, 3, 256, 256), torch.randn(k, 50, cfg.max_action_dim) * 0.3 + + def run(support_fn, desc_text): + ids_t = tok([desc_text] if desc_text else [""], padding="max_length", truncation=True, + max_length=cfg.morph_text_max_len, return_tensors="pt") + errs = [] + for ep in test: + s = int(m.episodes["dataset_from_index"][ep]) + e = int(m.episodes["dataset_to_index"][ep]) + for idx in range(s, e - 1, 30): + item = ds[idx] + task = item.get("task") or "" + t = tok([task], padding=True, truncation=True, max_length=48, return_tensors="pt") + img = torch.nn.functional.interpolate(item[prim][None], size=(256, 256), mode="bilinear")[0] + stt = torch.nn.functional.pad(item["observation.state"].float()[:6], (0, cfg.max_state_dim - 6)) + sup_img, sup_act = support_fn(K) + b = {"observation.images.cam0": img[None].cuda(), + "observation.images.cam1": torch.zeros_like(img)[None].cuda(), + "observation.state": stt[None].cuda(), + "observation.language.tokens": t["input_ids"].cuda(), + "observation.language.attention_mask": t["attention_mask"].bool().cuda(), + "morph_text_ids": ids_t["input_ids"].cuda(), + "morph_text_mask": ids_t["attention_mask"].bool().cuda(), + "support_images": sup_img[None].cuda(), "support_actions": sup_act[None].cuda(), + "embodiment_id": torch.tensor([0]).cuda()} + with torch.autocast("cuda", torch.bfloat16): + pr = pol.predict_action_chunk(b)[0].cpu().float().numpy() + gu = quantile_normalize(store.chunk_for(ep, idx - s), q01, q99)[:, :7] * span / 2 + mid + pu = pr[:, :7] * span / 2 + mid + errs.append(np.linalg.norm(integ(pu)[-1] - integ(gu)[-1]) * 1000) + return np.mean(errs) + + print("=== C-incontext (demo conditioning, no FT) on held-out LeKiwi ===") + print(f"REAL LeKiwi demos (K={K}) + desc=lekiwi endpoint {run(get_real_support, prompts['lekiwi']):.1f}mm") + print(f"REAL LeKiwi demos (K={K}) + desc=none endpoint {run(get_real_support, None):.1f}mm") + print(f"WRONG/garbage demos + desc=lekiwi endpoint {run(get_wrong_support, prompts['lekiwi']):.1f}mm") + print(f"WRONG/garbage demos + desc=none endpoint {run(get_wrong_support, None):.1f}mm") + print("\nprior baselines (no demo conditioning): C-diverse 295mm | C-qwen-morph desc=none 294mm") + + +if __name__ == "__main__": + main() diff --git a/tinyvla2/scripts/eval_incontext_v2_lekiwi.py b/tinyvla2/scripts/eval_incontext_v2_lekiwi.py new file mode 100644 index 0000000000000000000000000000000000000000..d455fc5b1d4c900954744b760335eaeeb01ec489 --- /dev/null +++ b/tinyvla2/scripts/eval_incontext_v2_lekiwi.py @@ -0,0 +1,141 @@ +#!/usr/bin/env python +"""Auto-eval for in-context demo conditioning on held-out LeKiwi. + +The key test: give the model K=3 REAL (obs, action) example pairs from LeKiwi's +OWN data at test time — NO gradient update, NO fine-tuning — and see if that beats +zero-shot with no examples, wrong examples, or the numeric/text descriptor alone. +""" + +from __future__ import annotations + +import numpy as np +import torch +import yaml +from lerobot.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata +from transformers import AutoTokenizer + +from tinyvla.data.canonical import CanonicalChunkStore, quantile_normalize +from tinyvla.modeling_tinyvla import TinyVLAPolicy + +CKPT = "outputs/tv2_C_incontext_v2/final" +DS = "lekiwi_cleanup" +ROOT = f"/home/alexw/tinyvla_data/lekiwi/{DS}" +K = 3 + + +@torch.no_grad() +def main(): + pol = TinyVLAPolicy.from_pretrained(CKPT).cuda().eval() + cfg = pol.config + tok = AutoTokenizer.from_pretrained(cfg.lm_model_name) + prompts = yaml.safe_load(open("configs/morphology/robot_prompts.yaml")) + + m = LeRobotDatasetMetadata(DS, root=ROOT) + ds = LeRobotDataset(DS, root=ROOT, delta_timestamps={"action": [t / m.fps for t in range(50)]}, + video_backend="torchcodec") + store = CanonicalChunkStore(DS, src_fps=m.fps, chunk=50) + st = store.compute_stats() + q01, q99 = np.asarray(st["q01"]), np.asarray(st["q99"]) + span = np.maximum(q99 - q01, 0.01 * np.median(np.abs(np.concatenate([q01, q99])) + 1e-6)) + mid = 0.5 * (q01 + q99) + + imk = sorted(k for k in ds.meta.features if k.startswith("observation.images")) + prim = next((k for k in imk if "front" in k or "base" in k), imk[0]) + + def integ(d): + return np.cumsum(d[:, :3], 0) + + def img_at(idx, ep, ep_start): + it = ds[idx] + im = torch.nn.functional.interpolate(it[prim][None], size=(256, 256), mode="bilinear")[0] + raw = quantile_normalize(store.chunk_for(ep, idx - ep_start), q01, q99) + act = torch.from_numpy(raw[:, :7]).float() + act = torch.nn.functional.pad(act, (0, cfg.max_action_dim - 7)) + return im, act + + # REAL LeKiwi demo pairs from early episodes (support pool), disjoint from test + support_pool_eps = range(0, min(20, m.total_episodes - 15)) + test = range(max(0, m.total_episodes - 15), m.total_episodes) + + def get_real_support(k): + imgs, acts = [], [] + for _ in range(k): + ep = int(np.random.choice(list(support_pool_eps))) + s = int(m.episodes["dataset_from_index"][ep]) + e = int(m.episodes["dataset_to_index"][ep]) + idx = int(np.random.randint(s, max(s + 1, e - 1))) + im, act = img_at(idx, ep, s) + imgs.append(im) + acts.append(act) + return torch.stack(imgs), torch.stack(acts) + + def get_wrong_support(k): + # random noise images + random actions — a garbage support set control + return torch.rand(k, 3, 256, 256), torch.randn(k, 50, cfg.max_action_dim) * 0.3 + + + # control: demos from a DIFFERENT ROBOT (jaco) — if the model truly extracts + # robot-specific dynamics from exemplars, these should be worse than LeKiwi's own + def get_otherrobot_support(k): + from lerobot.datasets.lerobot_dataset import LeRobotDataset as _LD, LeRobotDatasetMetadata as _LM + from tinyvla.data.canonical import CanonicalChunkStore as _CS + jr = "/home/alexw/tinyvla_data/heldout/heldout_jaco_play" + jm = _LM("heldout_jaco_play", root=jr) + jds = _LD("heldout_jaco_play", root=jr, delta_timestamps={"action": [t / jm.fps for t in range(50)]}, video_backend="torchcodec") + jstore = _CS("heldout_jaco_play", src_fps=jm.fps, chunk=50) + jst = jstore.compute_stats(200) + jq01, jq99 = np.asarray(jst["q01"]), np.asarray(jst["q99"]) + jimk = sorted(kk for kk in jds.meta.features if kk.startswith("observation.images"))[0] + imgs, acts = [], [] + for _ in range(k): + ep = int(np.random.randint(0, 50)) + s = int(jm.episodes["dataset_from_index"][ep]); e = int(jm.episodes["dataset_to_index"][ep]) + idx = int(np.random.randint(s, max(s + 1, e - 1))) + it = jds[idx] + im = torch.nn.functional.interpolate(it[jimk][None], size=(256, 256), mode="bilinear")[0] + raw = quantile_normalize(jstore.chunk_for(ep, idx - s), jq01, jq99) + act = torch.nn.functional.pad(torch.from_numpy(raw[:, :7]).float(), (0, cfg.max_action_dim - 7)) + imgs.append(im); acts.append(act) + return torch.stack(imgs), torch.stack(acts) + + def run(support_fn, desc_text): + ids_t = tok([desc_text] if desc_text else [""], padding="max_length", truncation=True, + max_length=cfg.morph_text_max_len, return_tensors="pt") + errs = [] + for ep in test: + s = int(m.episodes["dataset_from_index"][ep]) + e = int(m.episodes["dataset_to_index"][ep]) + for idx in range(s, e - 1, 30): + item = ds[idx] + task = item.get("task") or "" + t = tok([task], padding=True, truncation=True, max_length=48, return_tensors="pt") + img = torch.nn.functional.interpolate(item[prim][None], size=(256, 256), mode="bilinear")[0] + stt = torch.nn.functional.pad(item["observation.state"].float()[:6], (0, cfg.max_state_dim - 6)) + sup_img, sup_act = support_fn(K) + b = {"observation.images.cam0": img[None].cuda(), + "observation.images.cam1": torch.zeros_like(img)[None].cuda(), + "observation.state": stt[None].cuda(), + "observation.language.tokens": t["input_ids"].cuda(), + "observation.language.attention_mask": t["attention_mask"].bool().cuda(), + "morph_text_ids": ids_t["input_ids"].cuda(), + "morph_text_mask": ids_t["attention_mask"].bool().cuda(), + "support_images": sup_img[None].cuda(), "support_actions": sup_act[None].cuda(), + "embodiment_id": torch.tensor([0]).cuda()} + with torch.autocast("cuda", torch.bfloat16): + pr = pol.predict_action_chunk(b)[0].cpu().float().numpy() + gu = quantile_normalize(store.chunk_for(ep, idx - s), q01, q99)[:, :7] * span / 2 + mid + pu = pr[:, :7] * span / 2 + mid + errs.append(np.linalg.norm(integ(pu)[-1] - integ(gu)[-1]) * 1000) + return np.mean(errs) + + print("=== C-incontext-v2 (cross-task demos, no FT) on held-out LeKiwi ===") + print(f"REAL LeKiwi demos (K={K}) + desc=lekiwi endpoint {run(get_real_support, prompts['lekiwi']):.1f}mm") + print(f"REAL LeKiwi demos (K={K}) + desc=none endpoint {run(get_real_support, None):.1f}mm") + print(f"WRONG/garbage demos + desc=lekiwi endpoint {run(get_wrong_support, prompts['lekiwi']):.1f}mm") + print(f"WRONG/garbage demos + desc=none endpoint {run(get_wrong_support, None):.1f}mm") + print(f"OTHER-ROBOT (jaco) demos + desc=none endpoint {run(get_otherrobot_support, None):.1f}mm") + print("\nprior baselines: C-diverse 295mm | C-qwen-morph desc=none 294mm | v1 demo(same-task) real=290.6 garbage=288.5 (dead)") + + +if __name__ == "__main__": + main() diff --git a/tinyvla2/scripts/eval_incontext_v3_lekiwi.py b/tinyvla2/scripts/eval_incontext_v3_lekiwi.py new file mode 100644 index 0000000000000000000000000000000000000000..98c2e9ca8f51313156f958ab1cd1bbf1c213d782 --- /dev/null +++ b/tinyvla2/scripts/eval_incontext_v3_lekiwi.py @@ -0,0 +1,141 @@ +#!/usr/bin/env python +"""Auto-eval for in-context demo conditioning on held-out LeKiwi. + +The key test: give the model K=3 REAL (obs, action) example pairs from LeKiwi's +OWN data at test time — NO gradient update, NO fine-tuning — and see if that beats +zero-shot with no examples, wrong examples, or the numeric/text descriptor alone. +""" + +from __future__ import annotations + +import numpy as np +import torch +import yaml +from lerobot.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata +from transformers import AutoTokenizer + +from tinyvla.data.canonical import CanonicalChunkStore, quantile_normalize +from tinyvla.modeling_tinyvla import TinyVLAPolicy + +CKPT = "outputs/tv2_C_incontext_v3/final" +DS = "lekiwi_cleanup" +ROOT = f"/home/alexw/tinyvla_data/lekiwi/{DS}" +K = 3 + + +@torch.no_grad() +def main(): + pol = TinyVLAPolicy.from_pretrained(CKPT).cuda().eval() + cfg = pol.config + tok = AutoTokenizer.from_pretrained(cfg.lm_model_name) + prompts = yaml.safe_load(open("configs/morphology/robot_prompts.yaml")) + + m = LeRobotDatasetMetadata(DS, root=ROOT) + ds = LeRobotDataset(DS, root=ROOT, delta_timestamps={"action": [t / m.fps for t in range(50)]}, + video_backend="torchcodec") + store = CanonicalChunkStore(DS, src_fps=m.fps, chunk=50) + st = store.compute_stats() + q01, q99 = np.asarray(st["q01"]), np.asarray(st["q99"]) + span = np.maximum(q99 - q01, 0.01 * np.median(np.abs(np.concatenate([q01, q99])) + 1e-6)) + mid = 0.5 * (q01 + q99) + + imk = sorted(k for k in ds.meta.features if k.startswith("observation.images")) + prim = next((k for k in imk if "front" in k or "base" in k), imk[0]) + + def integ(d): + return np.cumsum(d[:, :3], 0) + + def img_at(idx, ep, ep_start): + it = ds[idx] + im = torch.nn.functional.interpolate(it[prim][None], size=(256, 256), mode="bilinear")[0] + raw = quantile_normalize(store.chunk_for(ep, idx - ep_start), q01, q99) + act = torch.from_numpy(raw[:, :7]).float() + act = torch.nn.functional.pad(act, (0, cfg.max_action_dim - 7)) + return im, act + + # REAL LeKiwi demo pairs from early episodes (support pool), disjoint from test + support_pool_eps = range(0, min(20, m.total_episodes - 15)) + test = range(max(0, m.total_episodes - 15), m.total_episodes) + + def get_real_support(k): + imgs, acts = [], [] + for _ in range(k): + ep = int(np.random.choice(list(support_pool_eps))) + s = int(m.episodes["dataset_from_index"][ep]) + e = int(m.episodes["dataset_to_index"][ep]) + idx = int(np.random.randint(s, max(s + 1, e - 1))) + im, act = img_at(idx, ep, s) + imgs.append(im) + acts.append(act) + return torch.stack(imgs), torch.stack(acts) + + def get_wrong_support(k): + # random noise images + random actions — a garbage support set control + return torch.rand(k, 3, 256, 256), torch.randn(k, 50, cfg.max_action_dim) * 0.3 + + + # control: demos from a DIFFERENT ROBOT (jaco) — if the model truly extracts + # robot-specific dynamics from exemplars, these should be worse than LeKiwi's own + def get_otherrobot_support(k): + from lerobot.datasets.lerobot_dataset import LeRobotDataset as _LD, LeRobotDatasetMetadata as _LM + from tinyvla.data.canonical import CanonicalChunkStore as _CS + jr = "/home/alexw/tinyvla_data/heldout/heldout_jaco_play" + jm = _LM("heldout_jaco_play", root=jr) + jds = _LD("heldout_jaco_play", root=jr, delta_timestamps={"action": [t / jm.fps for t in range(50)]}, video_backend="torchcodec") + jstore = _CS("heldout_jaco_play", src_fps=jm.fps, chunk=50) + jst = jstore.compute_stats(200) + jq01, jq99 = np.asarray(jst["q01"]), np.asarray(jst["q99"]) + jimk = sorted(kk for kk in jds.meta.features if kk.startswith("observation.images"))[0] + imgs, acts = [], [] + for _ in range(k): + ep = int(np.random.randint(0, 50)) + s = int(jm.episodes["dataset_from_index"][ep]); e = int(jm.episodes["dataset_to_index"][ep]) + idx = int(np.random.randint(s, max(s + 1, e - 1))) + it = jds[idx] + im = torch.nn.functional.interpolate(it[jimk][None], size=(256, 256), mode="bilinear")[0] + raw = quantile_normalize(jstore.chunk_for(ep, idx - s), jq01, jq99) + act = torch.nn.functional.pad(torch.from_numpy(raw[:, :7]).float(), (0, cfg.max_action_dim - 7)) + imgs.append(im); acts.append(act) + return torch.stack(imgs), torch.stack(acts) + + def run(support_fn, desc_text): + ids_t = tok([desc_text] if desc_text else [""], padding="max_length", truncation=True, + max_length=cfg.morph_text_max_len, return_tensors="pt") + errs = [] + for ep in test: + s = int(m.episodes["dataset_from_index"][ep]) + e = int(m.episodes["dataset_to_index"][ep]) + for idx in range(s, e - 1, 30): + item = ds[idx] + task = item.get("task") or "" + t = tok([task], padding=True, truncation=True, max_length=48, return_tensors="pt") + img = torch.nn.functional.interpolate(item[prim][None], size=(256, 256), mode="bilinear")[0] + stt = torch.nn.functional.pad(item["observation.state"].float()[:6], (0, cfg.max_state_dim - 6)) + sup_img, sup_act = support_fn(K) + b = {"observation.images.cam0": img[None].cuda(), + "observation.images.cam1": torch.zeros_like(img)[None].cuda(), + "observation.state": stt[None].cuda(), + "observation.language.tokens": t["input_ids"].cuda(), + "observation.language.attention_mask": t["attention_mask"].bool().cuda(), + "morph_text_ids": ids_t["input_ids"].cuda(), + "morph_text_mask": ids_t["attention_mask"].bool().cuda(), + "support_images": sup_img[None].cuda(), "support_actions": sup_act[None].cuda(), + "embodiment_id": torch.tensor([0]).cuda()} + with torch.autocast("cuda", torch.bfloat16): + pr = pol.predict_action_chunk(b)[0].cpu().float().numpy() + gu = quantile_normalize(store.chunk_for(ep, idx - s), q01, q99)[:, :7] * span / 2 + mid + pu = pr[:, :7] * span / 2 + mid + errs.append(np.linalg.norm(integ(pu)[-1] - integ(gu)[-1]) * 1000) + return np.mean(errs) + + print("=== C-incontext-v3 (high-capacity demos, no FT) on held-out LeKiwi ===") + print(f"REAL LeKiwi demos (K={K}) + desc=lekiwi endpoint {run(get_real_support, prompts['lekiwi']):.1f}mm") + print(f"REAL LeKiwi demos (K={K}) + desc=none endpoint {run(get_real_support, None):.1f}mm") + print(f"WRONG/garbage demos + desc=lekiwi endpoint {run(get_wrong_support, prompts['lekiwi']):.1f}mm") + print(f"WRONG/garbage demos + desc=none endpoint {run(get_wrong_support, None):.1f}mm") + print(f"OTHER-ROBOT (jaco) demos + desc=none endpoint {run(get_otherrobot_support, None):.1f}mm") + print("\nprior baselines: C-diverse 295mm | C-qwen-morph desc=none 294mm | v1 demo(same-task) real=290.6 garbage=288.5 (dead)") + + +if __name__ == "__main__": + main() diff --git a/tinyvla2/scripts/eval_libero.py b/tinyvla2/scripts/eval_libero.py new file mode 100644 index 0000000000000000000000000000000000000000..15b7ee0228ce7b78541a5f3b56b7faccfc4ff410 --- /dev/null +++ b/tinyvla2/scripts/eval_libero.py @@ -0,0 +1,193 @@ +#!/usr/bin/env python +"""Closed-loop LIBERO evaluation for canonical-schema TinyVLA checkpoints. + +Adapts env observations to the canonical schema the policy was trained on +(cam0/cam1, padded normalized state, embodiment_id) and unnormalizes the +predicted actions with the LIBERO dataset stats. + +Usage: + python scripts/eval_libero.py --checkpoint outputs/libero_ft/final \ + --suite libero_spatial --episodes 20 [--refresh-s 1.0] +""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +import torch + + +def make_normalizer(repo_id="HuggingFaceVLA/libero"): + from lerobot.datasets.lerobot_dataset import LeRobotDatasetMetadata + + meta = LeRobotDatasetMetadata(repo_id) + stats = meta.stats + + def norm(key, x): + s = stats[key] + mean = torch.as_tensor(s["mean"], dtype=torch.float32, device=x.device) + std = torch.as_tensor(s["std"], dtype=torch.float32, device=x.device).clamp(min=1e-6) + return (x - mean) / std + + def unnorm_action(x): + s = stats["action"] + mean = torch.as_tensor(s["mean"], dtype=torch.float32, device=x.device) + std = torch.as_tensor(s["std"], dtype=torch.float32, device=x.device) + return x * std + mean + + return norm, unnorm_action, meta + + +@torch.no_grad() +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--checkpoint", type=Path, required=True) + parser.add_argument("--suite", default="libero_spatial", + choices=["libero_spatial", "libero_object", "libero_goal", "libero_10", "libero_90"]) + parser.add_argument("--episodes", type=int, default=20) + parser.add_argument("--embodiment-id", type=int, default=2) + parser.add_argument("--refresh-s", type=float, default=None, + help="if set, refresh the semantic latent only every N seconds (dual-rate mode)") + parser.add_argument("--max-steps", type=int, default=520) + parser.add_argument("--n-action-steps", type=int, default=None, + help="execute only first N actions of each chunk before re-planning") + parser.add_argument("--action-repeat", type=int, default=1, + help="env steps per predicted action (dataset 10fps vs env 20Hz -> 2)") + parser.add_argument("--save-video-dir", type=Path, default=None, + help="save per-episode mp4s of the agentview camera here") + args = parser.parse_args() + + import numpy as np + from lerobot.envs.factory import make_env, make_env_config + from transformers import AutoTokenizer + from tinyvla.modeling_tinyvla import TinyVLAPolicy + + policy = TinyVLAPolicy.from_pretrained(args.checkpoint).cuda().eval() + cfg = policy.config + if args.n_action_steps: + cfg.n_action_steps = args.n_action_steps + tok = AutoTokenizer.from_pretrained(cfg.lm_model_name) + norm, unnorm_action, meta = make_normalizer() + fps = meta.fps + + env_cfg = make_env_config("libero", task=args.suite) + envs_dict = make_env(env_cfg, n_envs=1) + task_envs = envs_dict[args.suite] # {task_id: vec_env} + + from scipy.spatial.transform import Rotation + + def to_canonical(obs, task_text, latent=None): + imgs = {} + for slot, key in (("cam0", "image"), ("cam1", "image2")): + x = torch.as_tensor(np.asarray(obs["pixels"][key])) + if x.dim() == 4: # (1, H, W, C) + x = x[0] + # robosuite renders 180-degree rotated relative to the recorded dataset + x = x.flip(0).flip(1) + x = x.permute(2, 0, 1).float() / 255.0 + if x.shape[-2:] != (cfg.image_size, cfg.image_size): + x = torch.nn.functional.interpolate( + x[None], size=(cfg.image_size, cfg.image_size), mode="bilinear", align_corners=False + )[0] + imgs[slot] = x + + # dataset convention (OpenVLA-style): [eef pos(3), eef axis-angle(3), gripper qpos(2)] + rs = obs["robot_state"] + pos = np.asarray(rs["eef"]["pos"]).flatten() + quat = np.asarray(rs["eef"]["quat"]).flatten() # robosuite: (x, y, z, w) + rotvec = Rotation.from_quat(quat).as_rotvec() + # canonicalize antipodal representation to match dataset convention + # (dataset uses rotvec with positive x-component, ~+pi for downward gripper) + if rotvec[0] < 0: + theta = np.linalg.norm(rotvec) + if theta > 1e-6: + rotvec = rotvec * (theta - 2 * np.pi) / theta + grip = np.asarray(rs["gripper"]["qpos"]).flatten() + state = torch.tensor(np.concatenate([pos, rotvec, grip]), dtype=torch.float32) + state = norm("observation.state", state) + state = torch.nn.functional.pad(state, (0, cfg.max_state_dim - state.shape[-1])) + + t = tok([task_text], padding=True, truncation=True, + max_length=cfg.tokenizer_max_length, return_tensors="pt") + batch = { + "observation.images.cam0": imgs["cam0"][None].cuda(), + "observation.images.cam1": imgs["cam1"][None].cuda(), + "observation.state": state[None].cuda(), + "observation.language.tokens": t["input_ids"].cuda(), + "observation.language.attention_mask": t["attention_mask"].bool().cuda(), + "embodiment_id": torch.tensor([args.embodiment_id], device="cuda"), + } + if latent is not None: + batch["semantic_latent"] = latent + return batch + + action_dim = meta.features["action"]["shape"][0] + total, succ = 0, 0 + per_task = {} + + task_ids = sorted(task_envs.keys()) + ep_plan = [(tid, i) for i in range((args.episodes + len(task_ids) - 1) // len(task_ids)) for tid in task_ids] + ep_plan = ep_plan[: args.episodes] + + if args.save_video_dir: + args.save_video_dir.mkdir(parents=True, exist_ok=True) + + for ep, (tid, rep) in enumerate(ep_plan): + env = task_envs[tid] + obs, info = env.reset(seed=1000 + rep) + frames = [] if args.save_video_dir else None + try: + task_text = env.get_attr("task_description")[0] + except Exception: + task_text = getattr(getattr(env, "envs", [None])[0], "task_description", "") + policy.reset() + latent = None + last_refresh = -1e9 + done = False + step_i = 0 + ep_succ = False + while not done and step_i < args.max_steps: + t_now = step_i / fps + batch = to_canonical(obs, task_text) + if args.refresh_s is not None: + if t_now - last_refresh >= args.refresh_s: + latent = policy._semantic_latent(batch) + last_refresh = t_now + batch["semantic_latent"] = latent + act_norm = policy.select_action(batch) # (1, max_action_dim) normalized padded + act = unnorm_action(act_norm[0, :action_dim].cpu()).clamp(-1, 1) + for _ in range(args.action_repeat): + try: + obs, reward, terminated, truncated, info = env.step(act.numpy()[None]) + except ValueError: # stepped into env's internal horizon + done = True + break + done = bool(terminated[0] or truncated[0]) + if info.get("is_success") is not None: + ep_succ = ep_succ or bool(np.asarray(info["is_success"]).flatten()[0]) + if frames is not None: + frames.append(np.asarray(obs["pixels"]["image"])[0][::-1, ::-1]) + step_i += 1 + if done or step_i >= args.max_steps: + done = done or step_i >= args.max_steps + break + total += 1 + succ += int(ep_succ) + per_task.setdefault(task_text[:50], []).append(int(ep_succ)) + print(f"ep {ep}: {'SUCCESS' if ep_succ else 'fail'} ({step_i} steps) | {task_text[:60]}") + if frames: + import imageio.v2 as imageio + + tag = "succ" if ep_succ else "fail" + path = args.save_video_dir / f"ep{ep:02d}_task{tid}_{tag}.mp4" + imageio.mimwrite(path, frames, fps=20, quality=7) + + print(f"\n=== {args.suite} | {args.checkpoint} | refresh={args.refresh_s} ===") + print(f"success rate: {succ}/{total} = {succ/total:.1%}") + for t, v in per_task.items(): + print(f" {sum(v)}/{len(v)} {t}") + + +if __name__ == "__main__": + main() diff --git a/tinyvla2/scripts/eval_mega_lekiwi.py b/tinyvla2/scripts/eval_mega_lekiwi.py new file mode 100644 index 0000000000000000000000000000000000000000..1fe8289cdf30d64639d7c0cd59c51717997dbc2e --- /dev/null +++ b/tinyvla2/scripts/eval_mega_lekiwi.py @@ -0,0 +1,106 @@ +#!/usr/bin/env python +"""Auto-eval for C-mega on held-out LeKiwi: text-prompt x numeric-descriptor ablation. + +Compares (numeric descriptor correct/wrong) x (text prompt correct/none) to see +whether the natural-language robot description adds anything beyond the C-scheme +numeric descriptor, zero-shot, no fine-tuning. Also reports the zero-floor and the +prior C-diverse baseline (295mm, numeric-only) for context. +""" + +from __future__ import annotations + +import numpy as np +import torch +import yaml +from lerobot.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata +from scipy.spatial.transform import Rotation +from transformers import AutoTokenizer + +from tinyvla.data.canonical import CanonicalChunkStore, quantile_normalize +from tinyvla.modeling_tinyvla import TinyVLAPolicy +from tinyvla.modules.embodiment import MORPH_FIELDS + +CKPT = "outputs/tv2_C_mega/final" +DS = "lekiwi_cleanup" +ROOT = f"/home/alexw/tinyvla_data/lekiwi/{DS}" +_SC = {"arm_dof": 0.1, "reach_m": 2, "gripper_width_m": 10, "num_cameras": 1 / 3, + "control_hz": 1 / 30, "joint_lo_mean": 1 / 3.1416, "joint_hi_mean": 1 / 3.1416, + "workspace_x": 2, "workspace_y": 2, "workspace_z": 2, "payload_kg": 0.2} + + +def mvec(d): + return torch.tensor([d.get(f, 0) * _SC.get(f, 1) for f in MORPH_FIELDS], dtype=torch.float32) + + +@torch.no_grad() +def main(): + pol = TinyVLAPolicy.from_pretrained(CKPT).cuda().eval() + cfg = pol.config + tok = AutoTokenizer.from_pretrained(cfg.lm_model_name) + desc = yaml.safe_load(open("configs/morphology/descriptors.yaml")) + prompts = yaml.safe_load(open("configs/morphology/robot_prompts.yaml")) + + m = LeRobotDatasetMetadata(DS, root=ROOT) + ds = LeRobotDataset(DS, root=ROOT, delta_timestamps={"action": [t / m.fps for t in range(50)]}, + video_backend="torchcodec") + store = CanonicalChunkStore(DS, src_fps=m.fps, chunk=50) + st = store.compute_stats() + q01, q99 = np.asarray(st["q01"]), np.asarray(st["q99"]) + span = np.maximum(q99 - q01, 0.01 * np.median(np.abs(np.concatenate([q01, q99])) + 1e-6)) + mid = 0.5 * (q01 + q99) + + imk = sorted(k for k in ds.meta.features if k.startswith("observation.images")) + prim = next((k for k in imk if "front" in k or "base" in k), imk[0]) + + def integ(d): + return np.cumsum(d[:, :3], 0) + + n_eps = ds.num_episodes + test = range(max(0, n_eps - 15), n_eps) + + def run(morph, prompt_prefix): + errs, zf = [], [] + for ep in test: + s = int(m.episodes["dataset_from_index"][ep]) + e = int(m.episodes["dataset_to_index"][ep]) + for idx in range(s, e - 1, 30): + item = ds[idx] + task = item.get("task") or "" + text = f"{prompt_prefix} {task}" if prompt_prefix else task + t = tok([text], padding=True, truncation=True, max_length=48, return_tensors="pt") + img = torch.nn.functional.interpolate(item[prim][None], size=(256, 256), mode="bilinear")[0] + stt = torch.nn.functional.pad(item["observation.state"].float()[:6], (0, cfg.max_state_dim - 6)) + b = {"observation.images.cam0": img[None].cuda(), + "observation.images.cam1": torch.zeros_like(img)[None].cuda(), + "observation.state": stt[None].cuda(), + "observation.language.tokens": t["input_ids"].cuda(), + "observation.language.attention_mask": t["attention_mask"].bool().cuda(), + "morphology": morph[None].cuda(), "embodiment_id": torch.tensor([0]).cuda()} + with torch.autocast("cuda", torch.bfloat16): + pr = pol.predict_action_chunk(b)[0].cpu().float().numpy() + gu = quantile_normalize(store.chunk_for(ep, idx - s), q01, q99)[:, :7] * span / 2 + mid + pu = pr[:, :7] * span / 2 + mid + errs.append(np.linalg.norm(integ(pu)[-1] - integ(gu)[-1]) * 1000) + zf.append(np.linalg.norm(integ(gu)[-1]) * 1000) + return np.mean(errs), np.mean(zf) + + print(f"=== C-mega on held-out LeKiwi (n_eps_test={len(list(test))}) ===") + zero_floor = None + conditions = [ + ("numeric=lekiwi + text=lekiwi", mvec(desc["lekiwi"]), prompts["lekiwi"]), + ("numeric=lekiwi + text=none ", mvec(desc["lekiwi"]), None), + ("numeric=none + text=lekiwi", torch.zeros(16), prompts["lekiwi"]), + ("numeric=none + text=none ", torch.zeros(16), None), + ("numeric=so101(wrong,non-mobile) + text=lekiwi", mvec(desc["so101"]), prompts["lekiwi"]), + ] + for label, morph, prompt in conditions: + err, zf = run(morph, prompt) + if zero_floor is None: + zero_floor = zf + print(f"{label:48} endpoint {err:.1f}mm") + print(f"\nzero-floor: {zero_floor:.1f}mm") + print("prior C-diverse baseline (numeric-only, no mega training): 295.0mm") + + +if __name__ == "__main__": + main() diff --git a/tinyvla2/scripts/eval_offline.py b/tinyvla2/scripts/eval_offline.py new file mode 100644 index 0000000000000000000000000000000000000000..ead9de1fc22e329b814c9e466397260f8aedc1a3 --- /dev/null +++ b/tinyvla2/scripts/eval_offline.py @@ -0,0 +1,114 @@ +#!/usr/bin/env python +"""Offline eval for canonical-schema checkpoints (stage2+). + +Reports action-chunk MSE (normalized space) on held-out episodes, the +per-timestep error curve, and stale-latent degradation. + +Usage: + python scripts/eval_offline.py \ + --checkpoint outputs/stage2_mixture/final \ + --repo-id VoicAndrei__so100_kitchen \ + --root ~/tinyvla_data/so101_v3/VoicAndrei__so100_kitchen \ + --episodes 8 --stale-s 0 1 2 [--embodiment-id 0] [--no-latent] +""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +import torch + + +@torch.no_grad() +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--checkpoint", type=Path, required=True) + parser.add_argument("--repo-id", required=True) + parser.add_argument("--root", default=None) + parser.add_argument("--episodes", type=int, default=8) + parser.add_argument("--stride", type=int, default=30) + parser.add_argument("--stale-s", type=float, nargs="*", default=[0.0, 1.0, 2.0]) + parser.add_argument("--embodiment-id", type=int, default=0) + parser.add_argument("--split", choices=["first", "last"], default="last") + args = parser.parse_args() + + from lerobot.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata + from transformers import AutoTokenizer + from tinyvla.data.mixture import CanonicalSource + from tinyvla.modeling_tinyvla import TinyVLAPolicy + + policy = TinyVLAPolicy.from_pretrained(args.checkpoint).cuda().eval() + cfg = policy.config + chunk = cfg.chunk_size + tok = AutoTokenizer.from_pretrained(cfg.lm_model_name) + + meta = LeRobotDatasetMetadata(args.repo_id, root=args.root) + ds = LeRobotDataset( + args.repo_id, + root=args.root, + delta_timestamps={"action": [t / meta.fps for t in range(chunk)]}, + video_backend="torchcodec", + ) + src = CanonicalSource( + ds, args.embodiment_id, cfg.image_size, cfg.max_state_dim, cfg.max_action_dim + ) + + if args.split == "first": + eps = list(range(args.episodes)) + else: + eps = list(range(ds.num_episodes - args.episodes, ds.num_episodes)) + + def to_batch(item): + t = tok([item.pop("task")], padding=True, truncation=True, + max_length=cfg.tokenizer_max_length, return_tensors="pt") + b = {k: v[None].cuda() if torch.is_tensor(v) else v for k, v in item.items()} + b["observation.language.tokens"] = t["input_ids"].cuda() + b["observation.language.attention_mask"] = t["attention_mask"].bool().cuda() + return b + + results = {s: [] for s in args.stale_s} + per_t = torch.zeros(chunk) + n_chunks = 0 + + for ep in eps: + start = int(ds.meta.episodes["dataset_from_index"][ep]) + end = int(ds.meta.episodes["dataset_to_index"][ep]) + for idx in range(start, end - 1, args.stride): + if idx >= len(src): + break + item = src[idx - 0] + gt = item["action"].clone() # (chunk, A) normalized + mask = item["action_dim_mask"].clone() + pad = item.get("action_is_pad") + batch = to_batch(dict(item)) + for stale_s in args.stale_s: + b = dict(batch) + if stale_s > 0: + stale_idx = max(start, idx - int(stale_s * ds.fps)) + stale_item = src[stale_idx] + sb = to_batch(dict(stale_item)) + b["semantic_latent"] = policy._semantic_latent(sb) + pred = policy.predict_action_chunk(b)[0].cpu() # (chunk, A) normalized + err = (pred[:, mask] - gt[:, mask]) ** 2 + if pad is not None: + err = err[~pad] + mse = err.mean().item() + results[stale_s].append(mse) + if stale_s == 0: + e = ((pred - gt) ** 2)[:, mask].mean(dim=-1) + if pad is not None: + e = e * (~pad).float() + per_t += e + n_chunks += 1 + + print(f"\n=== {args.repo_id} | {len(eps)} {args.split} episodes | {n_chunks} chunks | ckpt {args.checkpoint} ===") + for s, vals in results.items(): + print(f"stale {s:.0f}s: normalized chunk MSE {sum(vals)/len(vals):.4f}") + curve = (per_t / max(n_chunks, 1)).sqrt() + print("per-timestep normalized RMSE (t=0,10,25,49):", + [round(curve[i].item(), 3) for i in (0, 10, 25, 49)]) + + +if __name__ == "__main__": + main() diff --git a/tinyvla2/scripts/eval_qwenmorph_lekiwi.py b/tinyvla2/scripts/eval_qwenmorph_lekiwi.py new file mode 100644 index 0000000000000000000000000000000000000000..9208f47dc73dad81ac0c79e209fc140728d107e0 --- /dev/null +++ b/tinyvla2/scripts/eval_qwenmorph_lekiwi.py @@ -0,0 +1,102 @@ +#!/usr/bin/env python +"""Auto-eval for C-qwen-morph (morphology via shared-Qwen text encode, LM unfrozen) +on held-out LeKiwi. Compares correct vs wrong vs no robot-description text, and +against prior baselines (C-diverse numeric-MLP 295mm; C-mega text-in-prompt 522mm). +""" + +from __future__ import annotations + +import numpy as np +import torch +import yaml +from lerobot.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata +from transformers import AutoTokenizer + +from tinyvla.data.canonical import CanonicalChunkStore, quantile_normalize +from tinyvla.modeling_tinyvla import TinyVLAPolicy + +CKPT = "outputs/tv2_C_qwen_morph/final" +DS = "lekiwi_cleanup" +ROOT = f"/home/alexw/tinyvla_data/lekiwi/{DS}" + + +@torch.no_grad() +def main(): + pol = TinyVLAPolicy.from_pretrained(CKPT).cuda().eval() + cfg = pol.config + tok = AutoTokenizer.from_pretrained(cfg.lm_model_name) + prompts = yaml.safe_load(open("configs/morphology/robot_prompts.yaml")) + + m = LeRobotDatasetMetadata(DS, root=ROOT) + ds = LeRobotDataset(DS, root=ROOT, delta_timestamps={"action": [t / m.fps for t in range(50)]}, + video_backend="torchcodec") + store = CanonicalChunkStore(DS, src_fps=m.fps, chunk=50) + st = store.compute_stats() + q01, q99 = np.asarray(st["q01"]), np.asarray(st["q99"]) + span = np.maximum(q99 - q01, 0.01 * np.median(np.abs(np.concatenate([q01, q99])) + 1e-6)) + mid = 0.5 * (q01 + q99) + + imk = sorted(k for k in ds.meta.features if k.startswith("observation.images")) + prim = next((k for k in imk if "front" in k or "base" in k), imk[0]) + + def integ(d): + return np.cumsum(d[:, :3], 0) + + n_eps = ds.num_episodes + test = range(max(0, n_eps - 15), n_eps) + + def tok_desc(text): + if text is None: + ids = torch.zeros(1, cfg.morph_text_max_len, dtype=torch.long) + mask = torch.zeros(1, cfg.morph_text_max_len, dtype=torch.bool) + mask[0, 0] = True + return ids, mask + t = tok([text], padding="max_length", truncation=True, + max_length=cfg.morph_text_max_len, return_tensors="pt") + return t["input_ids"], t["attention_mask"].bool() + + def run(desc_text): + ids, mask = tok_desc(desc_text) + errs, zf = [], [] + for ep in test: + s = int(m.episodes["dataset_from_index"][ep]) + e = int(m.episodes["dataset_to_index"][ep]) + for idx in range(s, e - 1, 30): + item = ds[idx] + task = item.get("task") or "" + t = tok([task], padding=True, truncation=True, max_length=48, return_tensors="pt") + img = torch.nn.functional.interpolate(item[prim][None], size=(256, 256), mode="bilinear")[0] + stt = torch.nn.functional.pad(item["observation.state"].float()[:6], (0, cfg.max_state_dim - 6)) + b = {"observation.images.cam0": img[None].cuda(), + "observation.images.cam1": torch.zeros_like(img)[None].cuda(), + "observation.state": stt[None].cuda(), + "observation.language.tokens": t["input_ids"].cuda(), + "observation.language.attention_mask": t["attention_mask"].bool().cuda(), + "morph_text_ids": ids.cuda(), "morph_text_mask": mask.cuda(), + "embodiment_id": torch.tensor([0]).cuda()} + with torch.autocast("cuda", torch.bfloat16): + pr = pol.predict_action_chunk(b)[0].cpu().float().numpy() + gu = quantile_normalize(store.chunk_for(ep, idx - s), q01, q99)[:, :7] * span / 2 + mid + pu = pr[:, :7] * span / 2 + mid + errs.append(np.linalg.norm(integ(pu)[-1] - integ(gu)[-1]) * 1000) + zf.append(np.linalg.norm(integ(gu)[-1]) * 1000) + return np.mean(errs), np.mean(zf) + + print(f"=== C-qwen-morph (LM unfrozen, morph via shared Qwen) on held-out LeKiwi ===") + zero_floor = None + conditions = [ + ("desc=lekiwi (correct)", prompts["lekiwi"]), + ("desc=so101 (wrong, non-mobile arm)", prompts["so101"]), + ("desc=none", None), + ] + for label, text in conditions: + err, zf = run(text) + if zero_floor is None: + zero_floor = zf + print(f"{label:38} endpoint {err:.1f}mm") + print(f"\nzero-floor: {zero_floor:.1f}mm") + print("prior baselines: C-diverse numeric-MLP 295.0mm | C-mega text-in-prompt(frozen LM) 298.5mm best") + + +if __name__ == "__main__": + main() diff --git a/tinyvla2/scripts/eval_richslow_lekiwi.py b/tinyvla2/scripts/eval_richslow_lekiwi.py new file mode 100644 index 0000000000000000000000000000000000000000..cfce782c623c1171201c15684cc0429f202edbd8 --- /dev/null +++ b/tinyvla2/scripts/eval_richslow_lekiwi.py @@ -0,0 +1,141 @@ +#!/usr/bin/env python +"""Auto-eval for in-context demo conditioning on held-out LeKiwi. + +The key test: give the model K=3 REAL (obs, action) example pairs from LeKiwi's +OWN data at test time — NO gradient update, NO fine-tuning — and see if that beats +zero-shot with no examples, wrong examples, or the numeric/text descriptor alone. +""" + +from __future__ import annotations + +import numpy as np +import torch +import yaml +from lerobot.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata +from transformers import AutoTokenizer + +from tinyvla.data.canonical import CanonicalChunkStore, quantile_normalize +from tinyvla.modeling_tinyvla import TinyVLAPolicy + +CKPT = "outputs/tv2_richslow/final" +DS = "lekiwi_cleanup" +ROOT = f"/home/alexw/tinyvla_data/lekiwi/{DS}" +K = 3 + + +@torch.no_grad() +def main(): + pol = TinyVLAPolicy.from_pretrained(CKPT).cuda().eval() + cfg = pol.config + tok = AutoTokenizer.from_pretrained(cfg.lm_model_name) + prompts = yaml.safe_load(open("configs/morphology/robot_prompts.yaml")) + + m = LeRobotDatasetMetadata(DS, root=ROOT) + ds = LeRobotDataset(DS, root=ROOT, delta_timestamps={"action": [t / m.fps for t in range(50)]}, + video_backend="torchcodec") + store = CanonicalChunkStore(DS, src_fps=m.fps, chunk=50) + st = store.compute_stats() + q01, q99 = np.asarray(st["q01"]), np.asarray(st["q99"]) + span = np.maximum(q99 - q01, 0.01 * np.median(np.abs(np.concatenate([q01, q99])) + 1e-6)) + mid = 0.5 * (q01 + q99) + + imk = sorted(k for k in ds.meta.features if k.startswith("observation.images")) + prim = next((k for k in imk if "front" in k or "base" in k), imk[0]) + + def integ(d): + return np.cumsum(d[:, :3], 0) + + def img_at(idx, ep, ep_start): + it = ds[idx] + im = torch.nn.functional.interpolate(it[prim][None], size=(256, 256), mode="bilinear")[0] + raw = quantile_normalize(store.chunk_for(ep, idx - ep_start), q01, q99) + act = torch.from_numpy(raw[:, :7]).float() + act = torch.nn.functional.pad(act, (0, cfg.max_action_dim - 7)) + return im, act + + # REAL LeKiwi demo pairs from early episodes (support pool), disjoint from test + support_pool_eps = range(0, min(20, m.total_episodes - 15)) + test = range(max(0, m.total_episodes - 15), m.total_episodes) + + def get_real_support(k): + imgs, acts = [], [] + for _ in range(k): + ep = int(np.random.choice(list(support_pool_eps))) + s = int(m.episodes["dataset_from_index"][ep]) + e = int(m.episodes["dataset_to_index"][ep]) + idx = int(np.random.randint(s, max(s + 1, e - 1))) + im, act = img_at(idx, ep, s) + imgs.append(im) + acts.append(act) + return torch.stack(imgs), torch.stack(acts) + + def get_wrong_support(k): + # random noise images + random actions — a garbage support set control + return torch.rand(k, 3, 256, 256), torch.randn(k, 50, cfg.max_action_dim) * 0.3 + + + # control: demos from a DIFFERENT ROBOT (jaco) — if the model truly extracts + # robot-specific dynamics from exemplars, these should be worse than LeKiwi's own + def get_otherrobot_support(k): + from lerobot.datasets.lerobot_dataset import LeRobotDataset as _LD, LeRobotDatasetMetadata as _LM + from tinyvla.data.canonical import CanonicalChunkStore as _CS + jr = "/home/alexw/tinyvla_data/heldout/heldout_jaco_play" + jm = _LM("heldout_jaco_play", root=jr) + jds = _LD("heldout_jaco_play", root=jr, delta_timestamps={"action": [t / jm.fps for t in range(50)]}, video_backend="torchcodec") + jstore = _CS("heldout_jaco_play", src_fps=jm.fps, chunk=50) + jst = jstore.compute_stats(200) + jq01, jq99 = np.asarray(jst["q01"]), np.asarray(jst["q99"]) + jimk = sorted(kk for kk in jds.meta.features if kk.startswith("observation.images"))[0] + imgs, acts = [], [] + for _ in range(k): + ep = int(np.random.randint(0, 50)) + s = int(jm.episodes["dataset_from_index"][ep]); e = int(jm.episodes["dataset_to_index"][ep]) + idx = int(np.random.randint(s, max(s + 1, e - 1))) + it = jds[idx] + im = torch.nn.functional.interpolate(it[jimk][None], size=(256, 256), mode="bilinear")[0] + raw = quantile_normalize(jstore.chunk_for(ep, idx - s), jq01, jq99) + act = torch.nn.functional.pad(torch.from_numpy(raw[:, :7]).float(), (0, cfg.max_action_dim - 7)) + imgs.append(im); acts.append(act) + return torch.stack(imgs), torch.stack(acts) + + def run(support_fn, desc_text): + ids_t = tok([desc_text] if desc_text else [""], padding="max_length", truncation=True, + max_length=cfg.morph_text_max_len, return_tensors="pt") + errs = [] + for ep in test: + s = int(m.episodes["dataset_from_index"][ep]) + e = int(m.episodes["dataset_to_index"][ep]) + for idx in range(s, e - 1, 30): + item = ds[idx] + task = item.get("task") or "" + t = tok([task], padding=True, truncation=True, max_length=48, return_tensors="pt") + img = torch.nn.functional.interpolate(item[prim][None], size=(256, 256), mode="bilinear")[0] + stt = torch.nn.functional.pad(item["observation.state"].float()[:6], (0, cfg.max_state_dim - 6)) + sup_img, sup_act = support_fn(K) + b = {"observation.images.cam0": img[None].cuda(), + "observation.images.cam1": torch.zeros_like(img)[None].cuda(), + "observation.state": stt[None].cuda(), + "observation.language.tokens": t["input_ids"].cuda(), + "observation.language.attention_mask": t["attention_mask"].bool().cuda(), + "morph_text_ids": ids_t["input_ids"].cuda(), + "morph_text_mask": ids_t["attention_mask"].bool().cuda(), + "support_images": sup_img[None].cuda(), "support_actions": sup_act[None].cuda(), + "embodiment_id": torch.tensor([0]).cuda()} + with torch.autocast("cuda", torch.bfloat16): + pr = pol.predict_action_chunk(b)[0].cpu().float().numpy() + gu = quantile_normalize(store.chunk_for(ep, idx - s), q01, q99)[:, :7] * span / 2 + mid + pu = pr[:, :7] * span / 2 + mid + errs.append(np.linalg.norm(integ(pu)[-1] - integ(gu)[-1]) * 1000) + return np.mean(errs) + + print("=== RICH-SLOW (demos in LM attention stream, no FT) on held-out LeKiwi ===") + print(f"REAL LeKiwi demos (K={K}) + desc=lekiwi endpoint {run(get_real_support, prompts['lekiwi']):.1f}mm") + print(f"REAL LeKiwi demos (K={K}) + desc=none endpoint {run(get_real_support, None):.1f}mm") + print(f"WRONG/garbage demos + desc=lekiwi endpoint {run(get_wrong_support, prompts['lekiwi']):.1f}mm") + print(f"WRONG/garbage demos + desc=none endpoint {run(get_wrong_support, None):.1f}mm") + print(f"OTHER-ROBOT (jaco) demos + desc=none endpoint {run(get_otherrobot_support, None):.1f}mm") + print("\nprior baselines: C-diverse 295mm | C-qwen-morph desc=none 294mm | v1 demo(same-task) real=290.6 garbage=288.5 (dead)") + + +if __name__ == "__main__": + main() diff --git a/tinyvla2/scripts/eval_robocasa_all.py b/tinyvla2/scripts/eval_robocasa_all.py new file mode 100644 index 0000000000000000000000000000000000000000..108564f6c4cd282321ac3d9f0c00e6878da3a83d --- /dev/null +++ b/tinyvla2/scripts/eval_robocasa_all.py @@ -0,0 +1,110 @@ +#!/usr/bin/env python +"""How does the best checkpoint do on SIMILAR robots (arms like the training ones) +vs the exotic held-out ones? Reports endpoint error AND the zero-prediction floor +for each, since the floor differs a lot per robot (it is what "predicting mean +motion" achieves) — only the ratio to floor is comparable across robots. +""" + +from __future__ import annotations + +import numpy as np +import torch +import yaml +from lerobot.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata +from transformers import AutoTokenizer + +from tinyvla.data.canonical import CanonicalChunkStore, quantile_normalize +from tinyvla.modeling_tinyvla import TinyVLAPolicy +from tinyvla.modules.embodiment import MORPH_FIELDS + +CKPT = "outputs/tv2_robocasa/final" +_SC = {"arm_dof": 0.1, "reach_m": 2, "gripper_width_m": 10, "num_cameras": 1 / 3, + "control_hz": 1 / 30, "joint_lo_mean": 1 / 3.1416, "joint_hi_mean": 1 / 3.1416, + "workspace_x": 2, "workspace_y": 2, "workspace_z": 2, "payload_kg": 0.2} + +# (label, dataset name, root, morph_key, in_training?) +TARGETS = [ + ("jaco (held-out ARM)", "heldout_jaco_play", "~/tinyvla_data/heldout/heldout_jaco_play", "jaco", False), + ("ur5 (in training)", "div_ur5", "~/tinyvla_data/diverse/div_ur5", "ur5", True), + ("xarm (in training)", "div_xarm", "~/tinyvla_data/diverse/div_xarm", "xarm", True), + ("dlr_edan (in training)", "div_dlr_edan", "~/tinyvla_data/diverse/div_dlr_edan", "dlr_edan", True), + ("stretch (in training, mobile)", "heldout_cmu_stretch", "~/tinyvla_data/heldout/heldout_cmu_stretch", "hello_stretch", True), + ("LeKiwi (held-out MOBILE)", "lekiwi_cleanup", "~/tinyvla_data/lekiwi/lekiwi_cleanup", "lekiwi", False), + ("RoboCasa (NEW domain)", "robocasa365", "~/tinyvla_data/robocasa365", "panda_omron", True), +] + + +@torch.no_grad() +def main(): + from pathlib import Path + + pol = TinyVLAPolicy.from_pretrained(CKPT).cuda().eval() + cfg = pol.config + tok = AutoTokenizer.from_pretrained(cfg.lm_model_name) + desc = yaml.safe_load(open("configs/morphology/descriptors.yaml")) + + print(f"=== {CKPT} across robots (endpoint err vs zero-floor) ===") + rows = [] + for label, name, root, mkey, in_train in TARGETS: + root = str(Path(root).expanduser()) + try: + m = LeRobotDatasetMetadata(name, root=root) + ds = LeRobotDataset(name, root=root, + delta_timestamps={"action": [t / m.fps for t in range(50)]}, + video_backend="torchcodec") + store = CanonicalChunkStore(name, src_fps=m.fps, chunk=50) + st = store.compute_stats() + q01, q99 = np.asarray(st["q01"]), np.asarray(st["q99"]) + span = np.maximum(q99 - q01, 0.01 * np.median(np.abs(np.concatenate([q01, q99])) + 1e-6)) + mid = 0.5 * (q01 + q99) + imk = sorted(k for k in ds.meta.features if k.startswith("observation.images")) + prim = next((k for k in imk if any(s in k for s in ("front", "base", "top", "image"))), imk[0]) + morph = torch.tensor([desc[mkey].get(f, 0) * _SC.get(f, 1) for f in MORPH_FIELDS], + dtype=torch.float32) + n_eps = ds.num_episodes + test = range(max(0, n_eps - 12), n_eps) + errs, floors = [], [] + for ep in test: + s = int(m.episodes["dataset_from_index"][ep]) + e = int(m.episodes["dataset_to_index"][ep]) + for idx in range(s, e - 1, 40): + item = ds[idx] + t = tok([item.get("task") or ""], padding=True, truncation=True, + max_length=48, return_tensors="pt") + img = torch.nn.functional.interpolate(item[prim][None], size=(256, 256), mode="bilinear")[0] + sd = item["observation.state"].float() + stt = torch.nn.functional.pad(sd[: min(len(sd), cfg.max_state_dim)], + (0, max(0, cfg.max_state_dim - len(sd)))) + b = {"observation.images.cam0": img[None].cuda(), + "observation.images.cam1": torch.zeros_like(img)[None].cuda(), + "observation.state": stt[None].cuda(), + "observation.language.tokens": t["input_ids"].cuda(), + "observation.language.attention_mask": t["attention_mask"].bool().cuda(), + "morphology": morph[None].cuda(), + "embodiment_id": torch.tensor([0]).cuda()} + with torch.autocast("cuda", torch.bfloat16): + pr = pol.predict_action_chunk(b)[0].cpu().float().numpy() + gu = quantile_normalize(store.chunk_for(ep, idx - s), q01, q99)[:, :7] * span / 2 + mid + pu = pr[:, :7] * span / 2 + mid + gp = np.cumsum(gu[:, :3], 0) + pp = np.cumsum(pu[:, :3], 0) + errs.append(np.linalg.norm(pp[-1] - gp[-1]) * 1000) + floors.append(np.linalg.norm(gp[-1]) * 1000) + err, floor = float(np.mean(errs)), float(np.mean(floors)) + rows.append((label, in_train, err, floor, err / floor)) + print(f"{label:32} {'[train]' if in_train else '[HELD-OUT]':11} " + f"endpoint {err:6.1f}mm floor {floor:6.1f}mm ratio {err/floor:.2f}") + except Exception as ex: + print(f"{label:32} FAILED {type(ex).__name__}: {str(ex)[:70]}") + + print("\nratio < 1.0 = better than predicting mean motion") + tr = [r for r in rows if r[1]] + ho = [r for r in rows if not r[1]] + if tr: + print(f"mean ratio, robots IN training : {np.mean([r[4] for r in tr]):.2f}") + if ho: + print(f"mean ratio, HELD-OUT robots : {np.mean([r[4] for r in ho]):.2f}") + + +if __name__ == "__main__": + main() diff --git a/tinyvla2/scripts/eval_similar_robots.py b/tinyvla2/scripts/eval_similar_robots.py new file mode 100644 index 0000000000000000000000000000000000000000..1d924497066ecf1fb46a8b87478c5b1ad4023588 --- /dev/null +++ b/tinyvla2/scripts/eval_similar_robots.py @@ -0,0 +1,109 @@ +#!/usr/bin/env python +"""How does the best checkpoint do on SIMILAR robots (arms like the training ones) +vs the exotic held-out ones? Reports endpoint error AND the zero-prediction floor +for each, since the floor differs a lot per robot (it is what "predicting mean +motion" achieves) — only the ratio to floor is comparable across robots. +""" + +from __future__ import annotations + +import numpy as np +import torch +import yaml +from lerobot.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata +from transformers import AutoTokenizer + +from tinyvla.data.canonical import CanonicalChunkStore, quantile_normalize +from tinyvla.modeling_tinyvla import TinyVLAPolicy +from tinyvla.modules.embodiment import MORPH_FIELDS + +CKPT = "outputs/tv2_C_diverse/final" +_SC = {"arm_dof": 0.1, "reach_m": 2, "gripper_width_m": 10, "num_cameras": 1 / 3, + "control_hz": 1 / 30, "joint_lo_mean": 1 / 3.1416, "joint_hi_mean": 1 / 3.1416, + "workspace_x": 2, "workspace_y": 2, "workspace_z": 2, "payload_kg": 0.2} + +# (label, dataset name, root, morph_key, in_training?) +TARGETS = [ + ("jaco (held-out ARM)", "heldout_jaco_play", "~/tinyvla_data/heldout/heldout_jaco_play", "jaco", False), + ("ur5 (in training)", "div_ur5", "~/tinyvla_data/diverse/div_ur5", "ur5", True), + ("xarm (in training)", "div_xarm", "~/tinyvla_data/diverse/div_xarm", "xarm", True), + ("dlr_edan (in training)", "div_dlr_edan", "~/tinyvla_data/diverse/div_dlr_edan", "dlr_edan", True), + ("stretch (in training, mobile)", "heldout_cmu_stretch", "~/tinyvla_data/heldout/heldout_cmu_stretch", "hello_stretch", True), + ("LeKiwi (held-out MOBILE)", "lekiwi_cleanup", "~/tinyvla_data/lekiwi/lekiwi_cleanup", "lekiwi", False), +] + + +@torch.no_grad() +def main(): + from pathlib import Path + + pol = TinyVLAPolicy.from_pretrained(CKPT).cuda().eval() + cfg = pol.config + tok = AutoTokenizer.from_pretrained(cfg.lm_model_name) + desc = yaml.safe_load(open("configs/morphology/descriptors.yaml")) + + print(f"=== {CKPT} across robots (endpoint err vs zero-floor) ===") + rows = [] + for label, name, root, mkey, in_train in TARGETS: + root = str(Path(root).expanduser()) + try: + m = LeRobotDatasetMetadata(name, root=root) + ds = LeRobotDataset(name, root=root, + delta_timestamps={"action": [t / m.fps for t in range(50)]}, + video_backend="torchcodec") + store = CanonicalChunkStore(name, src_fps=m.fps, chunk=50) + st = store.compute_stats() + q01, q99 = np.asarray(st["q01"]), np.asarray(st["q99"]) + span = np.maximum(q99 - q01, 0.01 * np.median(np.abs(np.concatenate([q01, q99])) + 1e-6)) + mid = 0.5 * (q01 + q99) + imk = sorted(k for k in ds.meta.features if k.startswith("observation.images")) + prim = next((k for k in imk if any(s in k for s in ("front", "base", "top", "image"))), imk[0]) + morph = torch.tensor([desc[mkey].get(f, 0) * _SC.get(f, 1) for f in MORPH_FIELDS], + dtype=torch.float32) + n_eps = ds.num_episodes + test = range(max(0, n_eps - 12), n_eps) + errs, floors = [], [] + for ep in test: + s = int(m.episodes["dataset_from_index"][ep]) + e = int(m.episodes["dataset_to_index"][ep]) + for idx in range(s, e - 1, 40): + item = ds[idx] + t = tok([item.get("task") or ""], padding=True, truncation=True, + max_length=48, return_tensors="pt") + img = torch.nn.functional.interpolate(item[prim][None], size=(256, 256), mode="bilinear")[0] + sd = item["observation.state"].float() + stt = torch.nn.functional.pad(sd[: min(len(sd), cfg.max_state_dim)], + (0, max(0, cfg.max_state_dim - len(sd)))) + b = {"observation.images.cam0": img[None].cuda(), + "observation.images.cam1": torch.zeros_like(img)[None].cuda(), + "observation.state": stt[None].cuda(), + "observation.language.tokens": t["input_ids"].cuda(), + "observation.language.attention_mask": t["attention_mask"].bool().cuda(), + "morphology": morph[None].cuda(), + "embodiment_id": torch.tensor([0]).cuda()} + with torch.autocast("cuda", torch.bfloat16): + pr = pol.predict_action_chunk(b)[0].cpu().float().numpy() + gu = quantile_normalize(store.chunk_for(ep, idx - s), q01, q99)[:, :7] * span / 2 + mid + pu = pr[:, :7] * span / 2 + mid + gp = np.cumsum(gu[:, :3], 0) + pp = np.cumsum(pu[:, :3], 0) + errs.append(np.linalg.norm(pp[-1] - gp[-1]) * 1000) + floors.append(np.linalg.norm(gp[-1]) * 1000) + err, floor = float(np.mean(errs)), float(np.mean(floors)) + rows.append((label, in_train, err, floor, err / floor)) + print(f"{label:32} {'[train]' if in_train else '[HELD-OUT]':11} " + f"endpoint {err:6.1f}mm floor {floor:6.1f}mm ratio {err/floor:.2f}") + except Exception as ex: + print(f"{label:32} FAILED {type(ex).__name__}: {str(ex)[:70]}") + + print("\nratio < 1.0 = better than predicting mean motion") + tr = [r for r in rows if r[1]] + ho = [r for r in rows if not r[1]] + if tr: + print(f"mean ratio, robots IN training : {np.mean([r[4] for r in tr]):.2f}") + if ho: + print(f"mean ratio, HELD-OUT robots : {np.mean([r[4] for r in ho]):.2f}") + + +if __name__ == "__main__": + main() diff --git a/tinyvla2/scripts/fewshot_grid.py b/tinyvla2/scripts/fewshot_grid.py new file mode 100644 index 0000000000000000000000000000000000000000..019f827d320e0a8848718be09633a62f287da2e9 --- /dev/null +++ b/tinyvla2/scripts/fewshot_grid.py @@ -0,0 +1,113 @@ +#!/usr/bin/env python +"""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 # jaco has 976 eps; eval on last 20, FT on first N (<= 936) +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() diff --git a/tinyvla2/scripts/label_spatial.py b/tinyvla2/scripts/label_spatial.py new file mode 100644 index 0000000000000000000000000000000000000000..bb872ab90521dd19c032ad3a2651a6de06c41833 --- /dev/null +++ b/tinyvla2/scripts/label_spatial.py @@ -0,0 +1,148 @@ +#!/usr/bin/env python +"""Stage 3: spatial distillation labels from a Qwen3-VL teacher (XS-VLA style). + +For every Nth frame of each episode, ask the teacher for the bounding box of +the task-relevant object; quantize the box center onto a GRID x GRID map and +store the cell index. The student later learns a linear classifier over its +fast-path spatial tokens with CE x 0.15 on labeled frames. + +Output: parquet with columns (dataset, episode_index, frame_index, cell, +cx, cy, confidence_ok) at ~/tinyvla_data/spatial_labels/.parquet + +Usage: + python scripts/label_spatial.py --teacher Qwen/Qwen3-VL-4B-Instruct \ + --frame-stride 10 --episode-frac 0.4 [--datasets-limit 2] [--pilot 20] +""" + +from __future__ import annotations + +import argparse +import json +import re +from pathlib import Path + +import torch + +GRID = 32 +OUT_DIR = Path.home() / "tinyvla_data" / "spatial_labels" + +POINT_RE = re.compile(r"\[?\s*(\d+)\s*,\s*(\d+)\s*\]?") + +PROMPT = ( + "Task: {task}\n" + "Look at the image. Where is the single object the robot must interact with next " + "to accomplish this task? Answer with ONLY its center point as [x, y] in 0-1000 " + "normalized coordinates. If unsure, output [0,0]." +) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--teacher", default="Qwen/Qwen3.5-4B") + parser.add_argument("--data-root", type=Path, default=Path.home() / "tinyvla_data/so101_v3") + parser.add_argument("--frame-stride", type=int, default=10) + parser.add_argument("--episode-frac", type=float, default=0.4) + parser.add_argument("--datasets-limit", type=int, default=None) + parser.add_argument("--pilot", type=int, default=None, help="label only N frames total, print results") + parser.add_argument("--batch-size", type=int, default=16) + args = parser.parse_args() + + import pyarrow as pa + import pyarrow.parquet as pq + from lerobot.datasets.lerobot_dataset import LeRobotDataset + from transformers import AutoModelForImageTextToText, AutoProcessor + + model = AutoModelForImageTextToText.from_pretrained( + args.teacher, dtype=torch.bfloat16, device_map="cuda" + ) + proc = AutoProcessor.from_pretrained(args.teacher) + OUT_DIR.mkdir(parents=True, exist_ok=True) + + roots = sorted(args.data_root.iterdir()) + if args.datasets_limit: + roots = roots[: args.datasets_limit] + + total_done = 0 + for root in roots: + if not (root / "meta/info.json").exists(): + continue + out_path = OUT_DIR / f"{root.name}.parquet" + if out_path.exists() and not args.pilot: + continue + ds = LeRobotDataset(root.name, root=root, video_backend="torchcodec") + image_key = sorted(k for k in ds.meta.features if k.startswith("observation.images"))[0] + n_eps = max(1, int(ds.num_episodes * args.episode_frac)) + + rows = [] + pending = [] # (ep, fi, image_pil, task) + + def flush(): + nonlocal total_done + if not pending: + return + msgs = [ + [{"role": "user", "content": [ + {"type": "image", "image": img}, + {"type": "text", "text": PROMPT.format(task=task)}, + ]}] + for _, _, img, task in pending + ] + texts = [ + proc.apply_chat_template( + m, tokenize=False, add_generation_prompt=True, enable_thinking=False + ) + for m in msgs + ] + images = [[p[2]] for p in pending] + inputs = proc(text=texts, images=images, return_tensors="pt", padding=True).to("cuda") + with torch.no_grad(): + out = model.generate(**inputs, max_new_tokens=16, do_sample=False) + answers = proc.batch_decode(out[:, inputs["input_ids"].shape[1]:], skip_special_tokens=True) + for (ep, fi, _, _), ans in zip(pending, answers): + m = POINT_RE.search(ans) + ok = False + cell, cx, cy = -1, -1.0, -1.0 + if m: + px, py = (int(g) for g in m.groups()) + if 0 < px <= 1000 and 0 < py <= 1000: + cx, cy = px / 1000.0, py / 1000.0 + gx, gy = min(GRID - 1, int(cx * GRID)), min(GRID - 1, int(cy * GRID)) + cell = gy * GRID + gx + ok = True + rows.append({"dataset": root.name, "episode_index": ep, "frame_index": fi, + "cell": cell, "cx": cx, "cy": cy, "confidence_ok": ok}) + if args.pilot: + print(f"ep{ep} f{fi}: '{ans.strip()[:60]}' -> cell {cell} ({cx:.2f},{cy:.2f})") + total_done += len(pending) + pending.clear() + + from torchvision.transforms.functional import to_pil_image + import torch.nn.functional as F + + for ep in range(n_eps): + start = int(ds.meta.episodes["dataset_from_index"][ep]) + end = int(ds.meta.episodes["dataset_to_index"][ep]) + for idx in range(start, end, args.frame_stride): + item = ds[idx] + # label on the same 256^2 view the student sees; ~8x fewer + # teacher vision tokens than full res + small = F.interpolate( + item[image_key][None].clamp(0, 1), size=(256, 256), + mode="bilinear", align_corners=False, + )[0] + img = to_pil_image(small) + pending.append((ep, idx - start, img, item.get("task") or "")) + if len(pending) >= args.batch_size: + flush() + if args.pilot and total_done + len(pending) >= args.pilot: + flush() + print(f"pilot done: {total_done} frames") + return + flush() + pq.write_table(pa.Table.from_pylist(rows), out_path) + ok_rate = sum(r["confidence_ok"] for r in rows) / max(len(rows), 1) + print(f"{root.name}: {len(rows)} labels -> {out_path} (ok {ok_rate:.1%})") + + +if __name__ == "__main__": + main() diff --git a/tinyvla2/scripts/prepare_community_v3.py b/tinyvla2/scripts/prepare_community_v3.py new file mode 100644 index 0000000000000000000000000000000000000000..0df30de0cab0a318b8f9b22fcee142aa52f8ba02 --- /dev/null +++ b/tinyvla2/scripts/prepare_community_v3.py @@ -0,0 +1,163 @@ +#!/usr/bin/env python +"""Prepare the SO-100/101 slice of HuggingFaceVLA/community_dataset_v3. + +Steps: + catalog — fetch all sub-dataset info.json, write catalog.json + select — filter SO100/101 single-arm, rank by episodes, write selection.json + convert — download each selected subdir + convert v2.1 -> v3.0 locally + +Usage: + python scripts/prepare_community_v3.py catalog + python scripts/prepare_community_v3.py select --target-episodes 12000 + python scripts/prepare_community_v3.py convert [--limit N] +""" + +from __future__ import annotations + +import argparse +import json +import shutil +from concurrent.futures import ThreadPoolExecutor, as_completed +from pathlib import Path + +AGG_REPO = "HuggingFaceVLA/community_dataset_v3" +DISK_BUDGET_GB = 400 +DATA_ROOT = Path.home() / "tinyvla_data" +CATALOG = DATA_ROOT / "community_v3_catalog.json" +SELECTION = DATA_ROOT / "community_v3_selection.json" +CONVERTED_DIR = DATA_ROOT / "so101_v3" +SO_ROBOTS = {"so100", "so101", "so100_follower", "so101_follower"} + + +def cmd_catalog(args): + from huggingface_hub import HfApi, hf_hub_download + + api = HfApi() + files = api.list_repo_files(AGG_REPO, repo_type="dataset") + subs = sorted(f.rsplit("/meta/info.json", 1)[0] for f in files if f.endswith("meta/info.json")) + print(f"{len(subs)} sub-datasets") + + def fetch(sub): + try: + p = hf_hub_download(AGG_REPO, f"{sub}/meta/info.json", repo_type="dataset") + info = json.load(open(p)) + return sub, { + "robot_type": info.get("robot_type"), + "episodes": info.get("total_episodes"), + "frames": info.get("total_frames"), + "fps": info.get("fps"), + "version": info.get("codebase_version"), + "action_shape": info.get("features", {}).get("action", {}).get("shape"), + "cameras": [k for k in info.get("features", {}) if k.startswith("observation.images")], + } + except Exception as e: + return sub, {"error": str(e)[:100]} + + catalog = {} + with ThreadPoolExecutor(16) as ex: + futs = [ex.submit(fetch, s) for s in subs] + for i, f in enumerate(as_completed(futs)): + sub, meta = f.result() + catalog[sub] = meta + if (i + 1) % 100 == 0: + print(f"{i+1}/{len(subs)}") + + DATA_ROOT.mkdir(parents=True, exist_ok=True) + CATALOG.write_text(json.dumps(catalog, indent=1)) + print(f"wrote {CATALOG}") + + +def cmd_select(args): + catalog = json.loads(CATALOG.read_text()) + rows = [ + (sub, m) + for sub, m in catalog.items() + if m.get("robot_type") in SO_ROBOTS + and m.get("action_shape") == [6] + and m.get("fps") == 30 + and m.get("episodes") + and m.get("cameras") + ] + # prefer larger datasets: fewer conversions per episode + rows.sort(key=lambda r: -r[1]["episodes"]) + picked, total = [], 0 + for sub, m in rows: + if total >= args.target_episodes: + break + picked.append({"sub": sub, **m}) + total += m["episodes"] + print(f"{len(rows)} eligible; picked {len(picked)} datasets, {total} episodes") + SELECTION.write_text(json.dumps(picked, indent=1)) + print(f"wrote {SELECTION}") + + +def free_gb(path: Path) -> float: + return shutil.disk_usage(path).free / 1e9 + + +def used_gb_cached(path: Path) -> float: + import subprocess + + out = subprocess.run(["du", "-s", "--block-size=1G", str(path)], capture_output=True, text=True) + return float(out.stdout.split()[0]) if out.returncode == 0 else 0.0 + + +def cmd_convert(args): + from huggingface_hub import snapshot_download + from lerobot.scripts.convert_dataset_v21_to_v30 import convert_dataset + + picked = json.loads(SELECTION.read_text()) + if args.limit: + picked = picked[: args.limit] + CONVERTED_DIR.mkdir(parents=True, exist_ok=True) + raw_dir = DATA_ROOT / "_raw_v21" + done, failed = 0, [] + for item in picked: + sub = item["sub"] + name = sub.replace("/", "__") + out = CONVERTED_DIR / name + if (out / "meta" / "info.json").exists(): + done += 1 + continue + local = raw_dir / name + try: + snapshot_download( + AGG_REPO, + repo_type="dataset", + allow_patterns=[f"{sub}/*"], + local_dir=raw_dir / "_dl", + ) + src = raw_dir / "_dl" / sub + if local.exists(): + shutil.rmtree(local) + shutil.move(str(src), str(local)) + convert_dataset(repo_id=name, root=local, push_to_hub=False, force_conversion=True) + # converter writes v3.0 in place at root + shutil.move(str(local), str(out)) + # drop the v2.1 originals the converter stashes as _old + old = local.parent / (local.name + "_old") + if old.exists(): + shutil.rmtree(old) + done += 1 + print(f"[{done}/{len(picked)}] {sub}: converted -> {out} (free {free_gb(DATA_ROOT):.0f}GB)") + if free_gb(DATA_ROOT) < 150 or used_gb_cached(DATA_ROOT) > DISK_BUDGET_GB: + print("STOP: disk budget reached") + break + except Exception as e: + failed.append(sub) + print(f"FAIL {sub}: {type(e).__name__}: {str(e)[:200]}") + print(f"done={done} failed={len(failed)}") + if failed: + (DATA_ROOT / "convert_failures.json").write_text(json.dumps(failed)) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + sub = parser.add_subparsers(dest="cmd", required=True) + sub.add_parser("catalog") + p_sel = sub.add_parser("select") + p_sel.add_argument("--target-episodes", type=int, default=12_000) + p_conv = sub.add_parser("convert") + p_conv.add_argument("--limit", type=int, default=None) + args = parser.parse_args() + {"catalog": cmd_catalog, "select": cmd_select, "convert": cmd_convert}[args.cmd](args) diff --git a/tinyvla2/scripts/replay_libero.py b/tinyvla2/scripts/replay_libero.py new file mode 100644 index 0000000000000000000000000000000000000000..4deafbd6996ad4b6f7aabad7e0525609e10a7037 --- /dev/null +++ b/tinyvla2/scripts/replay_libero.py @@ -0,0 +1,85 @@ +#!/usr/bin/env python +"""Replay ground-truth dataset actions in the LIBERO env. + +Diagnoses action-convention correctness independently of the policy: if GT +actions from a matching task episode achieve success (at least sometimes), +the env adapter/action semantics are right and remaining failures are model +quality; if GT never succeeds, the action mapping is wrong. + +Usage: + python scripts/replay_libero.py --suite libero_spatial --episodes 5 --action-repeat 2 +""" + +from __future__ import annotations + +import argparse + +import numpy as np +import torch + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--suite", default="libero_spatial") + parser.add_argument("--episodes", type=int, default=5) + parser.add_argument("--action-repeat", type=int, default=2) + args = parser.parse_args() + + from lerobot.datasets.lerobot_dataset import LeRobotDataset + from lerobot.envs.factory import make_env, make_env_config + + ds = LeRobotDataset("HuggingFaceVLA/libero", video_backend="torchcodec") + env_cfg = make_env_config("libero", task=args.suite) + task_envs = make_env(env_cfg, n_envs=1)[args.suite] + + # map task_description -> env + env_by_task = {} + for tid, env in task_envs.items(): + try: + desc = env.get_attr("task_description")[0] + except Exception: + desc = getattr(env.envs[0], "task_description", "") + env_by_task[desc.strip().lower()] = (tid, env) + + # first dataset episode per task <-> env's first reset (init_state_id=0), + # assuming dataset episode order preserves LIBERO demo order within a task + eps_meta = ds.meta.episodes + first_ep_by_task = {} + for ep in range(ds.num_episodes): + start = int(eps_meta["dataset_from_index"][ep]) + task = ds[start]["task"].strip().lower() + if task in env_by_task and task not in first_ep_by_task: + first_ep_by_task[task] = ep + + replayed, succ = 0, 0 + for task, ep in list(first_ep_by_task.items())[: args.episodes]: + start = int(eps_meta["dataset_from_index"][ep]) + end = int(eps_meta["dataset_to_index"][ep]) + tid, env = env_by_task[task] + obs, info = env.reset(seed=0) # fresh env: init_state_id starts at 0 + ep_succ = False + for idx in range(start, end): + act = ds[idx]["action"].numpy() + done = False + for _ in range(args.action_repeat): + try: + obs, reward, term, trunc, info = env.step(act[None]) + except ValueError: + done = True + break + if info.get("is_success") is not None: + ep_succ = ep_succ or bool(np.asarray(info["is_success"]).flatten()[0]) + if term[0] or trunc[0]: + done = True + break + if done or ep_succ: + break + replayed += 1 + succ += int(ep_succ) + print(f"ep {ep} (task {tid}): {'SUCCESS' if ep_succ else 'fail'} | {task[:60]}") + + print(f"\nGT replay success: {succ}/{replayed}") + + +if __name__ == "__main__": + main() diff --git a/tinyvla2/scripts/tidybot_fewshot.py b/tinyvla2/scripts/tidybot_fewshot.py new file mode 100644 index 0000000000000000000000000000000000000000..6e46b37d17604825573f47b18e820d9a4bc2a63e --- /dev/null +++ b/tinyvla2/scripts/tidybot_fewshot.py @@ -0,0 +1,104 @@ +#!/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() diff --git a/tinyvla2/scripts/train.py b/tinyvla2/scripts/train.py new file mode 100644 index 0000000000000000000000000000000000000000..9f64a4594217ea5cdc0cae81f2a8d9e43337ca2a --- /dev/null +++ b/tinyvla2/scripts/train.py @@ -0,0 +1,306 @@ +#!/usr/bin/env python +"""Stage-2+ training loop: weighted multi-dataset mixture + accelerate. + +Thin replacement for lerobot-train adding: +- weighted sampling across LeRobotDatasets (per-dataset embodiment ids) +- staleness augmentation (Stage 2b) +- spatial-distillation aux loss (Stage 3) + +Usage: + accelerate launch scripts/train.py --config configs/stage2_mixture.yaml +""" + +from __future__ import annotations + +import argparse +import math +import time +from pathlib import Path + +import torch +import yaml + + +def make_policy(cfg: dict): + """Canonical-schema policy: two fixed camera slots, padded state/action.""" + from lerobot.configs import FeatureType, PolicyFeature + from tinyvla.configuration_tinyvla import TinyVLAConfig + from tinyvla.modeling_tinyvla import TinyVLAPolicy + + pcfg = TinyVLAConfig(**cfg.get("policy", {})) + s = pcfg.image_size + pcfg.input_features = { + "observation.images.cam0": PolicyFeature(type=FeatureType.VISUAL, shape=(3, s, s)), + "observation.images.cam1": PolicyFeature(type=FeatureType.VISUAL, shape=(3, s, s)), + "observation.state": PolicyFeature(type=FeatureType.STATE, shape=(pcfg.max_state_dim,)), + } + pcfg.output_features = { + "action": PolicyFeature(type=FeatureType.ACTION, shape=(pcfg.max_action_dim,)), + } + pcfg.validate_features() + return TinyVLAPolicy(pcfg), pcfg + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--config", type=Path, required=True) + args = parser.parse_args() + cfg = yaml.safe_load(args.config.read_text()) + + from accelerate import Accelerator + from lerobot.datasets.lerobot_dataset import LeRobotDataset + + accelerator = Accelerator(mixed_precision=cfg.get("mixed_precision", "bf16")) + + # ---- datasets ------------------------------------------------------ + # spec forms: + # {repo_id, weight, root?, episodes?, revision?} — one dataset + # {root_glob, weight, embodiment_group?} — local converted dirs, + # weight is split across matches proportionally to episode count + datasets, weights, names, embodiment_ids = [], [], [], [] + chunk = cfg["policy"]["chunk_size"] + + def add(src, w, name, emb_id): + datasets.append(src) + weights.append(w) + names.append(name) + embodiment_ids.append(emb_id) + accelerator.print( + f"dataset[{len(datasets)-1}] {name}: eps={src.ds.num_episodes} w={w:.4f} emb={emb_id}" + ) + + policy, pcfg = make_policy(cfg) + + from tinyvla.data.mixture import CanonicalSource, WeightedMixtureDataset + + from lerobot.datasets.lerobot_dataset import LeRobotDatasetMetadata + + def make_ds(repo_id, root=None, episodes=None, revision=None): + # delta_timestamps must be set at construction (it feeds DatasetReader), + # and needs fps — read metadata first + meta = LeRobotDatasetMetadata(repo_id, root=root, revision=revision) + return LeRobotDataset( + repo_id, + root=root, + episodes=episodes, + revision=revision, + delta_timestamps={"action": [t / meta.fps for t in range(chunk)]}, + video_backend="torchcodec", + ) + + labels_dir = cfg.get("spatial_labels_dir") + action_space = pcfg.action_space # "native" (A) or "canonical" (B, C) + morph_yaml = cfg.get("morphology_descriptors") + morph_map = {} + if pcfg.conditioning == "morph" and morph_yaml: + import torch as _t + + from tinyvla.modules.embodiment import MORPH_FIELDS + raw = yaml.safe_load(Path(morph_yaml).read_text()) + # normalization applied here (see MORPH_FIELDS comment in embodiment.py) + _sc = {"arm_dof": 0.1, "reach_m": 2, "gripper_width_m": 10, "num_cameras": 1 / 3, + "control_hz": 1 / 30, "joint_lo_mean": 1 / 3.1416, "joint_hi_mean": 1 / 3.1416, + "workspace_x": 2, "workspace_y": 2, "workspace_z": 2, "payload_kg": 0.2} + for krobot, d in raw.items(): + vec = [d.get(f, 0) * _sc.get(f, 1) for f in MORPH_FIELDS] + morph_map[krobot] = _t.tensor(vec, dtype=_t.float32) + + # natural-language robot descriptions for the Qwen prompt (slow path) + prompt_map = {} + if cfg.get("robot_prompts"): + prompt_map = yaml.safe_load(Path(cfg["robot_prompts"]).read_text()) + + # conditioning="morph_qwen": tokenize each robot's text description once, + # fixed-length, for the shared-Qwen text-only morphology encoder + morph_text_map = {} + if pcfg.conditioning == "morph_qwen" and prompt_map: + from transformers import AutoTokenizer as _Tok + + _tok0 = _Tok.from_pretrained(pcfg.lm_model_name) + for krobot, text in prompt_map.items(): + t = _tok0([text], padding="max_length", truncation=True, + max_length=pcfg.morph_text_max_len, return_tensors="pt") + morph_text_map[krobot] = (t["input_ids"][0], t["attention_mask"][0].bool()) + + def wrap(ds, emb_id, morph_key=None): + name = ds.repo_id.split("/")[-1] + store = None + if labels_dir: + from tinyvla.data.spatial_labels import SpatialLabelStore + + store = SpatialLabelStore(labels_dir, name) + if len(store) == 0: + store = None + canon_store, canon_stats = None, None + if action_space == "canonical": + from tinyvla.data.canonical import CanonicalChunkStore + + canon_store = CanonicalChunkStore(name, src_fps=ds.fps, chunk=chunk) + canon_stats = canon_store.compute_stats() + return CanonicalSource( + ds, + embodiment_id=emb_id, + image_size=pcfg.image_size, + max_state_dim=pcfg.max_state_dim, + max_action_dim=pcfg.max_action_dim, + staleness_max_s=pcfg.staleness_max_s, + spatial_labels=store, + action_space=action_space, + canonical_store=canon_store, + canonical_stats=canon_stats, + morphology=morph_map.get(morph_key) if morph_key else None, + # prepend-to-task text conditioning (separate experiment, didn't help) + robot_prompt=prompt_map.get(morph_key) if (morph_key and pcfg.conditioning != "morph_qwen") else None, + morph_text_ids=morph_text_map.get(morph_key, (None, None))[0] if morph_key else None, + morph_text_mask=morph_text_map.get(morph_key, (None, None))[1] if morph_key else None, + # demos are needed either by the MLP demo-encoder or by the rich-slow + # LM sequence (vlm_native), so enable sampling for both + n_support=pcfg.n_support if (pcfg.use_demo_conditioning or pcfg.vlm_native) else 0, + support_other_task=pcfg.support_other_task, + ) + + next_emb = 0 + for spec in cfg["datasets"]: + mkey = spec.get("morph_key") # e.g. "so101"/"bridge"/"rt1" for variant C + if "root_glob" in spec: + roots = sorted(Path(p) for p in __import__("glob").glob(spec["root_glob"])) + subs = [ + make_ds(r.name, root=r) + for r in roots + if (r / "meta" / "info.json").exists() + ] + total_eps = sum(d.num_episodes for d in subs) or 1 + for d in subs: + add(wrap(d, next_emb, mkey), spec["weight"] * d.num_episodes / total_eps, d.root.name, next_emb) + else: + emb = spec.get("embodiment_id", next_emb) + ds = make_ds( + spec["repo_id"], + root=spec.get("root"), + episodes=list(range(spec["episodes"])) if spec.get("episodes") else None, + revision=spec.get("revision"), + ) + add(wrap(ds, emb, mkey), spec["weight"], spec["repo_id"], emb) + next_emb += 1 + + mixture = WeightedMixtureDataset(datasets, weights, seed=cfg.get("seed", 42)) + loader = torch.utils.data.DataLoader( + mixture, + batch_size=cfg["batch_size"], + num_workers=cfg.get("num_workers", 8), + pin_memory=True, + persistent_workers=True, + drop_last=True, + ) + + # tokenizer for task strings (per-source normalization already done in CanonicalSource) + from transformers import AutoTokenizer + + tokenizer = AutoTokenizer.from_pretrained(pcfg.lm_model_name) + + backbone_params = [ + p for n, p in policy.named_parameters() if p.requires_grad and "semantic.vlm" in n + ] + head_params = [ + p for n, p in policy.named_parameters() if p.requires_grad and "semantic.vlm" not in n + ] + groups = [{"params": head_params, "lr": cfg["lr"]}] + if backbone_params: + groups.append({"params": backbone_params, "lr": cfg["lr"] * cfg.get("backbone_lr_mult", 0.1)}) + accelerator.print(f"backbone group: {sum(p.numel() for p in backbone_params)/1e6:.1f}M params at {cfg.get('backbone_lr_mult', 0.1)}x lr") + opt = torch.optim.AdamW(groups, betas=(0.9, 0.95), weight_decay=1e-10) + steps = cfg["steps"] + warmup = cfg.get("warmup_steps", 1000) + + def lr_lambda(s): + if s < warmup: + return s / warmup + p = (s - warmup) / max(1, steps - warmup) + return 0.025 + 0.975 * 0.5 * (1 + math.cos(math.pi * p)) + + sched = torch.optim.lr_scheduler.LambdaLR(opt, lr_lambda) + + policy, opt, loader, sched = accelerator.prepare(policy, opt, loader, sched) + + out_dir = Path(cfg["output_dir"]) + out_dir.mkdir(parents=True, exist_ok=True) + if cfg.get("wandb") and accelerator.is_main_process: + import wandb + + wandb.init(project=cfg["wandb"], config=cfg) + + step, t0 = 0, time.time() + if cfg.get("resume_from"): + from safetensors.torch import load_file + + sd = load_file(Path(cfg["resume_from"]) / "model.safetensors") + missing, unexpected = accelerator.unwrap_model(policy).load_state_dict(sd, strict=False) + step = int(cfg.get("resume_step", 0)) + for _ in range(step): + sched.step() # fast-forward LR schedule + accelerator.print(f"resumed from {cfg['resume_from']} at step {step} " + f"(missing {len(missing)}, unexpected {len(unexpected)})") + grad_accum = cfg.get("grad_accum", 1) + staleness_start = cfg.get("staleness_start_step") + staleness_on = False + data_iter = iter(loader) + while step < steps: + if staleness_start is not None and not staleness_on and step >= staleness_start: + # persistent workers hold dataset copies — rebuild the loader + for src in datasets: + src.staleness_prob = cfg.get("staleness_prob", 0.5) + del data_iter + loader = torch.utils.data.DataLoader( + mixture, + batch_size=cfg["batch_size"], + num_workers=cfg.get("num_workers", 8), + pin_memory=True, + persistent_workers=True, + drop_last=True, + ) + data_iter = iter(loader) + staleness_on = True + accelerator.print(f"staleness augmentation ON at step {step}") + opt.zero_grad() + for _ in range(grad_accum): + try: + batch = next(data_iter) + except StopIteration: + data_iter = iter(loader) + batch = next(data_iter) + tok = tokenizer( + list(batch.pop("task")), + padding=True, + truncation=True, + max_length=pcfg.tokenizer_max_length, + return_tensors="pt", + ) + batch["observation.language.tokens"] = tok["input_ids"] + batch["observation.language.attention_mask"] = tok["attention_mask"].bool() + batch = { + k: v.to(accelerator.device, non_blocking=True) if torch.is_tensor(v) else v + for k, v in batch.items() + } + loss, info = policy(batch) + accelerator.backward(loss / grad_accum) + accelerator.clip_grad_norm_(policy.parameters(), cfg.get("grad_clip", 10.0)) + opt.step() + sched.step() + step += 1 + + if step % cfg.get("log_freq", 50) == 0: + it_s = cfg.get("log_freq", 50) / (time.time() - t0) + t0 = time.time() + accelerator.print(f"step {step}/{steps} loss {info['loss']:.4f} {it_s:.2f} it/s") + if cfg.get("wandb") and accelerator.is_main_process: + wandb.log({"loss": info["loss"], "lr": sched.get_last_lr()[0]}, step=step) + + if step % cfg.get("save_freq", 2000) == 0 and accelerator.is_main_process: + accelerator.unwrap_model(policy).save_pretrained(out_dir / f"step_{step}") + + if accelerator.is_main_process: + accelerator.unwrap_model(policy).save_pretrained(out_dir / "final") + + +if __name__ == "__main__": + main() diff --git a/tinyvla2/scripts/train_stage1_smoke.sh b/tinyvla2/scripts/train_stage1_smoke.sh new file mode 100644 index 0000000000000000000000000000000000000000..bed8b84ac3d5e61fd2c39ccba7e9ccb508c8a6bb --- /dev/null +++ b/tinyvla2/scripts/train_stage1_smoke.sh @@ -0,0 +1,20 @@ +#!/usr/bin/env bash +# Stage 1: overfit smoke on svla_so100_stacking via lerobot-train. +set -euo pipefail + +source ~/tinyvla_venv/bin/activate + +lerobot-train \ + --policy.type=tinyvla \ + --policy.discover_packages_path=tinyvla \ + --policy.device=cuda \ + --policy.push_to_hub=false \ + --dataset.repo_id=lerobot/svla_so100_stacking \ + --dataset.video_backend=torchcodec \ + --batch_size=32 \ + --num_workers=8 \ + --steps=5000 \ + --log_freq=50 \ + --save_freq=2500 \ + --output_dir=outputs/stage1_smoke \ + "$@" diff --git a/tinyvla2/scripts/uav_fewshot.py b/tinyvla2/scripts/uav_fewshot.py new file mode 100644 index 0000000000000000000000000000000000000000..726199a5886e194bcba14689b6de052be0fa3470 --- /dev/null +++ b/tinyvla2/scripts/uav_fewshot.py @@ -0,0 +1,101 @@ +#!/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() diff --git a/tinyvla2/src/tinyvla.egg-info/PKG-INFO b/tinyvla2/src/tinyvla.egg-info/PKG-INFO new file mode 100644 index 0000000000000000000000000000000000000000..8775ad393714c0745e9a4bc7b21f470ecc7c2f08 --- /dev/null +++ b/tinyvla2/src/tinyvla.egg-info/PKG-INFO @@ -0,0 +1,15 @@ +Metadata-Version: 2.4 +Name: tinyvla +Version: 0.1.0 +Summary: Dual-rate VLA: Qwen3-0.6B slow semantic path + ~90M fast visuomotor path for SO-100/SO-101 +Requires-Python: >=3.10 +Requires-Dist: lerobot>=0.4 +Requires-Dist: timm>=1.0.15 +Requires-Dist: torch>=2.6 +Requires-Dist: torchvision +Requires-Dist: transformers>=4.51 +Requires-Dist: accelerate>=1.0 +Requires-Dist: wandb +Requires-Dist: pyarrow +Provides-Extra: dev +Requires-Dist: pytest; extra == "dev" diff --git a/tinyvla2/src/tinyvla.egg-info/SOURCES.txt b/tinyvla2/src/tinyvla.egg-info/SOURCES.txt new file mode 100644 index 0000000000000000000000000000000000000000..26abda1a46185581fd360163d64018af2d6128db --- /dev/null +++ b/tinyvla2/src/tinyvla.egg-info/SOURCES.txt @@ -0,0 +1,7 @@ +pyproject.toml +src/tinyvla.egg-info/PKG-INFO +src/tinyvla.egg-info/SOURCES.txt +src/tinyvla.egg-info/dependency_links.txt +src/tinyvla.egg-info/requires.txt +src/tinyvla.egg-info/top_level.txt +src/tinyvla/modules/vision.py \ No newline at end of file diff --git a/tinyvla2/src/tinyvla.egg-info/dependency_links.txt b/tinyvla2/src/tinyvla.egg-info/dependency_links.txt new file mode 100644 index 0000000000000000000000000000000000000000..8b137891791fe96927ad78e64b0aad7bded08bdc --- /dev/null +++ b/tinyvla2/src/tinyvla.egg-info/dependency_links.txt @@ -0,0 +1 @@ + diff --git a/tinyvla2/src/tinyvla.egg-info/requires.txt b/tinyvla2/src/tinyvla.egg-info/requires.txt new file mode 100644 index 0000000000000000000000000000000000000000..6206b7b01bd66ee181c69d7ecfd713cc6c7ba14e --- /dev/null +++ b/tinyvla2/src/tinyvla.egg-info/requires.txt @@ -0,0 +1,11 @@ +lerobot>=0.4 +timm>=1.0.15 +torch>=2.6 +torchvision +transformers>=4.51 +accelerate>=1.0 +wandb +pyarrow + +[dev] +pytest diff --git a/tinyvla2/src/tinyvla.egg-info/top_level.txt b/tinyvla2/src/tinyvla.egg-info/top_level.txt new file mode 100644 index 0000000000000000000000000000000000000000..c723e8829bdae705410ea629b7bb4c33e21b6809 --- /dev/null +++ b/tinyvla2/src/tinyvla.egg-info/top_level.txt @@ -0,0 +1 @@ +tinyvla diff --git a/tinyvla2/src/tinyvla/__init__.py b/tinyvla2/src/tinyvla/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9b1e318d33509f5fee7f7c755ff127b42f8a6c11 --- /dev/null +++ b/tinyvla2/src/tinyvla/__init__.py @@ -0,0 +1,9 @@ +"""TinyVLA: dual-rate VLA (Qwen3-0.6B slow path + ~90M fast visuomotor path). + +Importing this package registers the "tinyvla" policy with LeRobot +(use --policy.type=tinyvla --policy.discover_packages_path=tinyvla). +""" + +from .configuration_tinyvla import TinyVLAConfig + +__all__ = ["TinyVLAConfig"] diff --git a/tinyvla2/src/tinyvla/__pycache__/__init__.cpython-312.pyc b/tinyvla2/src/tinyvla/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e0c1f0d3d80e698b29fb45092c50a27ced9b5602 Binary files /dev/null and b/tinyvla2/src/tinyvla/__pycache__/__init__.cpython-312.pyc differ diff --git a/tinyvla2/src/tinyvla/__pycache__/configuration_tinyvla.cpython-312.pyc b/tinyvla2/src/tinyvla/__pycache__/configuration_tinyvla.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2b00482a17f6b773d5b29257fa0b758b0a11449a Binary files /dev/null and b/tinyvla2/src/tinyvla/__pycache__/configuration_tinyvla.cpython-312.pyc differ diff --git a/tinyvla2/src/tinyvla/__pycache__/modeling_tinyvla.cpython-312.pyc b/tinyvla2/src/tinyvla/__pycache__/modeling_tinyvla.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..78f0a98a83b1d28686455955df0733c8078c5763 Binary files /dev/null and b/tinyvla2/src/tinyvla/__pycache__/modeling_tinyvla.cpython-312.pyc differ diff --git a/tinyvla2/src/tinyvla/__pycache__/processor_tinyvla.cpython-312.pyc b/tinyvla2/src/tinyvla/__pycache__/processor_tinyvla.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..50063fccb3d14b6a1270655da9a14dfe15b2b438 Binary files /dev/null and b/tinyvla2/src/tinyvla/__pycache__/processor_tinyvla.cpython-312.pyc differ diff --git a/tinyvla2/src/tinyvla/configuration_tinyvla.py b/tinyvla2/src/tinyvla/configuration_tinyvla.py new file mode 100644 index 0000000000000000000000000000000000000000..5cbfefe9b28919d1a15f815dae4ca3ca867c88dc --- /dev/null +++ b/tinyvla2/src/tinyvla/configuration_tinyvla.py @@ -0,0 +1,160 @@ +"""TinyVLA policy config (LeRobot plugin, --policy.type=tinyvla).""" + +from __future__ import annotations + +from dataclasses import dataclass, field + +from lerobot.configs import FeatureType, NormalizationMode, PreTrainedConfig +from lerobot.optim import AdamWConfig +from lerobot.optim.schedulers import CosineDecayWithWarmupSchedulerConfig + + +@PreTrainedConfig.register_subclass("tinyvla") +@dataclass +class TinyVLAConfig(PreTrainedConfig): + # Input / output structure + n_obs_steps: int = 1 + chunk_size: int = 50 + n_action_steps: int = 50 + + normalization_mapping: dict[str, NormalizationMode] = field( + default_factory=lambda: { + "VISUAL": NormalizationMode.IDENTITY, + "STATE": NormalizationMode.MEAN_STD, + "ACTION": NormalizationMode.MEAN_STD, + } + ) + + # Cross-embodiment padding + max_state_dim: int = 16 + max_action_dim: int = 8 + num_embodiments: int = 16 + max_cameras: int = 3 + + # TinyVLA-2 A/B/C variant switch + # conditioning "id" -> learned embodiment embedding (variants A, B) + # conditioning "morph" -> descriptor MLP -> morphology tokens (variant C) + conditioning: str = "id" # "id" | "morph" (MLP over numeric fields) | "morph_qwen" (text via shared Qwen) + morph_tokens: int = 2 + morph_to_slow: bool = False # also inject morphology into the slow Qwen prefix + num_morph_readout: int = 4 # for morph_qwen: readout slots on the robot-text encode + morph_text_max_len: int = 32 + # in-context demo conditioning: K (obs, action) exemplar pairs, no gradient + # updates for the new robot at deployment. Composable with morph/morph_qwen. + use_demo_conditioning: bool = False + n_support: int = 3 + # sample demo exemplars from a DIFFERENT task of the same robot, so the + # exemplar carries robot-dynamics info the identity channels don't duplicate + support_other_task: bool = False + # demo-encoder capacity (v3): wider path so exemplar content isn't destroyed + demo_tokens_per_example: int = 4 + demo_visual_groups: int = 4 + demo_action_keys: int = 8 + demo_hidden_mult: int = 4 + # Rich-slow mode: demos + robot text + instruction + image all go into the LM + # sequence (so demos live in the ATTENTION STREAM — the fix for v1-v3's + # "placed alongside → ignored" failure), and the LM emits latent tokens that + # condition the fast path. DUAL-RATE IS PRESERVED: the fast path still has its + # own spatial tokens + proprio + expert, and the LM latents are cacheable + # exactly like the old semantic latent. + vlm_native: bool = False + num_action_readout: int = 16 + demo_action_keys_lm: int = 4 + # DEMO-ONLY: strip every other robot-identity channel (embodiment-ID token, + # morphology descriptor/text, and the per-source normalization that leaks + # identity) so the demo exemplars are the ONLY way to know which robot this + # is. Four prior attempts failed because identity was available elsewhere and + # the gradient had no reason to read the demos. + demo_only: bool = False + # action_space "native" (A) vs "canonical" (B, C: base-frame EE deltas) + action_space: str = "native" + + # Vision (shared Qwen3.5 vision tower for both paths) + image_size: int = 256 + freeze_vision_encoder: bool = True # Stage 1; unfreeze at 0.1x lr in Stage 2 + + # Slow path (Qwen3.5 VLM with native vision tower) + lm_model_name: str = "Qwen/Qwen3.5-0.8B" + lm_num_layers: int = 12 # L/2 layer skip (24-layer language stack) + num_readout_tokens: int = 8 + freeze_lm: bool = True + tokenizer_max_length: int = 48 + pad_language_to: str = "longest" + + # Fast path (action expert) + expert_dim: int = 512 + expert_layers: int = 12 + expert_heads: int = 8 + + # Ablation switches (Stage-2 gate baselines) + use_semantic_latent: bool = True + use_spatial_tokens: bool = True + + # Spatial distillation aux task (Stage 3, XS-VLA style) + spatial_vocab: int = 1024 # 32x32 grid cells + spatial_loss_weight: float = 0.0 # 0 disables the aux head/loss + + # Flow matching + num_denoise_steps: int = 10 + flow_beta_alpha: float = 1.5 + flow_beta_beta: float = 1.0 + min_period: float = 4e-3 + max_period: float = 4.0 + + # Staleness augmentation (Stage 2b): probability of conditioning on a latent + # computed from an earlier frame. Consumed by the training loop, not here. + staleness_prob: float = 0.0 + staleness_max_s: float = 2.0 + + # Training presets + optimizer_lr: float = 1e-4 + optimizer_betas: tuple[float, float] = (0.9, 0.95) + optimizer_eps: float = 1e-8 + optimizer_weight_decay: float = 1e-10 + optimizer_grad_clip_norm: float = 10.0 + scheduler_warmup_steps: int = 1_000 + scheduler_decay_steps: int = 40_000 + scheduler_decay_lr: float = 2.5e-6 + + def __post_init__(self): + super().__post_init__() + if self.n_action_steps > self.chunk_size: + raise ValueError( + f"n_action_steps ({self.n_action_steps}) must be <= chunk_size ({self.chunk_size})" + ) + + def validate_features(self) -> None: + image_keys = [k for k, v in self.input_features.items() if v.type is FeatureType.VISUAL] + if not image_keys: + raise ValueError("TinyVLA requires at least one image feature") + if len(image_keys) > self.max_cameras: + raise ValueError(f"got {len(image_keys)} cameras, max_cameras={self.max_cameras}") + + def get_optimizer_preset(self) -> AdamWConfig: + return AdamWConfig( + lr=self.optimizer_lr, + betas=self.optimizer_betas, + eps=self.optimizer_eps, + weight_decay=self.optimizer_weight_decay, + grad_clip_norm=self.optimizer_grad_clip_norm, + ) + + def get_scheduler_preset(self): + return CosineDecayWithWarmupSchedulerConfig( + peak_lr=self.optimizer_lr, + decay_lr=self.scheduler_decay_lr, + num_warmup_steps=self.scheduler_warmup_steps, + num_decay_steps=self.scheduler_decay_steps, + ) + + @property + def observation_delta_indices(self) -> list: + return [0] + + @property + def action_delta_indices(self) -> list: + return list(range(self.chunk_size)) + + @property + def reward_delta_indices(self) -> None: + return None diff --git a/tinyvla2/src/tinyvla/data/__init__.py b/tinyvla2/src/tinyvla/data/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/tinyvla2/src/tinyvla/data/__pycache__/__init__.cpython-312.pyc b/tinyvla2/src/tinyvla/data/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3d346b303bf7272d09300304fa6a9b686fbe4ba3 Binary files /dev/null and b/tinyvla2/src/tinyvla/data/__pycache__/__init__.cpython-312.pyc differ diff --git a/tinyvla2/src/tinyvla/data/__pycache__/canonical.cpython-312.pyc b/tinyvla2/src/tinyvla/data/__pycache__/canonical.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ac6a18ce154a45fb73f8560b478a78d7e445c2e2 Binary files /dev/null and b/tinyvla2/src/tinyvla/data/__pycache__/canonical.cpython-312.pyc differ diff --git a/tinyvla2/src/tinyvla/data/__pycache__/kinematics_so101.cpython-312.pyc b/tinyvla2/src/tinyvla/data/__pycache__/kinematics_so101.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1915f3a57c5b359a66279fcc225535137062fae9 Binary files /dev/null and b/tinyvla2/src/tinyvla/data/__pycache__/kinematics_so101.cpython-312.pyc differ diff --git a/tinyvla2/src/tinyvla/data/__pycache__/mixture.cpython-312.pyc b/tinyvla2/src/tinyvla/data/__pycache__/mixture.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a91858bc716e56390dc9ab0a15f63c866692afdf Binary files /dev/null and b/tinyvla2/src/tinyvla/data/__pycache__/mixture.cpython-312.pyc differ diff --git a/tinyvla2/src/tinyvla/data/__pycache__/spatial_labels.cpython-312.pyc b/tinyvla2/src/tinyvla/data/__pycache__/spatial_labels.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3f1fc0e2c401de93285522b2d00c23ef3a49eb45 Binary files /dev/null and b/tinyvla2/src/tinyvla/data/__pycache__/spatial_labels.cpython-312.pyc differ diff --git a/tinyvla2/src/tinyvla/data/canonical.py b/tinyvla2/src/tinyvla/data/canonical.py new file mode 100644 index 0000000000000000000000000000000000000000..91490db8e412f17e3c3440ec5d9666f96c31f2ce --- /dev/null +++ b/tinyvla2/src/tinyvla/data/canonical.py @@ -0,0 +1,129 @@ +"""Canonical base-frame EE action assembly from EE-pose sidecars (variants B, C). + +Canonical action chunk (T steps @ 10Hz) per query frame i of an episode: + knot frames: i, i+step, ..., i+(T-1)*step (step = round(src_fps / 10)) + Δ_0 = a_pose[i] ⊖ s_pose[i] (state anchor: target vs measured) + Δ_k = a_pose[knot_k] ⊖ a_pose[knot_{k-1}] (prev-target, telescoping) +where pose ⊖ pose = [Δpos (base frame) | Δrot = log(R_k @ R_{k-1}^T) rotvec | grip_abs]. + +Rotation deltas are ALWAYS composed via matrices then logged (never differenced +parameterizations). Gripper is absolute openness in [0,1] (per-dataset q1-q99). +Normalization is quantile (q01/q99 -> [-1,1]) with a per-dim floor. + +Sidecar parquet cols (from scripts/build_ee_sidecar.py): + episode_index, frame_index, + s_p{xyz} s_q{xyzw} (state EE pose), a_p{xyz} a_q{xyzw} (target EE pose), grip +""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np + +SIDECAR_DIR = Path.home() / "tinyvla_data" / "ee_sidecar" +CANON_DIM = 7 # dx dy dz + rotvec(3) + grip +TARGET_HZ = 10.0 + + +def _pose_delta(p_prev, q_prev, p_cur, q_cur): + """Base-frame left-delta between two poses -> (6,) [dpos(3), drotvec(3)].""" + from scipy.spatial.transform import Rotation + + dpos = p_cur - p_prev + R_prev = Rotation.from_quat(q_prev) + R_cur = Rotation.from_quat(q_cur) + drot = (R_cur * R_prev.inv()).as_rotvec() # base-frame left-multiply + return np.concatenate([dpos, drot]) + + +class CanonicalChunkStore: + """Loads one dataset's EE-pose sidecar; assembles canonical chunks + stats.""" + + def __init__(self, dataset_name: str, src_fps: float, chunk: int = 50, + grip_lo: float | None = None, grip_hi: float | None = None): + import pyarrow.parquet as pq + + path = SIDECAR_DIR / f"{dataset_name}.parquet" + t = pq.read_table(path).to_pydict() + self.chunk = chunk + self.step = max(1, round(src_fps / TARGET_HZ)) + # index rows by (episode, frame) + self.ep = np.asarray(t["episode_index"]) + self.fr = np.asarray(t["frame_index"]) + self.s_p = np.stack([t["s_px"], t["s_py"], t["s_pz"]], 1) + self.s_q = np.stack([t["s_qx"], t["s_qy"], t["s_qz"], t["s_qw"]], 1) + self.a_p = np.stack([t["a_px"], t["a_py"], t["a_pz"]], 1) + self.a_q = np.stack([t["a_qx"], t["a_qy"], t["a_qz"], t["a_qw"]], 1) + self.grip = np.asarray(t["grip"], dtype=np.float64) + # (ep, frame) -> row + self._row = {(int(e), int(f)): i for i, (e, f) in enumerate(zip(self.ep, self.fr))} + # gripper openness range (q1-q99) unless provided + self.grip_lo = grip_lo if grip_lo is not None else float(np.quantile(self.grip, 0.01)) + self.grip_hi = grip_hi if grip_hi is not None else float(np.quantile(self.grip, 0.99)) + self._ep_last_frame = {} + for e, f in zip(self.ep, self.fr): + e, f = int(e), int(f) + if f > self._ep_last_frame.get(e, -1): + self._ep_last_frame[e] = f + + def _grip_norm(self, raw): + span = max(self.grip_hi - self.grip_lo, 1e-6) + return float(np.clip((raw - self.grip_lo) / span, 0.0, 1.0)) + + def chunk_for(self, episode: int, frame: int) -> np.ndarray: + """Canonical (chunk, 7). Clamps knots past episode end (zero-order hold).""" + last = self._ep_last_frame.get(episode, frame) + out = np.zeros((self.chunk, CANON_DIM), dtype=np.float32) + + def pose_at(fr, which): + fr = min(fr, last) + r = self._row.get((episode, fr)) + if r is None: + return None + if which == "a": + return self.a_p[r], self.a_q[r], self.grip[r] + return self.s_p[r], self.s_q[r], self.grip[r] + + # anchor: target[i] vs state[i] + s = pose_at(frame, "s") + a0 = pose_at(frame, "a") + if s is None or a0 is None: + return out + out[0, :6] = _pose_delta(s[0], s[1], a0[0], a0[1]) + out[0, 6] = self._grip_norm(a0[2]) + prev = a0 + for k in range(1, self.chunk): + cur = pose_at(frame + k * self.step, "a") + if cur is None: + cur = prev # zero-order hold + out[k, :6] = _pose_delta(prev[0], prev[1], cur[0], cur[1]) + out[k, 6] = self._grip_norm(cur[2]) + prev = cur + return out + + def compute_stats(self, n_samples: int = 2000, seed: int = 0) -> dict: + """q01/q99 per canonical dim over sampled chunks (for [-1,1] scaling).""" + rng = np.random.default_rng(seed) + pairs = list(self._row.keys()) + pick = rng.choice(len(pairs), size=min(n_samples, len(pairs)), replace=False) + acc = [] + for j in pick: + e, f = pairs[j] + acc.append(self.chunk_for(e, f)) + acc = np.concatenate(acc, 0) # (N*chunk, 7) + q01 = np.quantile(acc, 0.01, axis=0) + q99 = np.quantile(acc, 0.99, axis=0) + return {"q01": q01.tolist(), "q99": q99.tolist(), + "grip_lo": self.grip_lo, "grip_hi": self.grip_hi} + + +def quantile_normalize(chunk: np.ndarray, q01, q99, floor_frac: float = 0.01) -> np.ndarray: + """Map [q01,q99] -> [-1,1] per dim with a floor on the scale to avoid blowups.""" + q01 = np.asarray(q01) + q99 = np.asarray(q99) + span = q99 - q01 + floor = floor_frac * np.median(np.abs(np.concatenate([q01, q99])) + 1e-6) + span = np.maximum(span, floor) + mid = 0.5 * (q01 + q99) + return (2.0 * (chunk - mid) / span).astype(np.float32) diff --git a/tinyvla2/src/tinyvla/data/kinematics_so101.py b/tinyvla2/src/tinyvla/data/kinematics_so101.py new file mode 100644 index 0000000000000000000000000000000000000000..ba48881c56b1cf56e7ceb738baab34b09d22b604 --- /dev/null +++ b/tinyvla2/src/tinyvla/data/kinematics_so101.py @@ -0,0 +1,37 @@ +"""SO-100/SO-101 forward kinematics for canonical EE-space action derivation. + +Joint order matches community datasets' action/state feature names +(main_shoulder_pan, main_shoulder_lift, main_elbow_flex, main_wrist_flex, +main_wrist_roll, main_gripper) — the URDF joint names are the same sans prefix. +Inputs are in DEGREES (as recorded); placo wants float64. +""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np + +SO101_JOINTS = ["shoulder_pan", "shoulder_lift", "elbow_flex", "wrist_flex", "wrist_roll", "gripper"] +DEFAULT_URDF = Path(__file__).resolve().parents[3] / "assets" / "so_arm" / "so101_new_calib.urdf" +TARGET_FRAME = "gripper_frame_link" + + +class SO101FK: + def __init__(self, urdf_path: str | Path = DEFAULT_URDF, target_frame: str = TARGET_FRAME): + from lerobot.model import RobotKinematics + + self.kin = RobotKinematics( + urdf_path=str(urdf_path), target_frame_name=target_frame, joint_names=SO101_JOINTS + ) + + def ee_pose(self, joints_deg: np.ndarray) -> np.ndarray: + """(6,) joint degrees -> 4x4 base->gripper homogeneous transform.""" + return self.kin.forward_kinematics(np.asarray(joints_deg, dtype=np.float64)) + + def ee_pos_quat(self, joints_deg: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + """-> (pos xyz, quat xyzw).""" + from scipy.spatial.transform import Rotation + + T = self.ee_pose(joints_deg) + return T[:3, 3].copy(), Rotation.from_matrix(T[:3, :3]).as_quat() diff --git a/tinyvla2/src/tinyvla/data/mixture.py b/tinyvla2/src/tinyvla/data/mixture.py new file mode 100644 index 0000000000000000000000000000000000000000..4bd291b37069d78afc368cdb183d9ef06d73fdf4 --- /dev/null +++ b/tinyvla2/src/tinyvla/data/mixture.py @@ -0,0 +1,275 @@ +"""Weighted mixture over heterogeneous LeRobotDatasets. + +CanonicalSource maps every dataset to one schema: + observation.images.cam0 (3, S, S) float [0,1] primary/scene camera + observation.images.cam1 (3, S, S) float [0,1] wrist camera or zeros + camera_mask (2,) bool + observation.state (max_state_dim,) normalized, zero-padded + action (chunk, max_action_dim) normalized, zero-padded + action_dim_mask (max_action_dim,) bool — real action dims + action_is_pad (chunk,) bool — episode-end padding + embodiment_id () long + task str + +Normalization is per-source (each dataset's own mean/std), so the lerobot +pipeline normalize step must NOT be applied on top. +""" + +from __future__ import annotations + +import re + +import torch +import torch.nn.functional as F +from torch.utils.data import Dataset + +PRIMARY_RE = re.compile(r"top|front|base|main|scene|image_0$|images\.image$|image2$") +WRIST_RE = re.compile(r"wrist|gripper|hand|image_1$") + + +def _norm(x: torch.Tensor, stats: dict | None) -> torch.Tensor: + if not stats: + return x + mean = torch.as_tensor(stats["mean"], dtype=x.dtype) + std = torch.as_tensor(stats["std"], dtype=x.dtype).clamp(min=1e-6) + return (x - mean) / std + + +def _pad(x: torch.Tensor, dim: int) -> torch.Tensor: + if x.shape[-1] >= dim: + return x[..., :dim] + return F.pad(x, (0, dim - x.shape[-1])) + + +class CanonicalSource(Dataset): + def __init__( + self, + ds, # LeRobotDataset with delta_timestamps set for action chunk + embodiment_id: int, + image_size: int = 256, + max_state_dim: int = 16, + max_action_dim: int = 8, + staleness_prob: float = 0.0, + staleness_max_s: float = 2.0, + spatial_labels=None, # SpatialLabelStore | None + action_space: str = "native", # "native" (A) or "canonical" (B, C) + canonical_store=None, # CanonicalChunkStore for this dataset + canonical_stats=None, # {"q01","q99"} for quantile norm + morphology=None, # (16,) descriptor tensor for variant C + robot_prompt=None, # natural-language robot description prepended to task + morph_text_ids=None, # (L,) tokenized robot description (conditioning="morph_qwen") + morph_text_mask=None, # (L,) bool + n_support: int = 0, # in-context demo conditioning: K (obs, action) exemplars + support_other_task: bool = False, # sample exemplars from a DIFFERENT task of the same robot + ): + self.ds = ds + self.embodiment_id = embodiment_id + self.image_size = image_size + self.max_state_dim = max_state_dim + self.max_action_dim = max_action_dim + self.staleness_prob = staleness_prob + self.staleness_max_s = staleness_max_s + self.spatial_labels = spatial_labels + self.action_space = action_space + self.canonical_store = canonical_store + self.canonical_stats = canonical_stats + self.morphology = morphology + self.robot_prompt = robot_prompt + self.morph_text_ids = morph_text_ids + self.morph_text_mask = morph_text_mask + self.n_support = n_support + self.support_other_task = support_other_task + self._task_index = None # lazily built {task_index: [frame idx, ...]} + + image_keys = sorted(k for k in ds.meta.features if k.startswith("observation.images")) + primary = next((k for k in image_keys if PRIMARY_RE.search(k)), image_keys[0]) + wrist = next((k for k in image_keys if k != primary and WRIST_RE.search(k)), None) + if wrist is None and len(image_keys) > 1: + wrist = next(k for k in image_keys if k != primary) + self.cam_keys = (primary, wrist) + + self.stats = ds.meta.stats + self.action_dim = ds.meta.features["action"]["shape"][0] + self.state_dim = ds.meta.features.get("observation.state", {}).get("shape", (0,))[0] + + # some converted repos have fewer parquet rows than metadata claims — + # clamp so sampling never goes out of bounds + try: + actual_rows = len(ds.reader.hf_dataset) + except AttributeError: + actual_rows = len(ds) + if actual_rows != len(ds): + import logging + + logging.warning( + f"{ds.repo_id}: meta frames {len(ds)} != parquet rows {actual_rows}; clamping" + ) + self._len = min(len(ds), actual_rows) + + def __len__(self): + return self._len + + def _build_task_index(self): + """{task_index: np.array of frame indices} — lets us pull demo exemplars + from a DIFFERENT task of the same robot, so the exemplar carries info + (how this robot moves) that the identity channels don't already provide.""" + import numpy as np + + try: + hf = self.ds.reader.hf_dataset.with_format("numpy") + tidx = np.asarray(hf["task_index"])[: self._len] + except Exception: + self._task_index = {} + return + self._task_index = {int(t): np.flatnonzero(tidx == t) for t in np.unique(tidx)} + + def _sample_support_idx(self, query_task: int) -> int: + """Frame index for one support exemplar.""" + import numpy as np + + if not self.support_other_task: + return int(torch.randint(0, self._len, (1,)).item()) + if self._task_index is None: + self._build_task_index() + others = [t for t in self._task_index if t != query_task and len(self._task_index[t]) > 0] + if not others: # single-task dataset → fall back to any frame + return int(torch.randint(0, self._len, (1,)).item()) + t = others[int(torch.randint(0, len(others), (1,)).item())] + pool = self._task_index[t] + return int(pool[int(torch.randint(0, len(pool), (1,)).item())]) + + def _img(self, x: torch.Tensor) -> torch.Tensor: + if x.shape[-2:] != (self.image_size, self.image_size): + x = F.interpolate( + x[None], size=(self.image_size, self.image_size), mode="bilinear", align_corners=False + )[0] + return x + + def __getitem__(self, idx): + item = self.ds[idx] + out = {} + primary, wrist = self.cam_keys + out["observation.images.cam0"] = self._img(item[primary]) + if wrist is not None: + out["observation.images.cam1"] = self._img(item[wrist]) + cam_mask = torch.tensor([True, True]) + else: + out["observation.images.cam1"] = torch.zeros_like(out["observation.images.cam0"]) + cam_mask = torch.tensor([True, False]) + out["camera_mask"] = cam_mask + + if self.state_dim: + state = _norm(item["observation.state"], self.stats.get("observation.state")) + else: + state = torch.zeros(1) + out["observation.state"] = _pad(state, self.max_state_dim) + + if self.action_space == "canonical" and self.canonical_store is not None: + from .canonical import CANON_DIM, quantile_normalize + + ep = int(item["episode_index"]) + ep_start = int(self.ds.meta.episodes["dataset_from_index"][ep]) + raw = self.canonical_store.chunk_for(ep, idx - ep_start) # (chunk, 7) + raw = quantile_normalize(raw, self.canonical_stats["q01"], self.canonical_stats["q99"]) + action = torch.from_numpy(raw) + out["action"] = _pad(action, self.max_action_dim) + mask = torch.zeros(self.max_action_dim, dtype=torch.bool) + mask[:CANON_DIM] = True + out["action_dim_mask"] = mask + else: + action = _norm(item["action"], self.stats.get("action")) + out["action"] = _pad(action, self.max_action_dim) + mask = torch.zeros(self.max_action_dim, dtype=torch.bool) + mask[: self.action_dim] = True + out["action_dim_mask"] = mask + if "action_is_pad" in item: + out["action_is_pad"] = item["action_is_pad"] + + out["embodiment_id"] = torch.tensor(self.embodiment_id, dtype=torch.long) + if self.morphology is not None: + out["morphology"] = self.morphology + task = item.get("task") or "" + out["task"] = f"{self.robot_prompt} {task}" if self.robot_prompt else task + if self.morph_text_ids is not None: + out["morph_text_ids"] = self.morph_text_ids + out["morph_text_mask"] = self.morph_text_mask + + cell = -1 + if self.spatial_labels is not None: + ep = int(item["episode_index"]) + ep_start = int(self.ds.meta.episodes["dataset_from_index"][ep]) + cell = self.spatial_labels.lookup(ep, idx - ep_start) + out["spatial_cell"] = torch.tensor(cell, dtype=torch.long) + + # staleness augmentation: with prob p, the semantic latent is computed + # from a frame up to staleness_max_s earlier in the same episode. + # latent_image is always emitted when enabled so batch keys are uniform. + if self.staleness_prob > 0: + out["latent_image"] = out["observation.images.cam0"] + if torch.rand(1).item() < self.staleness_prob: + ep = int(item["episode_index"]) + ep_start = int(self.ds.meta.episodes["dataset_from_index"][ep]) + k = int(torch.randint(0, int(self.staleness_max_s * self.ds.fps) + 1, (1,)).item()) + stale_idx = max(ep_start, idx - k) + if stale_idx != idx: + stale = self.ds[stale_idx] + out["latent_image"] = self._img(stale[self.cam_keys[0]]) + + # in-context demo conditioning: K random (obs, action) exemplars from this + # SAME source (no gradient updates involved at deployment — this is what + # gets fed at inference for a robot with a handful of demos and no FT) + if self.n_support > 0: + sup_imgs, sup_acts, sup_masks = [], [], [] + q_task = int(item.get("task_index", -1)) + for _ in range(self.n_support): + sidx = self._sample_support_idx(q_task) + sitem = self.ds[sidx] + sup_imgs.append(self._img(sitem[self.cam_keys[0]])) + if self.action_space == "canonical" and self.canonical_store is not None: + from .canonical import CANON_DIM, quantile_normalize + + sep = int(sitem["episode_index"]) + sep_start = int(self.ds.meta.episodes["dataset_from_index"][sep]) + sraw = self.canonical_store.chunk_for(sep, sidx - sep_start) + sraw = quantile_normalize(sraw, self.canonical_stats["q01"], self.canonical_stats["q99"]) + sact = _pad(torch.from_numpy(sraw), self.max_action_dim) + smask = torch.zeros(self.max_action_dim, dtype=torch.bool) + smask[:CANON_DIM] = True + else: + sact = _pad(_norm(sitem["action"], self.stats.get("action")), self.max_action_dim) + smask = torch.zeros(self.max_action_dim, dtype=torch.bool) + smask[: self.action_dim] = True + sup_acts.append(sact) + sup_masks.append(smask) + out["support_images"] = torch.stack(sup_imgs) # (K, 3, H, W) + out["support_actions"] = torch.stack(sup_acts) # (K, chunk, max_action_dim) + out["support_action_mask"] = torch.stack(sup_masks) # (K, max_action_dim) + return out + + +class WeightedMixtureDataset(Dataset): + """Samples source i with probability weights[i], then a uniform frame.""" + + def __init__(self, sources: list[CanonicalSource], weights: list[float], seed: int = 42): + assert len(sources) == len(weights) + total = sum(weights) + self.sources = sources + self.weights = torch.tensor([w / total for w in weights], dtype=torch.double) + self.seed = seed + self._len = sum(len(s) for s in sources) + + def __len__(self): + return self._len + + def __getitem__(self, index): + g = torch.Generator().manual_seed(self.seed + index) + for attempt in range(3): + src_idx = int(torch.multinomial(self.weights, 1, generator=g).item()) + src = self.sources[src_idx] + frame_idx = int(torch.randint(len(src), (1,), generator=g).item()) + try: + return src[frame_idx] + except (IndexError, KeyError, RuntimeError): # corrupt frame/video — resample + if attempt == 2: + raise + raise RuntimeError("unreachable") diff --git a/tinyvla2/src/tinyvla/data/spatial_labels.py b/tinyvla2/src/tinyvla/data/spatial_labels.py new file mode 100644 index 0000000000000000000000000000000000000000..fc9c2582d7197fd0c135840ddd56afd16bd3dc4a --- /dev/null +++ b/tinyvla2/src/tinyvla/data/spatial_labels.py @@ -0,0 +1,44 @@ +"""Spatial-distillation label store. + +Labels come from scripts/label_spatial.py parquets: one row per labeled frame +(dataset, episode_index, frame_index RELATIVE to episode start, cell in +[0, GRID*GRID) or -1). Frames were labeled every `stride`; lookup rounds a +query frame to the nearest labeled one. +""" + +from __future__ import annotations + +from pathlib import Path + +GRID = 32 +STRIDE = 10 + + +class SpatialLabelStore: + def __init__(self, labels_dir: str | Path, dataset_name: str): + import pyarrow.parquet as pq + + self.table = {} + path = Path(labels_dir) / f"{dataset_name}.parquet" + if path.exists(): + t = pq.read_table(path) + for ep, fi, cell, ok in zip( + t["episode_index"].to_pylist(), + t["frame_index"].to_pylist(), + t["cell"].to_pylist(), + t["confidence_ok"].to_pylist(), + ): + if ok and cell >= 0: + self.table[(int(ep), int(fi))] = int(cell) + + def __len__(self): + return len(self.table) + + def lookup(self, episode_index: int, rel_frame: int) -> int: + """Nearest labeled frame's cell, or -1.""" + base = round(rel_frame / STRIDE) * STRIDE + for fi in (base, base - STRIDE, base + STRIDE): + cell = self.table.get((episode_index, fi)) + if cell is not None: + return cell + return -1 diff --git a/tinyvla2/src/tinyvla/modeling_tinyvla.py b/tinyvla2/src/tinyvla/modeling_tinyvla.py new file mode 100644 index 0000000000000000000000000000000000000000..35cf9692f22ae49fac8fb9f716db1dad1677f0ae --- /dev/null +++ b/tinyvla2/src/tinyvla/modeling_tinyvla.py @@ -0,0 +1,286 @@ +"""TinyVLA policy: dual-rate VLA on a single shared Qwen3.5 vision encoder. + +Slow path: image -> Qwen3.5 vision tower -> language stack (first 12 layers) + +instruction -> 8-token semantic latent (cached at deployment, refreshed at +0.5-2 Hz). + +Fast path: Qwen3.5 vision tower spatial tokens (64/camera, no language stack) ++ proprio + embodiment token + cached semantic latent -> ~40M flow-matching +action expert -> 50-step action chunk. + +During training both paths run jointly (single-rate); staleness augmentation +is added by the Stage-2 training loop. +""" + +from __future__ import annotations + +from collections import deque + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from lerobot.configs import FeatureType +from lerobot.policies.pretrained import PreTrainedPolicy +from lerobot.utils.constants import ( + ACTION, + OBS_LANGUAGE_ATTENTION_MASK, + OBS_LANGUAGE_TOKENS, + OBS_STATE, +) + +from .configuration_tinyvla import TinyVLAConfig +from .modules.embodiment import EmbodimentEmbedding, pad_vector +from .modules.expert import FlowMatchingExpert +from .modules.semantic import SemanticPath + + +class TinyVLAPolicy(PreTrainedPolicy): + config_class = TinyVLAConfig + name = "tinyvla" + + def __init__(self, config: TinyVLAConfig, dataset_stats=None, dataset_meta=None, **kwargs): + super().__init__(config) + self.config = config + + self.image_keys = sorted( + k for k, v in config.input_features.items() if v.type is FeatureType.VISUAL + ) + state_ft = config.input_features.get(OBS_STATE) + self.state_dim = state_ft.shape[0] if state_ft is not None else 0 + self.action_dim = config.output_features[ACTION].shape[0] + + d = config.expert_dim + + self.semantic = SemanticPath( + model_name=config.lm_model_name, + num_layers=config.lm_num_layers, + num_readout=config.num_readout_tokens, + out_dim=d, + image_size=config.image_size, + freeze_lm=config.freeze_lm, + freeze_vision=config.freeze_vision_encoder, + morph_to_slow=config.morph_to_slow, + morph_dim=16, + morph_qwen=(config.conditioning == "morph_qwen"), + num_morph_readout=config.num_morph_readout, + vlm_native=config.vlm_native, + num_action_readout=config.num_action_readout, + max_state_dim=config.max_state_dim, + max_action_dim=config.max_action_dim, + ) + + # fast-path conditioning projections + self.spatial_proj = nn.Linear(self.semantic.visual_hidden_size, d) + self.camera_emb = nn.Embedding(config.max_cameras, d) + self.state_proj = nn.Linear(config.max_state_dim, d) + self.embodiment_emb = EmbodimentEmbedding(config.num_embodiments, d) + if config.conditioning == "morph": + from .modules.embodiment import MorphologyEncoder + + self.morph_encoder = MorphologyEncoder(d, n_tokens=config.morph_tokens) + if config.use_demo_conditioning: + from .modules.embodiment import DemoConditioner + + self.demo_encoder = DemoConditioner( + d, + visual_dim=self.semantic.visual_hidden_size, + action_dim=config.max_action_dim, + n_tokens_per_example=config.demo_tokens_per_example, + n_visual_groups=config.demo_visual_groups, + n_action_keys=config.demo_action_keys, + hidden_mult=config.demo_hidden_mult, + ) + + self.spatial_head = ( + nn.Linear(d, config.spatial_vocab) if config.spatial_loss_weight > 0 else None + ) + + self.expert = FlowMatchingExpert( + action_dim=config.max_action_dim, + chunk_size=config.chunk_size, + d_model=d, + n_layers=config.expert_layers, + n_heads=config.expert_heads, + min_period=config.min_period, + max_period=config.max_period, + ) + + self._action_queue: deque = deque() + + def reset(self): + self._action_queue.clear() + + def get_optim_params(self): + return [p for p in self.parameters() if p.requires_grad] + + # ---------------------------------------------------------------- features + + @staticmethod + def _last_obs(x: torch.Tensor, ndim: int) -> torch.Tensor: + """Drop the n_obs_steps dim if present: (B, S, ...) -> (B, ...).""" + return x[:, -1] if x.dim() == ndim + 1 else x + + def _encode_spatial(self, batch: dict) -> torch.Tensor: + """Per-camera spatial tokens for the fast path: (B, num_cams*64, d).""" + cond = [] + for i, key in enumerate(self.image_keys): + tokens = self.semantic.encode_image(self._last_obs(batch[key], 4)) # (B, 64, C) + proj = self.spatial_proj(tokens) + self.camera_emb.weight[i][None, None, :] + if i == 0: + self._cam0_proj = proj # consumed by the spatial-distillation head + cond.append(proj) + return torch.cat(cond, dim=1) + + def _semantic_latent(self, batch: dict, image=None) -> torch.Tensor: + """Slow-path latent: (B, K, d), cacheable, conditions the fast path. + + vlm_native=True makes this latent RICH — the LM sequence also carries the + robot description and the demo (image, action) exemplars, so the readout + tokens can attend over them. The fast path is unchanged: it still gets its + own spatial tokens + proprio + expert, just a better-informed latent. + """ + img = image if image is not None else self._last_obs(batch[self.image_keys[0]], 4) + if self.config.vlm_native: + state = None + if self.state_dim: + state = pad_vector(self._last_obs(batch[OBS_STATE], 2), self.config.max_state_dim) + return self.semantic.forward_unified( + image=img, + lang_tokens=batch[OBS_LANGUAGE_TOKENS], + lang_mask=batch[OBS_LANGUAGE_ATTENTION_MASK], + state=state, + # demo_only: no robot description — demos are the sole identity cue + robot_ids=None if self.config.demo_only else batch.get("morph_text_ids"), + robot_mask=None if self.config.demo_only else batch.get("morph_text_mask"), + demo_images=batch.get("support_images"), + demo_actions=batch.get("support_actions"), + n_demo_action_keys=self.config.demo_action_keys_lm, + ) + morph = batch.get("morphology") if self.config.morph_to_slow else None + return self.semantic( + img, + batch[OBS_LANGUAGE_TOKENS], + batch[OBS_LANGUAGE_ATTENTION_MASK], + morphology=morph, + ) + + def _conditioning(self, batch: dict) -> torch.Tensor: + b = batch[self.image_keys[0]].shape[0] + device = batch[self.image_keys[0]].device + + parts = [] + if self.config.use_semantic_latent: + latent = batch.get("semantic_latent") + if latent is None: + if "latent_image" in batch: # staleness augmentation: latent from an older frame + latent = self._semantic_latent(batch, image=batch["latent_image"]) + else: + latent = self._semantic_latent(batch) + parts.append(latent) + if self.config.use_spatial_tokens: + parts.append(self._encode_spatial(batch)) + if self.state_dim: + state = pad_vector(self._last_obs(batch[OBS_STATE], 2), self.config.max_state_dim) + parts.append(self.state_proj(state)[:, None, :]) + if self.config.demo_only: + pass # no identity token at all — demos in the LM latent are the only cue + elif self.config.conditioning == "morph": + desc = batch.get("morphology") + if desc is None: + from .modules.embodiment import MORPH_DIM + + desc = torch.zeros(b, MORPH_DIM, device=device) + parts.append(self.morph_encoder(desc.float())) + elif self.config.conditioning == "morph_qwen": + ids = batch.get("morph_text_ids") + mask = batch.get("morph_text_mask") + if ids is None: + ids = torch.zeros(b, 1, dtype=torch.long, device=device) + mask = torch.ones(b, 1, dtype=torch.bool, device=device) + parts.append(self.semantic.encode_robot_text(ids, mask)) + else: + emb_id = batch.get("embodiment_id") + if emb_id is None: + emb_id = torch.zeros(b, dtype=torch.long, device=device) + parts.append(self.embodiment_emb(emb_id)) + + if self.config.use_demo_conditioning: + sup_imgs = batch.get("support_images") # (B, K, 3, H, W) + sup_acts = batch.get("support_actions") # (B, K, T, A) + if sup_imgs is not None and sup_acts is not None: + bb, k = sup_imgs.shape[:2] + flat_imgs = sup_imgs.reshape(bb * k, *sup_imgs.shape[2:]) + spatial = self.semantic.encode_image(flat_imgs) # (B*K, 64, visual_dim) + # keep the spatial structure — the encoder does grouped pooling + spatial = spatial.view(bb, k, spatial.shape[1], spatial.shape[2]) + parts.append(self.demo_encoder(spatial, sup_acts)) + return torch.cat(parts, dim=1) + + # ---------------------------------------------------------------- training + + def forward(self, batch: dict) -> tuple[torch.Tensor, dict]: + cond = self._conditioning(batch) + + actions = pad_vector(batch[ACTION], self.config.max_action_dim) # (B, T, A) + b, t, a = actions.shape + device = actions.device + + beta = torch.distributions.Beta(self.config.flow_beta_alpha, self.config.flow_beta_beta) + time = beta.sample((b,)).to(device) * 0.999 + 0.001 + noise = torch.randn_like(actions) + x_t = time[:, None, None] * noise + (1 - time[:, None, None]) * actions + target = noise - actions + + pred = self.expert(x_t, time, cond) + + loss_mask = torch.ones(b, t, a, device=device) + if "action_dim_mask" in batch: # per-sample real dims (mixture training) + loss_mask = loss_mask * batch["action_dim_mask"][:, None, :].float() + else: + loss_mask[:, :, self.action_dim :] = 0 # padded action dims + if "action_is_pad" in batch: # episode-end padding from delta_timestamps + loss_mask = loss_mask * (~batch["action_is_pad"])[:, :, None].float() + + loss = (F.mse_loss(pred, target, reduction="none") * loss_mask).sum() / loss_mask.sum() + info = {"flow_loss": loss.item()} + + if ( + self.spatial_head is not None + and "spatial_cell" in batch + and self.config.use_spatial_tokens + ): + cells = batch["spatial_cell"] + labeled = cells >= 0 + if labeled.any(): + logits = self.spatial_head(self._cam0_proj.mean(dim=1)) # (B, vocab) + ce = F.cross_entropy(logits[labeled], cells[labeled]) + loss = loss + self.config.spatial_loss_weight * ce + info["spatial_ce"] = ce.item() + + info["loss"] = loss.item() + return loss, info + + # --------------------------------------------------------------- inference + + @torch.no_grad() + def predict_action_chunk(self, batch: dict) -> torch.Tensor: + cond = self._conditioning(batch) + b = cond.shape[0] + device = cond.device + + x = torch.randn(b, self.config.chunk_size, self.config.max_action_dim, device=device) + n = self.config.num_denoise_steps + dt = 1.0 / n + for i in range(n): + t = torch.full((b,), 1.0 - i * dt, device=device) + v = self.expert(x, t, cond) + x = x - dt * v + return x[:, :, : self.action_dim] + + @torch.no_grad() + def select_action(self, batch: dict) -> torch.Tensor: + if len(self._action_queue) == 0: + chunk = self.predict_action_chunk(batch) # (B, T, A) + self._action_queue.extend(chunk.transpose(0, 1)[: self.config.n_action_steps]) + return self._action_queue.popleft() diff --git a/tinyvla2/src/tinyvla/modules/__pycache__/embodiment.cpython-312.pyc b/tinyvla2/src/tinyvla/modules/__pycache__/embodiment.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1bd861a067174f4cd8765f67418f207802ad9f22 Binary files /dev/null and b/tinyvla2/src/tinyvla/modules/__pycache__/embodiment.cpython-312.pyc differ diff --git a/tinyvla2/src/tinyvla/modules/__pycache__/expert.cpython-312.pyc b/tinyvla2/src/tinyvla/modules/__pycache__/expert.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..976ca6fece2f4ee213dd6946bc3e67821fbd9c9a Binary files /dev/null and b/tinyvla2/src/tinyvla/modules/__pycache__/expert.cpython-312.pyc differ diff --git a/tinyvla2/src/tinyvla/modules/__pycache__/semantic.cpython-312.pyc b/tinyvla2/src/tinyvla/modules/__pycache__/semantic.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d88f24198e7f98652fe364c86e4b878f486d83bd Binary files /dev/null and b/tinyvla2/src/tinyvla/modules/__pycache__/semantic.cpython-312.pyc differ diff --git a/tinyvla2/src/tinyvla/modules/__pycache__/vision.cpython-312.pyc b/tinyvla2/src/tinyvla/modules/__pycache__/vision.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b664728cc05548fe870a1b9c8902508d3a93c75c Binary files /dev/null and b/tinyvla2/src/tinyvla/modules/__pycache__/vision.cpython-312.pyc differ diff --git a/tinyvla2/src/tinyvla/modules/embodiment.py b/tinyvla2/src/tinyvla/modules/embodiment.py new file mode 100644 index 0000000000000000000000000000000000000000..5209183608173867018194083ef0915b027ad729 --- /dev/null +++ b/tinyvla2/src/tinyvla/modules/embodiment.py @@ -0,0 +1,136 @@ +"""Cross-embodiment harmonization: pad-to-max vectors + embodiment embedding.""" + +from __future__ import annotations + +import torch +import torch.nn as nn + + +def pad_vector(x: torch.Tensor, target_dim: int) -> torch.Tensor: + """Zero-pad the last dim of (..., D) to (..., target_dim).""" + if x.shape[-1] == target_dim: + return x + if x.shape[-1] > target_dim: + raise ValueError(f"dim {x.shape[-1]} exceeds max {target_dim}") + pad = x.new_zeros(*x.shape[:-1], target_dim - x.shape[-1]) + return torch.cat([x, pad], dim=-1) + + +def dim_validity_mask(orig_dim: int, target_dim: int, ref: torch.Tensor) -> torch.Tensor: + """(target_dim,) bool mask, True for real dims, False for padding.""" + mask = torch.zeros(target_dim, dtype=torch.bool, device=ref.device) + mask[:orig_dim] = True + return mask + + +class EmbodimentEmbedding(nn.Module): + """One learned token per source dataset / robot embodiment (variants A, B).""" + + def __init__(self, num_embodiments: int, dim: int): + super().__init__() + self.emb = nn.Embedding(num_embodiments, dim) + + def forward(self, embodiment_id: torch.Tensor) -> torch.Tensor: + """(B,) int -> (B, 1, dim).""" + return self.emb(embodiment_id)[:, None, :] + + +# Descriptor fields (hand-normalized to O(1)), order fixed — see configs/morphology/*.yaml +MORPH_FIELDS = [ + "arm_dof", # /10 + "reach_m", # meters ×2 + "gripper_width_m", # meters ×10 + "num_cameras", # /3 + "is_mobile", # ±1 + "control_hz", # /30 + "joint_lo_mean", # rad, /pi + "joint_hi_mean", # rad, /pi + "workspace_x", # meters ×2 + "workspace_y", + "workspace_z", + "payload_kg", # /5 + "ee_type_parallel", # ±1 (parallel jaw) + "ee_type_multi", # ±1 (>2 finger) + "base_holonomic", # ±1 + "reserved", +] +MORPH_DIM = len(MORPH_FIELDS) # 16 + + +class MorphologyEncoder(nn.Module): + """Explicit robot descriptor -> a few conditioning tokens (variant C). + + Cheap alternative to a per-robot learned ID: a fixed 16-number descriptor + (arm DOF, reach, gripper width, mobility, control rate, joint limits, ...) + through an MLP -> `n_tokens` morphology tokens. Gives a shot at unseen-robot + transfer (a new robot = write its descriptor, no new learned parameters). + """ + + def __init__(self, dim: int, n_tokens: int = 2, in_dim: int = MORPH_DIM): + super().__init__() + self.n_tokens = n_tokens + self.mlp = nn.Sequential( + nn.Linear(in_dim, dim), + nn.GELU(), + nn.Linear(dim, dim * n_tokens), + ) + + def forward(self, descriptor: torch.Tensor) -> torch.Tensor: + """(B, in_dim) float -> (B, n_tokens, dim).""" + b = descriptor.shape[0] + return self.mlp(descriptor).view(b, self.n_tokens, -1) + + +class DemoConditioner(nn.Module): + """In-context few-shot demo conditioning: K (observation, action-chunk) exemplar + pairs -> conditioning tokens, no gradient update on the new robot. + + v3 (high-capacity): v1/v2 squeezed each exemplar through mean-pool(64 spatial + tokens) + a 3-point action summary into ONE token via a 2-layer MLP — a + bottleneck so tight the model learned to ignore the channel (real ≈ garbage + exemplars). This version widens every stage: + - vision: `n_visual_groups` spatially-grouped pools instead of one global mean + - action: `n_action_keys` evenly-spaced chunk keyframes instead of 3 points + - MLP: 3 layers at `hidden_mult`x width, LayerNorm, -> `n_tokens_per_example` tokens + """ + + def __init__(self, dim: int, visual_dim: int, action_dim: int, + n_tokens_per_example: int = 4, n_visual_groups: int = 4, + n_action_keys: int = 8, hidden_mult: int = 4): + super().__init__() + self.n_tokens_per_example = n_tokens_per_example + self.n_visual_groups = n_visual_groups + self.n_action_keys = n_action_keys + in_dim = visual_dim * n_visual_groups + action_dim * n_action_keys + hidden = dim * hidden_mult + self.mlp = nn.Sequential( + nn.Linear(in_dim, hidden), + nn.GELU(), + nn.LayerNorm(hidden), + nn.Linear(hidden, hidden), + nn.GELU(), + nn.Linear(hidden, dim * n_tokens_per_example), + ) + + def forward(self, obs_tokens: torch.Tensor, action_chunks: torch.Tensor) -> torch.Tensor: + """obs_tokens: (B, K, N, visual_dim) per-exemplar spatial tokens (N=64), or + (B, K, visual_dim) if already pooled (falls back to repeating it). + action_chunks: (B, K, T, action_dim). + Returns (B, K*n_tokens_per_example, dim). + """ + b, k, t, a = action_chunks.shape + if obs_tokens.dim() == 4: # (B, K, N, C) -> spatial group pools + n = obs_tokens.shape[2] + g = self.n_visual_groups + per = max(1, n // g) + groups = [obs_tokens[:, :, i * per:(i + 1) * per].mean(dim=2) for i in range(g)] + vis = torch.cat(groups, dim=-1) # (B, K, C*g) + else: + vis = obs_tokens.repeat(1, 1, self.n_visual_groups) + + keys = torch.linspace(0, t - 1, self.n_action_keys, device=action_chunks.device).long() + act = action_chunks[:, :, keys].reshape(b, k, -1) # (B, K, n_keys*a) + + feat = torch.cat([vis, act], dim=-1) + tok = self.mlp(feat) + return tok.view(b, k * self.n_tokens_per_example, -1) diff --git a/tinyvla2/src/tinyvla/modules/expert.py b/tinyvla2/src/tinyvla/modules/expert.py new file mode 100644 index 0000000000000000000000000000000000000000..b482ed33db711ce4fbf503f324525ed290468708 --- /dev/null +++ b/tinyvla2/src/tinyvla/modules/expert.py @@ -0,0 +1,128 @@ +"""Flow-matching action expert. + +A small transformer over noisy action tokens with interleaved +cross-attention (to the conditioning set) and self-attention blocks, +following the SmolVLA expert pattern. Predicts the flow velocity. +""" + +from __future__ import annotations + +import math + +import torch +import torch.nn as nn + + +def sincos_time_embedding(t: torch.Tensor, dim: int, min_period: float = 4e-3, max_period: float = 4.0) -> torch.Tensor: + """t: (B,) in [0, 1] -> (B, dim) sine-cosine embedding.""" + half = dim // 2 + freqs = torch.exp( + torch.linspace(math.log(min_period), math.log(max_period), half, device=t.device) + ) + angles = 2 * math.pi * t[:, None] / freqs[None, :] + return torch.cat([angles.sin(), angles.cos()], dim=-1) + + +class _MLP(nn.Module): + def __init__(self, dim: int, ratio: int = 4): + super().__init__() + self.net = nn.Sequential( + nn.Linear(dim, dim * ratio), + nn.GELU(), + nn.Linear(dim * ratio, dim), + ) + + def forward(self, x): + return self.net(x) + + +class SelfAttnBlock(nn.Module): + def __init__(self, dim: int, n_heads: int): + super().__init__() + self.norm1 = nn.LayerNorm(dim) + self.attn = nn.MultiheadAttention(dim, n_heads, batch_first=True) + self.norm2 = nn.LayerNorm(dim) + self.mlp = _MLP(dim) + + def forward(self, x, cond, cond_mask): + h = self.norm1(x) + h, _ = self.attn(h, h, h, need_weights=False) + x = x + h + x = x + self.mlp(self.norm2(x)) + return x + + +class CrossAttnBlock(nn.Module): + def __init__(self, dim: int, n_heads: int): + super().__init__() + self.norm1 = nn.LayerNorm(dim) + self.norm_cond = nn.LayerNorm(dim) + self.attn = nn.MultiheadAttention(dim, n_heads, batch_first=True) + self.norm2 = nn.LayerNorm(dim) + self.mlp = _MLP(dim) + + def forward(self, x, cond, cond_mask): + h = self.norm1(x) + c = self.norm_cond(cond) + # key_padding_mask: True = ignore + kpm = ~cond_mask if cond_mask is not None else None + h, _ = self.attn(h, c, c, key_padding_mask=kpm, need_weights=False) + x = x + h + x = x + self.mlp(self.norm2(x)) + return x + + +class FlowMatchingExpert(nn.Module): + def __init__( + self, + action_dim: int = 8, + chunk_size: int = 50, + d_model: int = 512, + n_layers: int = 12, + n_heads: int = 8, + min_period: float = 4e-3, + max_period: float = 4.0, + ): + super().__init__() + self.action_dim = action_dim + self.chunk_size = chunk_size + self.d_model = d_model + self.min_period = min_period + self.max_period = max_period + + self.action_in = nn.Linear(action_dim, d_model) + self.pos_emb = nn.Parameter(torch.zeros(1, chunk_size, d_model)) + self.time_mlp = nn.Sequential( + nn.Linear(d_model, d_model), + nn.GELU(), + nn.Linear(d_model, d_model), + ) + # even blocks cross-attend to the conditioning set, odd blocks self-attend + self.blocks = nn.ModuleList( + [ + (CrossAttnBlock if i % 2 == 0 else SelfAttnBlock)(d_model, n_heads) + for i in range(n_layers) + ] + ) + self.norm_out = nn.LayerNorm(d_model) + self.action_out = nn.Linear(d_model, action_dim) + + nn.init.trunc_normal_(self.pos_emb, std=0.02) + nn.init.zeros_(self.action_out.weight) + nn.init.zeros_(self.action_out.bias) + + def forward( + self, + noisy_actions: torch.Tensor, # (B, T, A) + time: torch.Tensor, # (B,) in [0, 1] + cond: torch.Tensor, # (B, N, d_model) + cond_mask: torch.Tensor | None = None, # (B, N) True = valid + ) -> torch.Tensor: + x = self.action_in(noisy_actions) + self.pos_emb[:, : noisy_actions.shape[1]] + t_emb = self.time_mlp( + sincos_time_embedding(time, self.d_model, self.min_period, self.max_period) + ) + x = x + t_emb[:, None, :] + for block in self.blocks: + x = block(x, cond, cond_mask) + return self.action_out(self.norm_out(x)) diff --git a/tinyvla2/src/tinyvla/modules/semantic.py b/tinyvla2/src/tinyvla/modules/semantic.py new file mode 100644 index 0000000000000000000000000000000000000000..ffef404b2704ae9d7256ca8f953fb17e02e9acf7 --- /dev/null +++ b/tinyvla2/src/tinyvla/modules/semantic.py @@ -0,0 +1,315 @@ +"""Slow semantic path: Qwen3.5-0.8B VLM over [image | instruction | readout slots]. + +Qwen3.5-0.8B has a native vision tower (12-layer ViT, patch 16, spatial merge 2: +a 256x256 image becomes 64 visual tokens), so the slow path consumes the raw +image directly — no hand-rolled projector. The language stack is truncated to +the first `num_layers` blocks (L/2 layer-skip, SmolVLA-style). The semantic +latent is the hidden state of K trailing readout slots (pad tokens after the +instruction; causal attention lets them aggregate the full prefix), projected +to the expert dimension. At deployment this runs at 0.5-2 Hz and the latent is +cached between refreshes. +""" + +from __future__ import annotations + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +class SemanticPath(nn.Module): + def __init__( + self, + model_name: str = "Qwen/Qwen3.5-0.8B", + num_layers: int = 12, + num_readout: int = 8, + out_dim: int = 512, + image_size: int = 256, + freeze_lm: bool = True, + freeze_vision: bool = True, + morph_to_slow: bool = False, + morph_dim: int = 16, + morph_qwen: bool = False, + num_morph_readout: int = 4, + vlm_native: bool = False, + num_action_readout: int = 16, + max_state_dim: int = 16, + max_action_dim: int = 8, + ): + super().__init__() + from transformers import AutoModelForImageTextToText, AutoProcessor + + vlm = AutoModelForImageTextToText.from_pretrained(model_name, dtype=torch.float32) + model = vlm.model # Qwen3_5Model: visual tower + language stack + # L/2 layer-skip: keep the first num_layers language blocks; lm_head unused + model.language_model.layers = model.language_model.layers[:num_layers] + model.language_model.config.num_hidden_layers = num_layers + self.vlm = model + cfg = vlm.config + self.hidden_size = cfg.text_config.hidden_size + self.image_token_id = cfg.image_token_id + self.vision_start_id = cfg.vision_start_token_id + self.vision_end_id = cfg.vision_end_token_id + + proc = AutoProcessor.from_pretrained(model_name) + self.pad_token_id = proc.tokenizer.pad_token_id or proc.tokenizer.eos_token_id + ip = proc.image_processor + self.patch_size = ip.patch_size # 16 + self.merge_size = ip.merge_size # 2 + self.temporal_patch_size = ip.temporal_patch_size # 2 + self.register_buffer( + "pixel_mean", torch.tensor(ip.image_mean).view(1, 3, 1, 1), persistent=False + ) + self.register_buffer( + "pixel_std", torch.tensor(ip.image_std).view(1, 3, 1, 1), persistent=False + ) + self.image_size = image_size + grid = image_size // self.patch_size + self.num_image_tokens = (grid // self.merge_size) ** 2 # 64 @ 256px + + self.out_proj = nn.Linear(self.hidden_size, out_dim) + self.num_readout = num_readout + self.morph_to_slow = morph_to_slow + if morph_to_slow: + # inject robot morphology into the slow LM prefix so the semantic + # latent is morphology-aware ("grasp X" differs per embodiment) + self.morph_to_lm = nn.Linear(morph_dim, self.hidden_size) + + # LLaVA-style unified mode: every modality becomes tokens inside the LM + # sequence, and the action policy reads latent tokens off the LM output. + self.vlm_native = vlm_native + self.num_action_readout = num_action_readout + if vlm_native: + self.state_to_lm = nn.Linear(max_state_dim, self.hidden_size) + self.action_to_lm = nn.Linear(max_action_dim, self.hidden_size) + self.action_readout_proj = nn.Linear(self.hidden_size, out_dim) + self.readout_emb = nn.Parameter(torch.zeros(num_action_readout, self.hidden_size)) + nn.init.normal_(self.readout_emb, std=0.02) + + self.morph_qwen = morph_qwen + self.num_morph_readout = num_morph_readout + if morph_qwen: + # encode the robot TEXT description through the SAME Qwen language + # stack (shared weights, text-only, no image) instead of an MLP over + # numeric fields — lets the (unfrozen) LM's language understanding + # do the morphology encoding instead of a from-scratch projector. + self.morph_out_proj = nn.Linear(self.hidden_size, out_dim) + + if freeze_lm: + for p in self.vlm.language_model.parameters(): + p.requires_grad = False + # vocab embeddings stay frozen either way (254M tied params, negligible gain) + self.vlm.get_input_embeddings().weight.requires_grad = False + if freeze_vision: + for p in self.vlm.visual.parameters(): + p.requires_grad = False + + @property + def visual_hidden_size(self) -> int: + return self.vlm.visual.config.out_hidden_size + + def encode_image(self, images: torch.Tensor) -> torch.Tensor: + """Qwen vision tower spatial tokens: (B, 3, H, W) -> (B, 64, out_hidden). + + Shared encoder for both paths; the fast path cross-attends to these + directly, bypassing the language stack. + """ + b = images.shape[0] + pixel_values, grid_thw = self._patchify(images) + out = self.vlm.get_image_features(pixel_values.to(self.pixel_mean.dtype), grid_thw) + embeds = torch.cat(out.pooler_output, dim=0) # per-image tuple -> (B*64, C) + return embeds.reshape(b, self.num_image_tokens, -1) + + def _patchify(self, images: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """(B, 3, H, W) in [0,1] -> Qwen-VL flattened patches + grid_thw (B, 3). + + Mirrors the Qwen VL image processor: resize to image_size, normalize, + repeat temporally, unfold into merge-ordered patches. + """ + b = images.shape[0] + if images.shape[-2:] != (self.image_size, self.image_size): + images = F.interpolate( + images, size=(self.image_size, self.image_size), mode="bilinear", align_corners=False + ) + images = (images - self.pixel_mean) / self.pixel_std + + p, m, tp = self.patch_size, self.merge_size, self.temporal_patch_size + g = self.image_size // p + # (B, tp, C, H, W) + x = images.unsqueeze(1).expand(-1, tp, -1, -1, -1) + # split H,W into (g/m, m, p) blocks, order patches so that the m*m + # merge-group is contiguous: (gh/m, gw/m, m, m) patch order + x = x.reshape(b, tp, 3, g // m, m, p, g // m, m, p) + x = x.permute(0, 3, 6, 4, 7, 2, 1, 5, 8) # b, gh/m, gw/m, m, m, C, tp, p, p + x = x.reshape(b * g * g, 3 * tp * p * p) + grid_thw = torch.tensor([[1, g, g]], device=images.device).expand(b, -1) + return x, grid_thw + + def forward( + self, + images: torch.Tensor, # (B, 3, H, W) in [0, 1] + lang_tokens: torch.Tensor, # (B, L) + lang_mask: torch.Tensor, # (B, L) bool, True = valid + morphology: torch.Tensor | None = None, # (B, morph_dim) for morph_to_slow + ) -> torch.Tensor: + """Returns semantic latent (B, num_readout, out_dim).""" + b = images.shape[0] + device = images.device + + pixel_values, grid_thw = self._patchify(images) + n_img = self.num_image_tokens + + # *n instruction *K + ids = torch.cat( + [ + torch.full((b, 1), self.vision_start_id, dtype=torch.long, device=device), + torch.full((b, n_img), self.image_token_id, dtype=torch.long, device=device), + torch.full((b, 1), self.vision_end_id, dtype=torch.long, device=device), + lang_tokens, + torch.full((b, self.num_readout), self.pad_token_id, dtype=torch.long, device=device), + ], + dim=1, + ) + attn = torch.cat( + [ + torch.ones(b, n_img + 2, dtype=torch.bool, device=device), + lang_mask.bool(), + torch.ones(b, self.num_readout, dtype=torch.bool, device=device), + ], + dim=1, + ) + mm_type = (ids == self.image_token_id).int() # text 0, image 1 + + if self.morph_to_slow and morphology is not None: + # embed tokens, scatter image features, add morphology to readout slots, + # then run the LM on inputs_embeds (keeps M-RoPE via mm_token_type_ids) + inputs_embeds = self.vlm.get_input_embeddings()(ids) + img = self.vlm.get_image_features(pixel_values.to(self.pixel_mean.dtype), grid_thw) + img = torch.cat(img.pooler_output, dim=0).to(inputs_embeds.dtype) # (B*n_img, hidden) + imask = (ids == self.image_token_id).unsqueeze(-1).expand_as(inputs_embeds) + inputs_embeds = inputs_embeds.masked_scatter(imask, img) + morph_tok = self.morph_to_lm(morphology.to(inputs_embeds.dtype)) # (B, hidden) + inputs_embeds[:, -self.num_readout:] = inputs_embeds[:, -self.num_readout:] + morph_tok[:, None, :] + out = self.vlm( + inputs_embeds=inputs_embeds, image_grid_thw=grid_thw, attention_mask=attn, + mm_token_type_ids=mm_type, use_cache=False, + ) + else: + out = self.vlm( + input_ids=ids, pixel_values=pixel_values.to(self.pixel_mean.dtype), + image_grid_thw=grid_thw, attention_mask=attn, mm_token_type_ids=mm_type, + use_cache=False, + ) + readout_h = out.last_hidden_state[:, -self.num_readout :] + return self.out_proj(readout_h) + + def encode_robot_text(self, desc_ids: torch.Tensor, desc_mask: torch.Tensor) -> torch.Tensor: + """Encode a robot-description text through the SAME (shared-weight) Qwen + language stack, text-only (no image, no M-RoPE needed). Trailing pad-token + readout slots aggregate the description via causal attention; their hidden + states are projected to `num_morph_readout` conditioning tokens. + + desc_ids/desc_mask: (B, L) tokenized "Robot: X, a N-DOF arm on ..." text. + Returns (B, num_morph_readout, out_dim). + """ + b = desc_ids.shape[0] + device = desc_ids.device + ids = torch.cat( + [desc_ids, torch.full((b, self.num_morph_readout), self.pad_token_id, + dtype=torch.long, device=device)], + dim=1, + ) + attn = torch.cat( + [desc_mask.bool(), torch.ones(b, self.num_morph_readout, dtype=torch.bool, device=device)], + dim=1, + ) + out = self.vlm.language_model( + input_ids=ids, attention_mask=attn, use_cache=False, + ) + readout_h = out.last_hidden_state[:, -self.num_morph_readout:] + return self.morph_out_proj(readout_h) + + def forward_unified( + self, + image: torch.Tensor, # (B, 3, H, W) current observation + lang_tokens: torch.Tensor, # (B, L) task instruction + lang_mask: torch.Tensor, # (B, L) + state: torch.Tensor | None = None, # (B, max_state_dim) proprio + robot_ids: torch.Tensor | None = None, # (B, Lr) robot description text + robot_mask: torch.Tensor | None = None, + demo_images: torch.Tensor | None = None, # (B, K, 3, H, W) + demo_actions: torch.Tensor | None = None, # (B, K, T, A) + n_demo_action_keys: int = 4, + ) -> torch.Tensor: + """LLaVA-style unified pass: ALL modalities become tokens in one LM + sequence, and the policy reads `num_action_readout` latent tokens off the + LM output. Demo exemplars live INSIDE the attention stream, so the readout + tokens can attend to them selectively (the thing that "conditioning-set" + placement could not do). + + Sequence layout: + [ main-image <-vis>] [instruction] [robot description] + [ demo_k image <-vis> demo_k action-keyframe tokens ]*K + [state token] [action readout slots] + Returns (B, num_action_readout, out_dim). + """ + b = image.shape[0] + device = image.device + emb = self.vlm.get_input_embeddings() + n_img = self.num_image_tokens + + # ---- encode all images (main + demos) in one vision pass ------------- + imgs = [image] + k = 0 + if demo_images is not None: + k = demo_images.shape[1] + imgs.append(demo_images.reshape(b * k, *demo_images.shape[2:])) + all_imgs = torch.cat(imgs, dim=0) + pixel_values, grid_thw = self._patchify(all_imgs) + feats = self.vlm.get_image_features(pixel_values.to(self.pixel_mean.dtype), grid_thw) + feats = torch.cat(feats.pooler_output, dim=0) # (B*(1+K)*n_img, hidden) + feats = feats.view(b * (1 + k), n_img, -1) + main_feat = feats[:b] # (B, n_img, hidden) + demo_feat = feats[b:].view(b, k, n_img, -1) if k else None + + seq, mm, msk = [], [], [] + + def push(x, is_img=False, mask=None): + seq.append(x) + n = x.shape[1] + mm.append(torch.full((b, n), 1 if is_img else 0, dtype=torch.int, device=device)) + msk.append(mask if mask is not None else torch.ones(b, n, dtype=torch.bool, device=device)) + + vs = emb(torch.full((b, 1), self.vision_start_id, dtype=torch.long, device=device)) + ve = emb(torch.full((b, 1), self.vision_end_id, dtype=torch.long, device=device)) + + push(vs); push(main_feat, is_img=True); push(ve) + push(emb(lang_tokens), mask=lang_mask.bool()) + if robot_ids is not None: + push(emb(robot_ids), mask=robot_mask.bool()) + if demo_feat is not None: + keys = torch.linspace(0, demo_actions.shape[2] - 1, n_demo_action_keys, + device=device).long() + for i in range(k): + push(vs); push(demo_feat[:, i], is_img=True); push(ve) + act_tok = self.action_to_lm(demo_actions[:, i, keys].to(main_feat.dtype)) + push(act_tok) # (B, n_keys, hidden) + if state is not None: + push(self.state_to_lm(state.to(main_feat.dtype))[:, None, :]) + push(self.readout_emb[None].expand(b, -1, -1).to(main_feat.dtype)) + + inputs_embeds = torch.cat(seq, dim=1) + mm_type = torch.cat(mm, dim=1) + attn = torch.cat(msk, dim=1) + # image_grid_thw must cover only the images actually present in-sequence + grid_for_seq = grid_thw[: b * (1 + k)] + + out = self.vlm( + inputs_embeds=inputs_embeds, + image_grid_thw=grid_for_seq, + attention_mask=attn, + mm_token_type_ids=mm_type, + use_cache=False, + ) + readout = out.last_hidden_state[:, -self.num_action_readout:] + return self.action_readout_proj(readout) diff --git a/tinyvla2/src/tinyvla/processor_tinyvla.py b/tinyvla2/src/tinyvla/processor_tinyvla.py new file mode 100644 index 0000000000000000000000000000000000000000..dadb4fe3c6c04b1373c8b3f190e7a8013a6e8fba --- /dev/null +++ b/tinyvla2/src/tinyvla/processor_tinyvla.py @@ -0,0 +1,45 @@ +"""Pre/post-processor pipelines for TinyVLA (normalization + language tokenization).""" + +from __future__ import annotations + +from typing import Any + +import torch + +from lerobot.processor import ( + PolicyAction, + PolicyProcessorPipeline, + TokenizerProcessorStep, + make_default_policy_processor_steps, + make_policy_processor_pipelines, +) + +from .configuration_tinyvla import TinyVLAConfig + + +def make_tinyvla_pre_post_processors( + config: TinyVLAConfig, + dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None, +) -> tuple[ + PolicyProcessorPipeline[dict[str, Any], dict[str, Any]], + PolicyProcessorPipeline[PolicyAction, PolicyAction], +]: + steps = make_default_policy_processor_steps(config, dataset_stats) + + input_steps = [ + steps.rename_observations, + steps.add_batch_dim, + TokenizerProcessorStep( + tokenizer_name=config.lm_model_name, + padding=config.pad_language_to, + padding_side="right", + max_length=config.tokenizer_max_length, + ), + steps.to_device, + steps.normalize, + ] + output_steps = [ + steps.unnormalize, + steps.to_cpu, + ] + return make_policy_processor_pipelines(input_steps=input_steps, output_steps=output_steps)