Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,9 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+videos/ep00_task0_succ.mp4 filter=lfs diff=lfs merge=lfs -text
+videos/ep01_task1_succ.mp4 filter=lfs diff=lfs merge=lfs -text
+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

README.md

@@ -0,0 +1,148 @@
+---
+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-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 |
+|---|---|---|
+| π0 | 3.3B | 98% |
+| SmolVLA-2.25B | 2.25B | 93% |
+| XS-VLA | 0.26B | ~93 (avg 90.3) |
+| **TinyVLA (this)** | **0.64B (140M online)** | **86% dual-rate / 79% full** |
+| OpenVLA | 7B | 84.7% |
+
+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 <ckpt> \
+#   --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.

checkpoints/stage2_final/config.json

@@ -0,0 +1,88 @@
+{
+    "type": "tinyvla",
+    "n_obs_steps": 1,
+    "input_features": {
+        "observation.images.cam0": {
+            "type": "VISUAL",
+            "shape": [
+                3,
+                256,
+                256
+            ]
+        },
+        "observation.images.cam1": {
+            "type": "VISUAL",
+            "shape": [
+                3,
+                256,
+                256
+            ]
+        },
+        "observation.state": {
+            "type": "STATE",
+            "shape": [
+                16
+            ]
+        }
+    },
+    "output_features": {
+        "action": {
+            "type": "ACTION",
+            "shape": [
+                8
+            ]
+        }
+    },
+    "device": "cuda",
+    "use_amp": false,
+    "use_peft": false,
+    "push_to_hub": true,
+    "repo_id": null,
+    "private": null,
+    "tags": null,
+    "license": null,
+    "pretrained_path": null,
+    "pretrained_revision": null,
+    "chunk_size": 50,
+    "n_action_steps": 50,
+    "normalization_mapping": {
+        "VISUAL": "IDENTITY",
+        "STATE": "MEAN_STD",
+        "ACTION": "MEAN_STD"
+    },
+    "max_state_dim": 16,
+    "max_action_dim": 8,
+    "num_embodiments": 8,
+    "max_cameras": 3,
+    "image_size": 256,
+    "freeze_vision_encoder": true,
+    "lm_model_name": "Qwen/Qwen3.5-0.8B",
+    "lm_num_layers": 12,
+    "num_readout_tokens": 8,
+    "freeze_lm": true,
+    "tokenizer_max_length": 48,
+    "pad_language_to": "longest",
+    "expert_dim": 512,
+    "expert_layers": 12,
+    "expert_heads": 8,
+    "use_semantic_latent": true,
+    "use_spatial_tokens": true,
+    "num_denoise_steps": 10,
+    "flow_beta_alpha": 1.5,
+    "flow_beta_beta": 1.0,
+    "min_period": 0.004,
+    "max_period": 4.0,
+    "staleness_prob": 0.0,
+    "staleness_max_s": 2.0,
+    "optimizer_lr": 0.0001,
+    "optimizer_betas": [
+        0.9,
+        0.95
+    ],
+    "optimizer_eps": 1e-08,
+    "optimizer_weight_decay": 1e-10,
+    "optimizer_grad_clip_norm": 10.0,
+    "scheduler_warmup_steps": 1000,
+    "scheduler_decay_steps": 40000,
+    "scheduler_decay_lr": 2.5e-06
+}

checkpoints/stage2_final/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1df919883d5f012a4b7b218f9d38ac165202dfaf6b3cd9515ab083458f2b0494
+size 2573613872

checkpoints/stage3_final/config.json

@@ -0,0 +1,90 @@
+{
+    "type": "tinyvla",
+    "n_obs_steps": 1,
+    "input_features": {
+        "observation.images.cam0": {
+            "type": "VISUAL",
+            "shape": [
+                3,
+                256,
+                256
+            ]
+        },
+        "observation.images.cam1": {
+            "type": "VISUAL",
+            "shape": [
+                3,
+                256,
+                256
+            ]
+        },
+        "observation.state": {
+            "type": "STATE",
+            "shape": [
+                16
+            ]
+        }
+    },
+    "output_features": {
+        "action": {
+            "type": "ACTION",
+            "shape": [
+                8
+            ]
+        }
+    },
+    "device": "cuda",
+    "use_amp": false,
+    "use_peft": false,
+    "push_to_hub": true,
+    "repo_id": null,
+    "private": null,
+    "tags": null,
+    "license": null,
+    "pretrained_path": null,
+    "pretrained_revision": null,
+    "chunk_size": 50,
+    "n_action_steps": 50,
+    "normalization_mapping": {
+        "VISUAL": "IDENTITY",
+        "STATE": "MEAN_STD",
+        "ACTION": "MEAN_STD"
+    },
+    "max_state_dim": 16,
+    "max_action_dim": 8,
+    "num_embodiments": 8,
+    "max_cameras": 3,
+    "image_size": 256,
+    "freeze_vision_encoder": true,
+    "lm_model_name": "Qwen/Qwen3.5-0.8B",
+    "lm_num_layers": 12,
+    "num_readout_tokens": 8,
+    "freeze_lm": true,
+    "tokenizer_max_length": 48,
+    "pad_language_to": "longest",
+    "expert_dim": 512,
+    "expert_layers": 12,
+    "expert_heads": 8,
+    "use_semantic_latent": true,
+    "use_spatial_tokens": true,
+    "spatial_vocab": 1024,
+    "spatial_loss_weight": 0.15,
+    "num_denoise_steps": 10,
+    "flow_beta_alpha": 1.5,
+    "flow_beta_beta": 1.0,
+    "min_period": 0.004,
+    "max_period": 4.0,
+    "staleness_prob": 0.0,
+    "staleness_max_s": 2.0,
+    "optimizer_lr": 0.0001,
+    "optimizer_betas": [
+        0.9,
+        0.95
+    ],
+    "optimizer_eps": 1e-08,
+    "optimizer_weight_decay": 1e-10,
+    "optimizer_grad_clip_norm": 10.0,
+    "scheduler_warmup_steps": 1000,
+    "scheduler_decay_steps": 40000,
+    "scheduler_decay_lr": 2.5e-06
+}

checkpoints/stage3_final/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:631b94a8ea3d279477ce0f79efc852d9c21bb41e483197aa0f77b19ba1f3b8e5
+size 2575715312

config.json

@@ -0,0 +1,90 @@
+{
+    "type": "tinyvla",
+    "n_obs_steps": 1,
+    "input_features": {
+        "observation.images.cam0": {
+            "type": "VISUAL",
+            "shape": [
+                3,
+                256,
+                256
+            ]
+        },
+        "observation.images.cam1": {
+            "type": "VISUAL",
+            "shape": [
+                3,
+                256,
+                256
+            ]
+        },
+        "observation.state": {
+            "type": "STATE",
+            "shape": [
+                16
+            ]
+        }
+    },
+    "output_features": {
+        "action": {
+            "type": "ACTION",
+            "shape": [
+                8
+            ]
+        }
+    },
+    "device": "cuda",
+    "use_amp": false,
+    "use_peft": false,
+    "push_to_hub": true,
+    "repo_id": null,
+    "private": null,
+    "tags": null,
+    "license": null,
+    "pretrained_path": null,
+    "pretrained_revision": null,
+    "chunk_size": 50,
+    "n_action_steps": 50,
+    "normalization_mapping": {
+        "VISUAL": "IDENTITY",
+        "STATE": "MEAN_STD",
+        "ACTION": "MEAN_STD"
+    },
+    "max_state_dim": 16,
+    "max_action_dim": 8,
+    "num_embodiments": 8,
+    "max_cameras": 3,
+    "image_size": 256,
+    "freeze_vision_encoder": false,
+    "lm_model_name": "Qwen/Qwen3.5-0.8B",
+    "lm_num_layers": 12,
+    "num_readout_tokens": 8,
+    "freeze_lm": false,
+    "tokenizer_max_length": 48,
+    "pad_language_to": "longest",
+    "expert_dim": 512,
+    "expert_layers": 12,
+    "expert_heads": 8,
+    "use_semantic_latent": true,
+    "use_spatial_tokens": true,
+    "spatial_vocab": 1024,
+    "spatial_loss_weight": 0.0,
+    "num_denoise_steps": 10,
+    "flow_beta_alpha": 1.5,
+    "flow_beta_beta": 1.0,
+    "min_period": 0.004,
+    "max_period": 4.0,
+    "staleness_prob": 0.0,
+    "staleness_max_s": 2.0,
+    "optimizer_lr": 0.0001,
+    "optimizer_betas": [
+        0.9,
+        0.95
+    ],
+    "optimizer_eps": 1e-08,
+    "optimizer_weight_decay": 1e-10,
+    "optimizer_grad_clip_norm": 10.0,
+    "scheduler_warmup_steps": 1000,
+    "scheduler_decay_steps": 40000,
+    "scheduler_decay_lr": 2.5e-06
+}

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

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

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

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

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

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

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

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

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:00712c0573cffd4af565a82d2b1e94f74be7d4de5ca9518228522cb311b47dc7
+size 2573613872

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

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

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

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

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/<dataset>.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()

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 <root>_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)

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

scripts/train.py

@@ -0,0 +1,255 @@
+#!/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")
+
+    def wrap(ds, emb_id):
+        store = None
+        if labels_dir:
+            from tinyvla.data.spatial_labels import SpatialLabelStore
+
+            store = SpatialLabelStore(labels_dir, ds.repo_id.split("/")[-1])
+            if len(store) == 0:
+                store = None
+        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,
+            # staleness_prob switched on at cfg["staleness_start_step"]
+            spatial_labels=store,
+        )
+
+    next_emb = 0
+    for spec in cfg["datasets"]:
+        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), 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), spec["weight"], spec["repo_id"], emb)
+        next_emb += 1
+
+    mixture = WeightedMixtureDataset(datasets, weights)
+    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()

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 \
+  "$@"

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"

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

src/tinyvla.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

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

src/tinyvla.egg-info/top_level.txt

@@ -0,0 +1 @@
+tinyvla

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"]

src/tinyvla/configuration_tinyvla.py

@@ -0,0 +1,122 @@
+"""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 = 8
+    max_cameras: int = 3
+
+    # 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

src/tinyvla/data/mixture.py

@@ -0,0 +1,178 @@
+"""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
+    ):
+        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
+
+        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 _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)
+
+        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)
+        out["task"] = item.get("task") or ""
+
+        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]])
+        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")

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

src/tinyvla/modeling_tinyvla.py

@@ -0,0 +1,213 @@
+"""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,
+        )
+
+        # 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)
+
+        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) -> torch.Tensor:
+        """Slow-path latent from the first (main) camera + instruction: (B, K, d)."""
+        return self.semantic(
+            self._last_obs(batch[self.image_keys[0]], 4),
+            batch[OBS_LANGUAGE_TOKENS],
+            batch[OBS_LANGUAGE_ATTENTION_MASK],
+        )
+
+    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(
+                        batch["latent_image"],
+                        batch[OBS_LANGUAGE_TOKENS],
+                        batch[OBS_LANGUAGE_ATTENTION_MASK],
+                    )
+                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, :])
+        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))
+        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()

src/tinyvla/modules/embodiment.py

@@ -0,0 +1,35 @@
+"""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."""
+
+    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, :]

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))

src/tinyvla/modules/semantic.py

@@ -0,0 +1,157 @@
+"""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,
+    ):
+        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
+
+        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
+    ) -> 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
+
+        # <vision_start> <image>*n <vision_end> instruction <pad>*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,
+        )
+
+        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=(ids == self.image_token_id).int(),  # text 0, image 1
+            use_cache=False,
+        )
+        readout_h = out.last_hidden_state[:, -self.num_readout :]
+        return self.out_proj(readout_h)