diff --git a/.gitattributes b/.gitattributes index a6344aac8c09253b3b630fb776ae94478aa0275b..92e08d7ac8aa20258069506f45b3340b42eaeae5 100644 --- a/.gitattributes +++ b/.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 diff --git a/README.md b/README.md new file mode 100644 index 0000000000000000000000000000000000000000..ccae059ee0703243d022a78471cfb5ee7efeea4e --- /dev/null +++ b/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 \ +# --suite libero_spatial --episodes 100 --n-action-steps 10 --refresh-s 1.0 +``` + +Trained 2026-08-14 → 2026-08-19 on 1× RTX A6000 (48GB), lerobot 0.6.1, torch 2.10, transformers 5.5.4. diff --git a/checkpoints/stage2_final/config.json b/checkpoints/stage2_final/config.json new file mode 100644 index 0000000000000000000000000000000000000000..6cb7b35bfb472230ff41d03d7e338add3c90e6eb --- /dev/null +++ b/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 +} \ No newline at end of file diff --git a/checkpoints/stage2_final/model.safetensors b/checkpoints/stage2_final/model.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..e48f662e9d86d2124707a00c6fa4b54f3be3a7a7 --- /dev/null +++ b/checkpoints/stage2_final/model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:1df919883d5f012a4b7b218f9d38ac165202dfaf6b3cd9515ab083458f2b0494 +size 2573613872 diff --git a/checkpoints/stage3_final/config.json b/checkpoints/stage3_final/config.json new file mode 100644 index 0000000000000000000000000000000000000000..deb6aecd78ed7749f70010a2c4e66d78aef09819 --- /dev/null +++ b/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 +} \ No newline at end of file diff --git a/checkpoints/stage3_final/model.safetensors b/checkpoints/stage3_final/model.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..047940131f18ee0f2b86bee93196eb4fe4c8a8ab --- /dev/null +++ b/checkpoints/stage3_final/model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:631b94a8ea3d279477ce0f79efc852d9c21bb41e483197aa0f77b19ba1f3b8e5 +size 2575715312 diff --git a/config.json b/config.json new file mode 100644 index 0000000000000000000000000000000000000000..f4017497d255e6d1e72c06d30a7f1baeb9221e90 --- /dev/null +++ b/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 +} \ No newline at end of file diff --git a/configs/ablation_nolatent.yaml b/configs/ablation_nolatent.yaml new file mode 100644 index 0000000000000000000000000000000000000000..6516fdaa74765a5c3927946c292ad25e7c025b04 --- /dev/null +++ b/configs/ablation_nolatent.yaml @@ -0,0 +1,25 @@ +# Stage-2 gate baseline: same mixture, 5k steps, nolatent ablation +output_dir: /home/alexw/tinyvla/outputs/ablation_nolatent +wandb: null +datasets: + - root_glob: /home/alexw/tinyvla_data/so101_v3/* + weight: 0.65 + - repo_id: nvidia/BridgeData2_LeRobot_v3 + root: /home/alexw/tinyvla_data/bridge_v3_10k + revision: main + episodes: 10000 + weight: 0.35 +policy: + chunk_size: 50 + n_action_steps: 50 + image_size: 256 + use_semantic_latent: false +batch_size: 64 +grad_accum: 2 +num_workers: 12 +lr: 1.0e-4 +warmup_steps: 500 +steps: 5000 +log_freq: 50 +save_freq: 2500 +mixed_precision: bf16 diff --git a/configs/ablation_nospatial.yaml b/configs/ablation_nospatial.yaml new file mode 100644 index 0000000000000000000000000000000000000000..e2bcc0ea6123d64cac7ec1c555d953e9b54b0fc1 --- /dev/null +++ b/configs/ablation_nospatial.yaml @@ -0,0 +1,25 @@ +# Stage-2 gate baseline: same mixture, 5k steps, nospatial ablation +output_dir: /home/alexw/tinyvla/outputs/ablation_nospatial +wandb: null +datasets: + - root_glob: /home/alexw/tinyvla_data/so101_v3/* + weight: 0.65 + - repo_id: nvidia/BridgeData2_LeRobot_v3 + root: /home/alexw/tinyvla_data/bridge_v3_10k + revision: main + episodes: 10000 + weight: 0.35 +policy: + chunk_size: 50 + n_action_steps: 50 + image_size: 256 + use_spatial_tokens: false +batch_size: 64 +grad_accum: 2 +num_workers: 12 +lr: 1.0e-4 +warmup_steps: 500 +steps: 5000 +log_freq: 50 +save_freq: 2500 +mixed_precision: bf16 diff --git a/configs/libero_ft.yaml b/configs/libero_ft.yaml new file mode 100644 index 0000000000000000000000000000000000000000..da4208f1273a4ad9c3388ece78ea7438472fef08 --- /dev/null +++ b/configs/libero_ft.yaml @@ -0,0 +1,31 @@ +# LIBERO fine-tune from the Stage-2 pretrained checkpoint (XS-VLA/SmolVLA protocol). +# LIBERO = new embodiment (id 2): Franka 7-DoF EE @10fps, two 256^2 cameras. + +output_dir: /home/alexw/tinyvla/outputs/libero_ft +wandb: null + +datasets: + - repo_id: HuggingFaceVLA/libero + weight: 1.0 + embodiment_id: 2 + +policy: + chunk_size: 50 + n_action_steps: 50 + image_size: 256 + freeze_lm: true + freeze_vision_encoder: true + +resume_from: /home/alexw/tinyvla/outputs/stage2_mixture/final +resume_step: 0 + +batch_size: 64 +grad_accum: 1 +num_workers: 12 +lr: 5.0e-5 +warmup_steps: 500 +steps: 15000 +grad_clip: 10.0 +log_freq: 50 +save_freq: 2500 +mixed_precision: bf16 diff --git a/configs/libero_ft2.yaml b/configs/libero_ft2.yaml new file mode 100644 index 0000000000000000000000000000000000000000..6eab4bdded3eb6fd9e90a984561afa326f16d26b --- /dev/null +++ b/configs/libero_ft2.yaml @@ -0,0 +1,33 @@ +# LIBERO fine-tune v2: unfreeze the vision tower (sim renders are OOD for the +# real-world-pretrained encoder), longer schedule. + +output_dir: /home/alexw/tinyvla/outputs/libero_ft2 +wandb: null + +datasets: + - repo_id: HuggingFaceVLA/libero + weight: 1.0 + embodiment_id: 2 + +policy: + chunk_size: 50 + n_action_steps: 50 + image_size: 256 + freeze_lm: true + freeze_vision_encoder: false + +backbone_lr_mult: 0.1 + +resume_from: /home/alexw/tinyvla/outputs/stage2_mixture/final +resume_step: 0 + +batch_size: 48 +grad_accum: 1 +num_workers: 12 +lr: 1.0e-4 +warmup_steps: 1000 +steps: 30000 +grad_clip: 10.0 +log_freq: 50 +save_freq: 5000 +mixed_precision: bf16 diff --git a/configs/libero_ft3.yaml b/configs/libero_ft3.yaml new file mode 100644 index 0000000000000000000000000000000000000000..d86dcf85332e0e0cb2b0e3d4434fa15fb978a5f5 --- /dev/null +++ b/configs/libero_ft3.yaml @@ -0,0 +1,33 @@ +# LIBERO fine-tune v2: unfreeze the vision tower (sim renders are OOD for the +# real-world-pretrained encoder), longer schedule. + +output_dir: /home/alexw/tinyvla/outputs/libero_ft3 +wandb: null + +datasets: + - repo_id: HuggingFaceVLA/libero + weight: 1.0 + embodiment_id: 2 + +policy: + chunk_size: 50 + n_action_steps: 50 + image_size: 256 + freeze_lm: true + freeze_vision_encoder: false + +backbone_lr_mult: 0.1 + +resume_from: /home/alexw/tinyvla/outputs/stage3_distill/final +resume_step: 0 + +batch_size: 48 +grad_accum: 1 +num_workers: 12 +lr: 1.0e-4 +warmup_steps: 1000 +steps: 30000 +grad_clip: 10.0 +log_freq: 50 +save_freq: 5000 +mixed_precision: bf16 diff --git a/configs/libero_ft4.yaml b/configs/libero_ft4.yaml new file mode 100644 index 0000000000000000000000000000000000000000..a719f9ef4c5abec2f9d3d699241c54f4064d0fe0 --- /dev/null +++ b/configs/libero_ft4.yaml @@ -0,0 +1,33 @@ +# LIBERO fine-tune v4: unfreeze BOTH vision tower and language stack +# (12 truncated layers; vocab embeddings stay frozen) at 0.1x lr, longer schedule. + +output_dir: /home/alexw/tinyvla/outputs/libero_ft4 +wandb: null + +datasets: + - repo_id: HuggingFaceVLA/libero + weight: 1.0 + embodiment_id: 2 + +policy: + chunk_size: 50 + n_action_steps: 50 + image_size: 256 + freeze_lm: false + freeze_vision_encoder: false + +backbone_lr_mult: 0.1 + +resume_from: /home/alexw/tinyvla/outputs/stage3_distill/final +resume_step: 0 + +batch_size: 48 +grad_accum: 1 +num_workers: 12 +lr: 1.0e-4 +warmup_steps: 1000 +steps: 40000 +grad_clip: 10.0 +log_freq: 50 +save_freq: 5000 +mixed_precision: bf16 diff --git a/configs/stage2_mixture.yaml b/configs/stage2_mixture.yaml new file mode 100644 index 0000000000000000000000000000000000000000..7f0498bee6ba0cb7ad7b15aa825dc903a3885a4b --- /dev/null +++ b/configs/stage2_mixture.yaml @@ -0,0 +1,37 @@ +# Stage 2: small-scale mixture pretraining on one A6000. +# SO-100/101 community slice (converted v3.0) 0.65 + Bridge V2 0.35. +# Fractal dropped for MVP (v2.0, unconvertible in lerobot 0.6.1). + +output_dir: /home/alexw/tinyvla/outputs/stage2_mixture +wandb: null # set a project name to enable + +datasets: + - root_glob: /home/alexw/tinyvla_data/so101_v3/* + weight: 0.65 + - repo_id: nvidia/BridgeData2_LeRobot_v3 + root: /home/alexw/tinyvla_data/bridge_v3_10k + revision: main + episodes: 10000 + weight: 0.35 + +policy: + chunk_size: 50 + n_action_steps: 50 + image_size: 256 + freeze_lm: true + freeze_vision_encoder: true + +batch_size: 64 +grad_accum: 2 +num_workers: 14 +resume_from: /home/alexw/tinyvla/outputs/stage2_mixture/step_2000 +resume_step: 2000 +lr: 1.0e-4 +warmup_steps: 1000 +steps: 40000 +grad_clip: 10.0 +log_freq: 50 +save_freq: 2000 +mixed_precision: bf16 +staleness_start_step: 20000 +staleness_prob: 0.5 diff --git a/configs/stage3_distill.yaml b/configs/stage3_distill.yaml new file mode 100644 index 0000000000000000000000000000000000000000..362662289a994c0c6d1d948985af39c24d8b64f8 --- /dev/null +++ b/configs/stage3_distill.yaml @@ -0,0 +1,40 @@ +# Stage 3: continue SO101+Bridge pretraining with the spatial-distillation +# aux CE loss (teacher labels from Qwen3.5-4B) and staleness aug from the start. + +output_dir: /home/alexw/tinyvla/outputs/stage3_distill +wandb: null + +spatial_labels_dir: /home/alexw/tinyvla_data/spatial_labels + +datasets: + - root_glob: /home/alexw/tinyvla_data/so101_v3/* + weight: 0.65 + - repo_id: nvidia/BridgeData2_LeRobot_v3 + root: /home/alexw/tinyvla_data/bridge_v3_10k + revision: main + episodes: 10000 + weight: 0.35 + +policy: + chunk_size: 50 + n_action_steps: 50 + image_size: 256 + freeze_lm: true + freeze_vision_encoder: true + spatial_loss_weight: 0.15 + +resume_from: /home/alexw/tinyvla/outputs/stage2_mixture/final +resume_step: 0 + +batch_size: 64 +grad_accum: 2 +num_workers: 14 +lr: 5.0e-5 +warmup_steps: 500 +steps: 20000 +grad_clip: 10.0 +log_freq: 50 +save_freq: 2500 +mixed_precision: bf16 +staleness_start_step: 0 +staleness_prob: 0.5 diff --git a/model.safetensors b/model.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..7172b9f5d6c04b7f409e99c8e955fb56ae4341e7 --- /dev/null +++ b/model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:00712c0573cffd4af565a82d2b1e94f74be7d4de5ca9518228522cb311b47dc7 +size 2573613872 diff --git a/scripts/diag_obs_gap.py b/scripts/diag_obs_gap.py new file mode 100644 index 0000000000000000000000000000000000000000..7d6b06096dbd602123412eee029aa3353bc02151 --- /dev/null +++ b/scripts/diag_obs_gap.py @@ -0,0 +1,131 @@ +#!/usr/bin/env python +"""Diagnose the env-obs adapter: compare policy predictions from env-rendered +observations vs dataset observations at the SAME init state. + +If pred(dataset obs) is close to GT but pred(env obs) differs, the observation +adapter (image orientation/cameras/state) is the remaining gap. +""" + +from __future__ import annotations + +import argparse + +import numpy as np +import torch +from scipy.spatial.transform import Rotation + + +@torch.no_grad() +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--checkpoint", default="outputs/libero_ft2/step_15000") + parser.add_argument("--suite", default="libero_spatial") + parser.add_argument("--tasks", type=int, default=3) + args = parser.parse_args() + + from lerobot.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata + from lerobot.envs.factory import make_env, make_env_config + from transformers import AutoTokenizer + from tinyvla.data.mixture import CanonicalSource + from tinyvla.modeling_tinyvla import TinyVLAPolicy + + policy = TinyVLAPolicy.from_pretrained(args.checkpoint).cuda().eval() + cfg = policy.config + tok = AutoTokenizer.from_pretrained(cfg.lm_model_name) + + meta = LeRobotDatasetMetadata("HuggingFaceVLA/libero") + ds = LeRobotDataset( + "HuggingFaceVLA/libero", + delta_timestamps={"action": [t / meta.fps for t in range(cfg.chunk_size)]}, + video_backend="torchcodec", + ) + src = CanonicalSource(ds, 2, cfg.image_size, cfg.max_state_dim, cfg.max_action_dim) + s_stats = meta.stats["observation.state"] + s_mean = torch.as_tensor(s_stats["mean"]).flatten().float() + s_std = torch.as_tensor(s_stats["std"]).flatten().float().clamp(min=1e-6) + + env_cfg = make_env_config("libero", task=args.suite) + task_envs = make_env(env_cfg, n_envs=1)[args.suite] + env_by_task = {} + for tid, env in task_envs.items(): + desc = env.get_attr("task_description")[0] + env_by_task[desc.strip().lower()] = (tid, env) + + eps_meta = ds.meta.episodes + first_ep_by_task = {} + for ep in range(ds.num_episodes): + start = int(eps_meta["dataset_from_index"][ep]) + task = ds[start]["task"].strip().lower() + if task in env_by_task and task not in first_ep_by_task: + first_ep_by_task[task] = ep + + def tok_batch(task_text): + t = tok([task_text], padding=True, truncation=True, + max_length=cfg.tokenizer_max_length, return_tensors="pt") + return t["input_ids"].cuda(), t["attention_mask"].bool().cuda() + + def env_to_batch(obs, task_text): + imgs = {} + for slot, key in (("cam0", "image"), ("cam1", "image2")): + x = torch.as_tensor(np.asarray(obs["pixels"][key]))[0].flip(0).flip(1) + x = x.permute(2, 0, 1).float() / 255.0 + x = torch.nn.functional.interpolate(x[None], size=(cfg.image_size, cfg.image_size), + mode="bilinear", align_corners=False)[0] + imgs[slot] = x + rs = obs["robot_state"] + pos = np.asarray(rs["eef"]["pos"]).flatten() + quat = np.asarray(rs["eef"]["quat"]).flatten() + rotvec = Rotation.from_quat(quat).as_rotvec() + if rotvec[0] < 0: + th = np.linalg.norm(rotvec) + rotvec = rotvec * (th - 2 * np.pi) / th + grip = np.asarray(rs["gripper"]["qpos"]).flatten() + state = torch.tensor(np.concatenate([pos, rotvec, grip]), dtype=torch.float32) + state = (state - s_mean) / s_std + state = torch.nn.functional.pad(state, (0, cfg.max_state_dim - state.shape[-1])) + ids, mask = tok_batch(task_text) + return { + "observation.images.cam0": imgs["cam0"][None].cuda(), + "observation.images.cam1": imgs["cam1"][None].cuda(), + "observation.state": state[None].cuda(), + "observation.language.tokens": ids, + "observation.language.attention_mask": mask, + "embodiment_id": torch.tensor([2], device="cuda"), + } + + def ds_to_batch(item): + ids, mask = tok_batch(item.pop("task")) + b = {k: v[None].cuda() for k, v in item.items() if torch.is_tensor(v)} + b["observation.language.tokens"] = ids + b["observation.language.attention_mask"] = mask + return b + + for task, ep in list(first_ep_by_task.items())[: args.tasks]: + tid, env = env_by_task[task] + obs, _ = env.reset(seed=0) + start = int(eps_meta["dataset_from_index"][ep]) + item = src[start] + gt = item["action"].clone()[None].cuda() + + env_b = env_to_batch(obs, task) + ds_b = ds_to_batch(dict(item)) + + torch.manual_seed(0) + pred_env = policy.predict_action_chunk(env_b) + torch.manual_seed(0) + pred_ds = policy.predict_action_chunk(ds_b) + + m = item["action_dim_mask"] + d_env_gt = ((pred_env[0, :, m].cpu() - gt[0, :, m].cpu()) ** 2).mean().sqrt().item() + d_ds_gt = ((pred_ds[0, :, m].cpu() - gt[0, :, m].cpu()) ** 2).mean().sqrt().item() + d_env_ds = ((pred_env[0, :, m] - pred_ds[0, :, m]) ** 2).mean().sqrt().item() + + # also compare the raw inputs + img_gap = (env_b["observation.images.cam0"] - ds_b["observation.images.cam0"]).abs().mean().item() + st_gap = (env_b["observation.state"] - ds_b["observation.state"]).abs().max().item() + print(f"task {tid}: RMSE(pred_env, GT)={d_env_gt:.3f} | RMSE(pred_ds, GT)={d_ds_gt:.3f} " + f"| RMSE(pred_env, pred_ds)={d_env_ds:.3f} | img L1 gap={img_gap:.4f} | state max gap={st_gap:.2f}") + + +if __name__ == "__main__": + main() diff --git a/scripts/download_subsets.py b/scripts/download_subsets.py new file mode 100644 index 0000000000000000000000000000000000000000..d013e0fbdac1e92850f0990b99e3082a4f46a8d0 --- /dev/null +++ b/scripts/download_subsets.py @@ -0,0 +1,77 @@ +#!/usr/bin/env python +"""Download Stage-2 dataset subsets under a hard disk budget. + +Usage: + python scripts/download_subsets.py --root ~/tinyvla_data [--dry-run] + +Downloads metadata first, measures actual on-disk size of a small episode +sample, then extrapolates before committing to the full subset download. +""" + +from __future__ import annotations + +import argparse +import shutil +from pathlib import Path + +DISK_BUDGET_GB = 400 + +# (repo_id, num_episodes or None for all) +SUBSETS = [ + ("HuggingFaceVLA/community_dataset_v1", None), + ("nvidia/BridgeData2_LeRobot_v3", 10_000), + ("IPEC-COMMUNITY/fractal20220817_data_lerobot", 8_000), +] + + +def free_gb(path: Path) -> float: + return shutil.disk_usage(path).free / 1e9 + + +def used_gb(path: Path) -> float: + return sum(f.stat().st_size for f in path.rglob("*") if f.is_file()) / 1e9 + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--root", type=Path, default=Path.home() / "tinyvla_data") + parser.add_argument("--dry-run", action="store_true") + parser.add_argument("--only", type=str, default=None, help="substring filter on repo_id") + args = parser.parse_args() + + args.root.mkdir(parents=True, exist_ok=True) + + from lerobot.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata + + for repo_id, n_eps in SUBSETS: + if args.only and args.only not in repo_id: + continue + print(f"\n=== {repo_id} (episodes: {n_eps or 'all'}) ===") + meta = LeRobotDatasetMetadata(repo_id) + total_eps = meta.total_episodes + print(f"total episodes: {total_eps}, fps: {meta.fps}") + + episodes = list(range(min(n_eps, total_eps))) if n_eps else None + + # probe with 1% to estimate size + probe_n = max(10, (len(episodes) if episodes else total_eps) // 100) + probe_dir = args.root / "_probe" / repo_id.replace("/", "__") + if not args.dry_run and not probe_dir.exists(): + LeRobotDataset(repo_id, root=probe_dir, episodes=list(range(probe_n))) + probe_gb = used_gb(probe_dir) + est_gb = probe_gb / probe_n * (len(episodes) if episodes else total_eps) + print(f"probe: {probe_n} eps = {probe_gb:.2f}GB -> estimated full subset {est_gb:.0f}GB") + if est_gb > free_gb(args.root) - 100 or used_gb(args.root) + est_gb > DISK_BUDGET_GB: + print(f"SKIP {repo_id}: would exceed budget ({DISK_BUDGET_GB}GB) or disk") + continue + + if args.dry_run: + continue + + target = args.root / repo_id.replace("/", "__") + LeRobotDataset(repo_id, root=target, episodes=episodes) + print(f"downloaded {repo_id}: {used_gb(target):.1f}GB, total used {used_gb(args.root):.1f}GB") + + +if __name__ == "__main__": + main() diff --git a/scripts/eval_libero.py b/scripts/eval_libero.py new file mode 100644 index 0000000000000000000000000000000000000000..15b7ee0228ce7b78541a5f3b56b7faccfc4ff410 --- /dev/null +++ b/scripts/eval_libero.py @@ -0,0 +1,193 @@ +#!/usr/bin/env python +"""Closed-loop LIBERO evaluation for canonical-schema TinyVLA checkpoints. + +Adapts env observations to the canonical schema the policy was trained on +(cam0/cam1, padded normalized state, embodiment_id) and unnormalizes the +predicted actions with the LIBERO dataset stats. + +Usage: + python scripts/eval_libero.py --checkpoint outputs/libero_ft/final \ + --suite libero_spatial --episodes 20 [--refresh-s 1.0] +""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +import torch + + +def make_normalizer(repo_id="HuggingFaceVLA/libero"): + from lerobot.datasets.lerobot_dataset import LeRobotDatasetMetadata + + meta = LeRobotDatasetMetadata(repo_id) + stats = meta.stats + + def norm(key, x): + s = stats[key] + mean = torch.as_tensor(s["mean"], dtype=torch.float32, device=x.device) + std = torch.as_tensor(s["std"], dtype=torch.float32, device=x.device).clamp(min=1e-6) + return (x - mean) / std + + def unnorm_action(x): + s = stats["action"] + mean = torch.as_tensor(s["mean"], dtype=torch.float32, device=x.device) + std = torch.as_tensor(s["std"], dtype=torch.float32, device=x.device) + return x * std + mean + + return norm, unnorm_action, meta + + +@torch.no_grad() +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--checkpoint", type=Path, required=True) + parser.add_argument("--suite", default="libero_spatial", + choices=["libero_spatial", "libero_object", "libero_goal", "libero_10", "libero_90"]) + parser.add_argument("--episodes", type=int, default=20) + parser.add_argument("--embodiment-id", type=int, default=2) + parser.add_argument("--refresh-s", type=float, default=None, + help="if set, refresh the semantic latent only every N seconds (dual-rate mode)") + parser.add_argument("--max-steps", type=int, default=520) + parser.add_argument("--n-action-steps", type=int, default=None, + help="execute only first N actions of each chunk before re-planning") + parser.add_argument("--action-repeat", type=int, default=1, + help="env steps per predicted action (dataset 10fps vs env 20Hz -> 2)") + parser.add_argument("--save-video-dir", type=Path, default=None, + help="save per-episode mp4s of the agentview camera here") + args = parser.parse_args() + + import numpy as np + from lerobot.envs.factory import make_env, make_env_config + from transformers import AutoTokenizer + from tinyvla.modeling_tinyvla import TinyVLAPolicy + + policy = TinyVLAPolicy.from_pretrained(args.checkpoint).cuda().eval() + cfg = policy.config + if args.n_action_steps: + cfg.n_action_steps = args.n_action_steps + tok = AutoTokenizer.from_pretrained(cfg.lm_model_name) + norm, unnorm_action, meta = make_normalizer() + fps = meta.fps + + env_cfg = make_env_config("libero", task=args.suite) + envs_dict = make_env(env_cfg, n_envs=1) + task_envs = envs_dict[args.suite] # {task_id: vec_env} + + from scipy.spatial.transform import Rotation + + def to_canonical(obs, task_text, latent=None): + imgs = {} + for slot, key in (("cam0", "image"), ("cam1", "image2")): + x = torch.as_tensor(np.asarray(obs["pixels"][key])) + if x.dim() == 4: # (1, H, W, C) + x = x[0] + # robosuite renders 180-degree rotated relative to the recorded dataset + x = x.flip(0).flip(1) + x = x.permute(2, 0, 1).float() / 255.0 + if x.shape[-2:] != (cfg.image_size, cfg.image_size): + x = torch.nn.functional.interpolate( + x[None], size=(cfg.image_size, cfg.image_size), mode="bilinear", align_corners=False + )[0] + imgs[slot] = x + + # dataset convention (OpenVLA-style): [eef pos(3), eef axis-angle(3), gripper qpos(2)] + rs = obs["robot_state"] + pos = np.asarray(rs["eef"]["pos"]).flatten() + quat = np.asarray(rs["eef"]["quat"]).flatten() # robosuite: (x, y, z, w) + rotvec = Rotation.from_quat(quat).as_rotvec() + # canonicalize antipodal representation to match dataset convention + # (dataset uses rotvec with positive x-component, ~+pi for downward gripper) + if rotvec[0] < 0: + theta = np.linalg.norm(rotvec) + if theta > 1e-6: + rotvec = rotvec * (theta - 2 * np.pi) / theta + grip = np.asarray(rs["gripper"]["qpos"]).flatten() + state = torch.tensor(np.concatenate([pos, rotvec, grip]), dtype=torch.float32) + state = norm("observation.state", state) + state = torch.nn.functional.pad(state, (0, cfg.max_state_dim - state.shape[-1])) + + t = tok([task_text], padding=True, truncation=True, + max_length=cfg.tokenizer_max_length, return_tensors="pt") + batch = { + "observation.images.cam0": imgs["cam0"][None].cuda(), + "observation.images.cam1": imgs["cam1"][None].cuda(), + "observation.state": state[None].cuda(), + "observation.language.tokens": t["input_ids"].cuda(), + "observation.language.attention_mask": t["attention_mask"].bool().cuda(), + "embodiment_id": torch.tensor([args.embodiment_id], device="cuda"), + } + if latent is not None: + batch["semantic_latent"] = latent + return batch + + action_dim = meta.features["action"]["shape"][0] + total, succ = 0, 0 + per_task = {} + + task_ids = sorted(task_envs.keys()) + ep_plan = [(tid, i) for i in range((args.episodes + len(task_ids) - 1) // len(task_ids)) for tid in task_ids] + ep_plan = ep_plan[: args.episodes] + + if args.save_video_dir: + args.save_video_dir.mkdir(parents=True, exist_ok=True) + + for ep, (tid, rep) in enumerate(ep_plan): + env = task_envs[tid] + obs, info = env.reset(seed=1000 + rep) + frames = [] if args.save_video_dir else None + try: + task_text = env.get_attr("task_description")[0] + except Exception: + task_text = getattr(getattr(env, "envs", [None])[0], "task_description", "") + policy.reset() + latent = None + last_refresh = -1e9 + done = False + step_i = 0 + ep_succ = False + while not done and step_i < args.max_steps: + t_now = step_i / fps + batch = to_canonical(obs, task_text) + if args.refresh_s is not None: + if t_now - last_refresh >= args.refresh_s: + latent = policy._semantic_latent(batch) + last_refresh = t_now + batch["semantic_latent"] = latent + act_norm = policy.select_action(batch) # (1, max_action_dim) normalized padded + act = unnorm_action(act_norm[0, :action_dim].cpu()).clamp(-1, 1) + for _ in range(args.action_repeat): + try: + obs, reward, terminated, truncated, info = env.step(act.numpy()[None]) + except ValueError: # stepped into env's internal horizon + done = True + break + done = bool(terminated[0] or truncated[0]) + if info.get("is_success") is not None: + ep_succ = ep_succ or bool(np.asarray(info["is_success"]).flatten()[0]) + if frames is not None: + frames.append(np.asarray(obs["pixels"]["image"])[0][::-1, ::-1]) + step_i += 1 + if done or step_i >= args.max_steps: + done = done or step_i >= args.max_steps + break + total += 1 + succ += int(ep_succ) + per_task.setdefault(task_text[:50], []).append(int(ep_succ)) + print(f"ep {ep}: {'SUCCESS' if ep_succ else 'fail'} ({step_i} steps) | {task_text[:60]}") + if frames: + import imageio.v2 as imageio + + tag = "succ" if ep_succ else "fail" + path = args.save_video_dir / f"ep{ep:02d}_task{tid}_{tag}.mp4" + imageio.mimwrite(path, frames, fps=20, quality=7) + + print(f"\n=== {args.suite} | {args.checkpoint} | refresh={args.refresh_s} ===") + print(f"success rate: {succ}/{total} = {succ/total:.1%}") + for t, v in per_task.items(): + print(f" {sum(v)}/{len(v)} {t}") + + +if __name__ == "__main__": + main() diff --git a/scripts/eval_offline.py b/scripts/eval_offline.py new file mode 100644 index 0000000000000000000000000000000000000000..ead9de1fc22e329b814c9e466397260f8aedc1a3 --- /dev/null +++ b/scripts/eval_offline.py @@ -0,0 +1,114 @@ +#!/usr/bin/env python +"""Offline eval for canonical-schema checkpoints (stage2+). + +Reports action-chunk MSE (normalized space) on held-out episodes, the +per-timestep error curve, and stale-latent degradation. + +Usage: + python scripts/eval_offline.py \ + --checkpoint outputs/stage2_mixture/final \ + --repo-id VoicAndrei__so100_kitchen \ + --root ~/tinyvla_data/so101_v3/VoicAndrei__so100_kitchen \ + --episodes 8 --stale-s 0 1 2 [--embodiment-id 0] [--no-latent] +""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +import torch + + +@torch.no_grad() +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--checkpoint", type=Path, required=True) + parser.add_argument("--repo-id", required=True) + parser.add_argument("--root", default=None) + parser.add_argument("--episodes", type=int, default=8) + parser.add_argument("--stride", type=int, default=30) + parser.add_argument("--stale-s", type=float, nargs="*", default=[0.0, 1.0, 2.0]) + parser.add_argument("--embodiment-id", type=int, default=0) + parser.add_argument("--split", choices=["first", "last"], default="last") + args = parser.parse_args() + + from lerobot.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata + from transformers import AutoTokenizer + from tinyvla.data.mixture import CanonicalSource + from tinyvla.modeling_tinyvla import TinyVLAPolicy + + policy = TinyVLAPolicy.from_pretrained(args.checkpoint).cuda().eval() + cfg = policy.config + chunk = cfg.chunk_size + tok = AutoTokenizer.from_pretrained(cfg.lm_model_name) + + meta = LeRobotDatasetMetadata(args.repo_id, root=args.root) + ds = LeRobotDataset( + args.repo_id, + root=args.root, + delta_timestamps={"action": [t / meta.fps for t in range(chunk)]}, + video_backend="torchcodec", + ) + src = CanonicalSource( + ds, args.embodiment_id, cfg.image_size, cfg.max_state_dim, cfg.max_action_dim + ) + + if args.split == "first": + eps = list(range(args.episodes)) + else: + eps = list(range(ds.num_episodes - args.episodes, ds.num_episodes)) + + def to_batch(item): + t = tok([item.pop("task")], padding=True, truncation=True, + max_length=cfg.tokenizer_max_length, return_tensors="pt") + b = {k: v[None].cuda() if torch.is_tensor(v) else v for k, v in item.items()} + b["observation.language.tokens"] = t["input_ids"].cuda() + b["observation.language.attention_mask"] = t["attention_mask"].bool().cuda() + return b + + results = {s: [] for s in args.stale_s} + per_t = torch.zeros(chunk) + n_chunks = 0 + + for ep in eps: + start = int(ds.meta.episodes["dataset_from_index"][ep]) + end = int(ds.meta.episodes["dataset_to_index"][ep]) + for idx in range(start, end - 1, args.stride): + if idx >= len(src): + break + item = src[idx - 0] + gt = item["action"].clone() # (chunk, A) normalized + mask = item["action_dim_mask"].clone() + pad = item.get("action_is_pad") + batch = to_batch(dict(item)) + for stale_s in args.stale_s: + b = dict(batch) + if stale_s > 0: + stale_idx = max(start, idx - int(stale_s * ds.fps)) + stale_item = src[stale_idx] + sb = to_batch(dict(stale_item)) + b["semantic_latent"] = policy._semantic_latent(sb) + pred = policy.predict_action_chunk(b)[0].cpu() # (chunk, A) normalized + err = (pred[:, mask] - gt[:, mask]) ** 2 + if pad is not None: + err = err[~pad] + mse = err.mean().item() + results[stale_s].append(mse) + if stale_s == 0: + e = ((pred - gt) ** 2)[:, mask].mean(dim=-1) + if pad is not None: + e = e * (~pad).float() + per_t += e + n_chunks += 1 + + print(f"\n=== {args.repo_id} | {len(eps)} {args.split} episodes | {n_chunks} chunks | ckpt {args.checkpoint} ===") + for s, vals in results.items(): + print(f"stale {s:.0f}s: normalized chunk MSE {sum(vals)/len(vals):.4f}") + curve = (per_t / max(n_chunks, 1)).sqrt() + print("per-timestep normalized RMSE (t=0,10,25,49):", + [round(curve[i].item(), 3) for i in (0, 10, 25, 49)]) + + +if __name__ == "__main__": + main() diff --git a/scripts/label_spatial.py b/scripts/label_spatial.py new file mode 100644 index 0000000000000000000000000000000000000000..bb872ab90521dd19c032ad3a2651a6de06c41833 --- /dev/null +++ b/scripts/label_spatial.py @@ -0,0 +1,148 @@ +#!/usr/bin/env python +"""Stage 3: spatial distillation labels from a Qwen3-VL teacher (XS-VLA style). + +For every Nth frame of each episode, ask the teacher for the bounding box of +the task-relevant object; quantize the box center onto a GRID x GRID map and +store the cell index. The student later learns a linear classifier over its +fast-path spatial tokens with CE x 0.15 on labeled frames. + +Output: parquet with columns (dataset, episode_index, frame_index, cell, +cx, cy, confidence_ok) at ~/tinyvla_data/spatial_labels/.parquet + +Usage: + python scripts/label_spatial.py --teacher Qwen/Qwen3-VL-4B-Instruct \ + --frame-stride 10 --episode-frac 0.4 [--datasets-limit 2] [--pilot 20] +""" + +from __future__ import annotations + +import argparse +import json +import re +from pathlib import Path + +import torch + +GRID = 32 +OUT_DIR = Path.home() / "tinyvla_data" / "spatial_labels" + +POINT_RE = re.compile(r"\[?\s*(\d+)\s*,\s*(\d+)\s*\]?") + +PROMPT = ( + "Task: {task}\n" + "Look at the image. Where is the single object the robot must interact with next " + "to accomplish this task? Answer with ONLY its center point as [x, y] in 0-1000 " + "normalized coordinates. If unsure, output [0,0]." +) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--teacher", default="Qwen/Qwen3.5-4B") + parser.add_argument("--data-root", type=Path, default=Path.home() / "tinyvla_data/so101_v3") + parser.add_argument("--frame-stride", type=int, default=10) + parser.add_argument("--episode-frac", type=float, default=0.4) + parser.add_argument("--datasets-limit", type=int, default=None) + parser.add_argument("--pilot", type=int, default=None, help="label only N frames total, print results") + parser.add_argument("--batch-size", type=int, default=16) + args = parser.parse_args() + + import pyarrow as pa + import pyarrow.parquet as pq + from lerobot.datasets.lerobot_dataset import LeRobotDataset + from transformers import AutoModelForImageTextToText, AutoProcessor + + model = AutoModelForImageTextToText.from_pretrained( + args.teacher, dtype=torch.bfloat16, device_map="cuda" + ) + proc = AutoProcessor.from_pretrained(args.teacher) + OUT_DIR.mkdir(parents=True, exist_ok=True) + + roots = sorted(args.data_root.iterdir()) + if args.datasets_limit: + roots = roots[: args.datasets_limit] + + total_done = 0 + for root in roots: + if not (root / "meta/info.json").exists(): + continue + out_path = OUT_DIR / f"{root.name}.parquet" + if out_path.exists() and not args.pilot: + continue + ds = LeRobotDataset(root.name, root=root, video_backend="torchcodec") + image_key = sorted(k for k in ds.meta.features if k.startswith("observation.images"))[0] + n_eps = max(1, int(ds.num_episodes * args.episode_frac)) + + rows = [] + pending = [] # (ep, fi, image_pil, task) + + def flush(): + nonlocal total_done + if not pending: + return + msgs = [ + [{"role": "user", "content": [ + {"type": "image", "image": img}, + {"type": "text", "text": PROMPT.format(task=task)}, + ]}] + for _, _, img, task in pending + ] + texts = [ + proc.apply_chat_template( + m, tokenize=False, add_generation_prompt=True, enable_thinking=False + ) + for m in msgs + ] + images = [[p[2]] for p in pending] + inputs = proc(text=texts, images=images, return_tensors="pt", padding=True).to("cuda") + with torch.no_grad(): + out = model.generate(**inputs, max_new_tokens=16, do_sample=False) + answers = proc.batch_decode(out[:, inputs["input_ids"].shape[1]:], skip_special_tokens=True) + for (ep, fi, _, _), ans in zip(pending, answers): + m = POINT_RE.search(ans) + ok = False + cell, cx, cy = -1, -1.0, -1.0 + if m: + px, py = (int(g) for g in m.groups()) + if 0 < px <= 1000 and 0 < py <= 1000: + cx, cy = px / 1000.0, py / 1000.0 + gx, gy = min(GRID - 1, int(cx * GRID)), min(GRID - 1, int(cy * GRID)) + cell = gy * GRID + gx + ok = True + rows.append({"dataset": root.name, "episode_index": ep, "frame_index": fi, + "cell": cell, "cx": cx, "cy": cy, "confidence_ok": ok}) + if args.pilot: + print(f"ep{ep} f{fi}: '{ans.strip()[:60]}' -> cell {cell} ({cx:.2f},{cy:.2f})") + total_done += len(pending) + pending.clear() + + from torchvision.transforms.functional import to_pil_image + import torch.nn.functional as F + + for ep in range(n_eps): + start = int(ds.meta.episodes["dataset_from_index"][ep]) + end = int(ds.meta.episodes["dataset_to_index"][ep]) + for idx in range(start, end, args.frame_stride): + item = ds[idx] + # label on the same 256^2 view the student sees; ~8x fewer + # teacher vision tokens than full res + small = F.interpolate( + item[image_key][None].clamp(0, 1), size=(256, 256), + mode="bilinear", align_corners=False, + )[0] + img = to_pil_image(small) + pending.append((ep, idx - start, img, item.get("task") or "")) + if len(pending) >= args.batch_size: + flush() + if args.pilot and total_done + len(pending) >= args.pilot: + flush() + print(f"pilot done: {total_done} frames") + return + flush() + pq.write_table(pa.Table.from_pylist(rows), out_path) + ok_rate = sum(r["confidence_ok"] for r in rows) / max(len(rows), 1) + print(f"{root.name}: {len(rows)} labels -> {out_path} (ok {ok_rate:.1%})") + + +if __name__ == "__main__": + main() diff --git a/scripts/prepare_community_v3.py b/scripts/prepare_community_v3.py new file mode 100644 index 0000000000000000000000000000000000000000..0df30de0cab0a318b8f9b22fcee142aa52f8ba02 --- /dev/null +++ b/scripts/prepare_community_v3.py @@ -0,0 +1,163 @@ +#!/usr/bin/env python +"""Prepare the SO-100/101 slice of HuggingFaceVLA/community_dataset_v3. + +Steps: + catalog — fetch all sub-dataset info.json, write catalog.json + select — filter SO100/101 single-arm, rank by episodes, write selection.json + convert — download each selected subdir + convert v2.1 -> v3.0 locally + +Usage: + python scripts/prepare_community_v3.py catalog + python scripts/prepare_community_v3.py select --target-episodes 12000 + python scripts/prepare_community_v3.py convert [--limit N] +""" + +from __future__ import annotations + +import argparse +import json +import shutil +from concurrent.futures import ThreadPoolExecutor, as_completed +from pathlib import Path + +AGG_REPO = "HuggingFaceVLA/community_dataset_v3" +DISK_BUDGET_GB = 400 +DATA_ROOT = Path.home() / "tinyvla_data" +CATALOG = DATA_ROOT / "community_v3_catalog.json" +SELECTION = DATA_ROOT / "community_v3_selection.json" +CONVERTED_DIR = DATA_ROOT / "so101_v3" +SO_ROBOTS = {"so100", "so101", "so100_follower", "so101_follower"} + + +def cmd_catalog(args): + from huggingface_hub import HfApi, hf_hub_download + + api = HfApi() + files = api.list_repo_files(AGG_REPO, repo_type="dataset") + subs = sorted(f.rsplit("/meta/info.json", 1)[0] for f in files if f.endswith("meta/info.json")) + print(f"{len(subs)} sub-datasets") + + def fetch(sub): + try: + p = hf_hub_download(AGG_REPO, f"{sub}/meta/info.json", repo_type="dataset") + info = json.load(open(p)) + return sub, { + "robot_type": info.get("robot_type"), + "episodes": info.get("total_episodes"), + "frames": info.get("total_frames"), + "fps": info.get("fps"), + "version": info.get("codebase_version"), + "action_shape": info.get("features", {}).get("action", {}).get("shape"), + "cameras": [k for k in info.get("features", {}) if k.startswith("observation.images")], + } + except Exception as e: + return sub, {"error": str(e)[:100]} + + catalog = {} + with ThreadPoolExecutor(16) as ex: + futs = [ex.submit(fetch, s) for s in subs] + for i, f in enumerate(as_completed(futs)): + sub, meta = f.result() + catalog[sub] = meta + if (i + 1) % 100 == 0: + print(f"{i+1}/{len(subs)}") + + DATA_ROOT.mkdir(parents=True, exist_ok=True) + CATALOG.write_text(json.dumps(catalog, indent=1)) + print(f"wrote {CATALOG}") + + +def cmd_select(args): + catalog = json.loads(CATALOG.read_text()) + rows = [ + (sub, m) + for sub, m in catalog.items() + if m.get("robot_type") in SO_ROBOTS + and m.get("action_shape") == [6] + and m.get("fps") == 30 + and m.get("episodes") + and m.get("cameras") + ] + # prefer larger datasets: fewer conversions per episode + rows.sort(key=lambda r: -r[1]["episodes"]) + picked, total = [], 0 + for sub, m in rows: + if total >= args.target_episodes: + break + picked.append({"sub": sub, **m}) + total += m["episodes"] + print(f"{len(rows)} eligible; picked {len(picked)} datasets, {total} episodes") + SELECTION.write_text(json.dumps(picked, indent=1)) + print(f"wrote {SELECTION}") + + +def free_gb(path: Path) -> float: + return shutil.disk_usage(path).free / 1e9 + + +def used_gb_cached(path: Path) -> float: + import subprocess + + out = subprocess.run(["du", "-s", "--block-size=1G", str(path)], capture_output=True, text=True) + return float(out.stdout.split()[0]) if out.returncode == 0 else 0.0 + + +def cmd_convert(args): + from huggingface_hub import snapshot_download + from lerobot.scripts.convert_dataset_v21_to_v30 import convert_dataset + + picked = json.loads(SELECTION.read_text()) + if args.limit: + picked = picked[: args.limit] + CONVERTED_DIR.mkdir(parents=True, exist_ok=True) + raw_dir = DATA_ROOT / "_raw_v21" + done, failed = 0, [] + for item in picked: + sub = item["sub"] + name = sub.replace("/", "__") + out = CONVERTED_DIR / name + if (out / "meta" / "info.json").exists(): + done += 1 + continue + local = raw_dir / name + try: + snapshot_download( + AGG_REPO, + repo_type="dataset", + allow_patterns=[f"{sub}/*"], + local_dir=raw_dir / "_dl", + ) + src = raw_dir / "_dl" / sub + if local.exists(): + shutil.rmtree(local) + shutil.move(str(src), str(local)) + convert_dataset(repo_id=name, root=local, push_to_hub=False, force_conversion=True) + # converter writes v3.0 in place at root + shutil.move(str(local), str(out)) + # drop the v2.1 originals the converter stashes as _old + old = local.parent / (local.name + "_old") + if old.exists(): + shutil.rmtree(old) + done += 1 + print(f"[{done}/{len(picked)}] {sub}: converted -> {out} (free {free_gb(DATA_ROOT):.0f}GB)") + if free_gb(DATA_ROOT) < 150 or used_gb_cached(DATA_ROOT) > DISK_BUDGET_GB: + print("STOP: disk budget reached") + break + except Exception as e: + failed.append(sub) + print(f"FAIL {sub}: {type(e).__name__}: {str(e)[:200]}") + print(f"done={done} failed={len(failed)}") + if failed: + (DATA_ROOT / "convert_failures.json").write_text(json.dumps(failed)) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + sub = parser.add_subparsers(dest="cmd", required=True) + sub.add_parser("catalog") + p_sel = sub.add_parser("select") + p_sel.add_argument("--target-episodes", type=int, default=12_000) + p_conv = sub.add_parser("convert") + p_conv.add_argument("--limit", type=int, default=None) + args = parser.parse_args() + {"catalog": cmd_catalog, "select": cmd_select, "convert": cmd_convert}[args.cmd](args) diff --git a/scripts/replay_libero.py b/scripts/replay_libero.py new file mode 100644 index 0000000000000000000000000000000000000000..4deafbd6996ad4b6f7aabad7e0525609e10a7037 --- /dev/null +++ b/scripts/replay_libero.py @@ -0,0 +1,85 @@ +#!/usr/bin/env python +"""Replay ground-truth dataset actions in the LIBERO env. + +Diagnoses action-convention correctness independently of the policy: if GT +actions from a matching task episode achieve success (at least sometimes), +the env adapter/action semantics are right and remaining failures are model +quality; if GT never succeeds, the action mapping is wrong. + +Usage: + python scripts/replay_libero.py --suite libero_spatial --episodes 5 --action-repeat 2 +""" + +from __future__ import annotations + +import argparse + +import numpy as np +import torch + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--suite", default="libero_spatial") + parser.add_argument("--episodes", type=int, default=5) + parser.add_argument("--action-repeat", type=int, default=2) + args = parser.parse_args() + + from lerobot.datasets.lerobot_dataset import LeRobotDataset + from lerobot.envs.factory import make_env, make_env_config + + ds = LeRobotDataset("HuggingFaceVLA/libero", video_backend="torchcodec") + env_cfg = make_env_config("libero", task=args.suite) + task_envs = make_env(env_cfg, n_envs=1)[args.suite] + + # map task_description -> env + env_by_task = {} + for tid, env in task_envs.items(): + try: + desc = env.get_attr("task_description")[0] + except Exception: + desc = getattr(env.envs[0], "task_description", "") + env_by_task[desc.strip().lower()] = (tid, env) + + # first dataset episode per task <-> env's first reset (init_state_id=0), + # assuming dataset episode order preserves LIBERO demo order within a task + eps_meta = ds.meta.episodes + first_ep_by_task = {} + for ep in range(ds.num_episodes): + start = int(eps_meta["dataset_from_index"][ep]) + task = ds[start]["task"].strip().lower() + if task in env_by_task and task not in first_ep_by_task: + first_ep_by_task[task] = ep + + replayed, succ = 0, 0 + for task, ep in list(first_ep_by_task.items())[: args.episodes]: + start = int(eps_meta["dataset_from_index"][ep]) + end = int(eps_meta["dataset_to_index"][ep]) + tid, env = env_by_task[task] + obs, info = env.reset(seed=0) # fresh env: init_state_id starts at 0 + ep_succ = False + for idx in range(start, end): + act = ds[idx]["action"].numpy() + done = False + for _ in range(args.action_repeat): + try: + obs, reward, term, trunc, info = env.step(act[None]) + except ValueError: + done = True + break + if info.get("is_success") is not None: + ep_succ = ep_succ or bool(np.asarray(info["is_success"]).flatten()[0]) + if term[0] or trunc[0]: + done = True + break + if done or ep_succ: + break + replayed += 1 + succ += int(ep_succ) + print(f"ep {ep} (task {tid}): {'SUCCESS' if ep_succ else 'fail'} | {task[:60]}") + + print(f"\nGT replay success: {succ}/{replayed}") + + +if __name__ == "__main__": + main() diff --git a/scripts/train.py b/scripts/train.py new file mode 100644 index 0000000000000000000000000000000000000000..66527abce14f676addbab2c91c78ea1287c11206 --- /dev/null +++ b/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() diff --git a/scripts/train_stage1_smoke.sh b/scripts/train_stage1_smoke.sh new file mode 100644 index 0000000000000000000000000000000000000000..bed8b84ac3d5e61fd2c39ccba7e9ccb508c8a6bb --- /dev/null +++ b/scripts/train_stage1_smoke.sh @@ -0,0 +1,20 @@ +#!/usr/bin/env bash +# Stage 1: overfit smoke on svla_so100_stacking via lerobot-train. +set -euo pipefail + +source ~/tinyvla_venv/bin/activate + +lerobot-train \ + --policy.type=tinyvla \ + --policy.discover_packages_path=tinyvla \ + --policy.device=cuda \ + --policy.push_to_hub=false \ + --dataset.repo_id=lerobot/svla_so100_stacking \ + --dataset.video_backend=torchcodec \ + --batch_size=32 \ + --num_workers=8 \ + --steps=5000 \ + --log_freq=50 \ + --save_freq=2500 \ + --output_dir=outputs/stage1_smoke \ + "$@" diff --git a/src/tinyvla.egg-info/PKG-INFO b/src/tinyvla.egg-info/PKG-INFO new file mode 100644 index 0000000000000000000000000000000000000000..8775ad393714c0745e9a4bc7b21f470ecc7c2f08 --- /dev/null +++ b/src/tinyvla.egg-info/PKG-INFO @@ -0,0 +1,15 @@ +Metadata-Version: 2.4 +Name: tinyvla +Version: 0.1.0 +Summary: Dual-rate VLA: Qwen3-0.6B slow semantic path + ~90M fast visuomotor path for SO-100/SO-101 +Requires-Python: >=3.10 +Requires-Dist: lerobot>=0.4 +Requires-Dist: timm>=1.0.15 +Requires-Dist: torch>=2.6 +Requires-Dist: torchvision +Requires-Dist: transformers>=4.51 +Requires-Dist: accelerate>=1.0 +Requires-Dist: wandb +Requires-Dist: pyarrow +Provides-Extra: dev +Requires-Dist: pytest; extra == "dev" diff --git a/src/tinyvla.egg-info/SOURCES.txt b/src/tinyvla.egg-info/SOURCES.txt new file mode 100644 index 0000000000000000000000000000000000000000..26abda1a46185581fd360163d64018af2d6128db --- /dev/null +++ b/src/tinyvla.egg-info/SOURCES.txt @@ -0,0 +1,7 @@ +pyproject.toml +src/tinyvla.egg-info/PKG-INFO +src/tinyvla.egg-info/SOURCES.txt +src/tinyvla.egg-info/dependency_links.txt +src/tinyvla.egg-info/requires.txt +src/tinyvla.egg-info/top_level.txt +src/tinyvla/modules/vision.py \ No newline at end of file diff --git a/src/tinyvla.egg-info/dependency_links.txt b/src/tinyvla.egg-info/dependency_links.txt new file mode 100644 index 0000000000000000000000000000000000000000..8b137891791fe96927ad78e64b0aad7bded08bdc --- /dev/null +++ b/src/tinyvla.egg-info/dependency_links.txt @@ -0,0 +1 @@ + diff --git a/src/tinyvla.egg-info/requires.txt b/src/tinyvla.egg-info/requires.txt new file mode 100644 index 0000000000000000000000000000000000000000..6206b7b01bd66ee181c69d7ecfd713cc6c7ba14e --- /dev/null +++ b/src/tinyvla.egg-info/requires.txt @@ -0,0 +1,11 @@ +lerobot>=0.4 +timm>=1.0.15 +torch>=2.6 +torchvision +transformers>=4.51 +accelerate>=1.0 +wandb +pyarrow + +[dev] +pytest diff --git a/src/tinyvla.egg-info/top_level.txt b/src/tinyvla.egg-info/top_level.txt new file mode 100644 index 0000000000000000000000000000000000000000..c723e8829bdae705410ea629b7bb4c33e21b6809 --- /dev/null +++ b/src/tinyvla.egg-info/top_level.txt @@ -0,0 +1 @@ +tinyvla diff --git a/src/tinyvla/__init__.py b/src/tinyvla/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9b1e318d33509f5fee7f7c755ff127b42f8a6c11 --- /dev/null +++ b/src/tinyvla/__init__.py @@ -0,0 +1,9 @@ +"""TinyVLA: dual-rate VLA (Qwen3-0.6B slow path + ~90M fast visuomotor path). + +Importing this package registers the "tinyvla" policy with LeRobot +(use --policy.type=tinyvla --policy.discover_packages_path=tinyvla). +""" + +from .configuration_tinyvla import TinyVLAConfig + +__all__ = ["TinyVLAConfig"] diff --git a/src/tinyvla/__pycache__/__init__.cpython-312.pyc b/src/tinyvla/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e0c1f0d3d80e698b29fb45092c50a27ced9b5602 Binary files /dev/null and b/src/tinyvla/__pycache__/__init__.cpython-312.pyc differ diff --git a/src/tinyvla/__pycache__/configuration_tinyvla.cpython-312.pyc b/src/tinyvla/__pycache__/configuration_tinyvla.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5adc6f251a8cfd267a1bfa693eac47ec185266c8 Binary files /dev/null and b/src/tinyvla/__pycache__/configuration_tinyvla.cpython-312.pyc differ diff --git a/src/tinyvla/__pycache__/modeling_tinyvla.cpython-312.pyc b/src/tinyvla/__pycache__/modeling_tinyvla.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3ff31a1b522048d03621828a28caab45c71a5f09 Binary files /dev/null and b/src/tinyvla/__pycache__/modeling_tinyvla.cpython-312.pyc differ diff --git a/src/tinyvla/__pycache__/processor_tinyvla.cpython-312.pyc b/src/tinyvla/__pycache__/processor_tinyvla.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..50063fccb3d14b6a1270655da9a14dfe15b2b438 Binary files /dev/null and b/src/tinyvla/__pycache__/processor_tinyvla.cpython-312.pyc differ diff --git a/src/tinyvla/configuration_tinyvla.py b/src/tinyvla/configuration_tinyvla.py new file mode 100644 index 0000000000000000000000000000000000000000..7b47b0a9beca409f80818048090e408c65f01de5 --- /dev/null +++ b/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 diff --git a/src/tinyvla/data/__init__.py b/src/tinyvla/data/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/src/tinyvla/data/__pycache__/__init__.cpython-312.pyc b/src/tinyvla/data/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3d346b303bf7272d09300304fa6a9b686fbe4ba3 Binary files /dev/null and b/src/tinyvla/data/__pycache__/__init__.cpython-312.pyc differ diff --git a/src/tinyvla/data/__pycache__/mixture.cpython-312.pyc b/src/tinyvla/data/__pycache__/mixture.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ffd90534b8cc45d9ce3fc8d5162bd2876d143b2c Binary files /dev/null and b/src/tinyvla/data/__pycache__/mixture.cpython-312.pyc differ diff --git a/src/tinyvla/data/__pycache__/spatial_labels.cpython-312.pyc b/src/tinyvla/data/__pycache__/spatial_labels.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3f1fc0e2c401de93285522b2d00c23ef3a49eb45 Binary files /dev/null and b/src/tinyvla/data/__pycache__/spatial_labels.cpython-312.pyc differ diff --git a/src/tinyvla/data/mixture.py b/src/tinyvla/data/mixture.py new file mode 100644 index 0000000000000000000000000000000000000000..065cedc892c1c5f908621a19848f61cb64d31b0f --- /dev/null +++ b/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") diff --git a/src/tinyvla/data/spatial_labels.py b/src/tinyvla/data/spatial_labels.py new file mode 100644 index 0000000000000000000000000000000000000000..fc9c2582d7197fd0c135840ddd56afd16bd3dc4a --- /dev/null +++ b/src/tinyvla/data/spatial_labels.py @@ -0,0 +1,44 @@ +"""Spatial-distillation label store. + +Labels come from scripts/label_spatial.py parquets: one row per labeled frame +(dataset, episode_index, frame_index RELATIVE to episode start, cell in +[0, GRID*GRID) or -1). Frames were labeled every `stride`; lookup rounds a +query frame to the nearest labeled one. +""" + +from __future__ import annotations + +from pathlib import Path + +GRID = 32 +STRIDE = 10 + + +class SpatialLabelStore: + def __init__(self, labels_dir: str | Path, dataset_name: str): + import pyarrow.parquet as pq + + self.table = {} + path = Path(labels_dir) / f"{dataset_name}.parquet" + if path.exists(): + t = pq.read_table(path) + for ep, fi, cell, ok in zip( + t["episode_index"].to_pylist(), + t["frame_index"].to_pylist(), + t["cell"].to_pylist(), + t["confidence_ok"].to_pylist(), + ): + if ok and cell >= 0: + self.table[(int(ep), int(fi))] = int(cell) + + def __len__(self): + return len(self.table) + + def lookup(self, episode_index: int, rel_frame: int) -> int: + """Nearest labeled frame's cell, or -1.""" + base = round(rel_frame / STRIDE) * STRIDE + for fi in (base, base - STRIDE, base + STRIDE): + cell = self.table.get((episode_index, fi)) + if cell is not None: + return cell + return -1 diff --git a/src/tinyvla/modeling_tinyvla.py b/src/tinyvla/modeling_tinyvla.py new file mode 100644 index 0000000000000000000000000000000000000000..7c2b580afa60f9ae670b40fbe01f3efdc9817610 --- /dev/null +++ b/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() diff --git a/src/tinyvla/modules/__pycache__/embodiment.cpython-312.pyc b/src/tinyvla/modules/__pycache__/embodiment.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..31f04e5e67c7d2a78bec4d72d032378a31b64bab Binary files /dev/null and b/src/tinyvla/modules/__pycache__/embodiment.cpython-312.pyc differ diff --git a/src/tinyvla/modules/__pycache__/expert.cpython-312.pyc b/src/tinyvla/modules/__pycache__/expert.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..976ca6fece2f4ee213dd6946bc3e67821fbd9c9a Binary files /dev/null and b/src/tinyvla/modules/__pycache__/expert.cpython-312.pyc differ diff --git a/src/tinyvla/modules/__pycache__/semantic.cpython-312.pyc b/src/tinyvla/modules/__pycache__/semantic.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f4155b4be58ab3c64f467451e09f9695c4cc4ab3 Binary files /dev/null and b/src/tinyvla/modules/__pycache__/semantic.cpython-312.pyc differ diff --git a/src/tinyvla/modules/__pycache__/vision.cpython-312.pyc b/src/tinyvla/modules/__pycache__/vision.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b664728cc05548fe870a1b9c8902508d3a93c75c Binary files /dev/null and b/src/tinyvla/modules/__pycache__/vision.cpython-312.pyc differ diff --git a/src/tinyvla/modules/embodiment.py b/src/tinyvla/modules/embodiment.py new file mode 100644 index 0000000000000000000000000000000000000000..543d2283564736bf871aebdb9a67a471fc9dae3c --- /dev/null +++ b/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, :] diff --git a/src/tinyvla/modules/expert.py b/src/tinyvla/modules/expert.py new file mode 100644 index 0000000000000000000000000000000000000000..b482ed33db711ce4fbf503f324525ed290468708 --- /dev/null +++ b/src/tinyvla/modules/expert.py @@ -0,0 +1,128 @@ +"""Flow-matching action expert. + +A small transformer over noisy action tokens with interleaved +cross-attention (to the conditioning set) and self-attention blocks, +following the SmolVLA expert pattern. Predicts the flow velocity. +""" + +from __future__ import annotations + +import math + +import torch +import torch.nn as nn + + +def sincos_time_embedding(t: torch.Tensor, dim: int, min_period: float = 4e-3, max_period: float = 4.0) -> torch.Tensor: + """t: (B,) in [0, 1] -> (B, dim) sine-cosine embedding.""" + half = dim // 2 + freqs = torch.exp( + torch.linspace(math.log(min_period), math.log(max_period), half, device=t.device) + ) + angles = 2 * math.pi * t[:, None] / freqs[None, :] + return torch.cat([angles.sin(), angles.cos()], dim=-1) + + +class _MLP(nn.Module): + def __init__(self, dim: int, ratio: int = 4): + super().__init__() + self.net = nn.Sequential( + nn.Linear(dim, dim * ratio), + nn.GELU(), + nn.Linear(dim * ratio, dim), + ) + + def forward(self, x): + return self.net(x) + + +class SelfAttnBlock(nn.Module): + def __init__(self, dim: int, n_heads: int): + super().__init__() + self.norm1 = nn.LayerNorm(dim) + self.attn = nn.MultiheadAttention(dim, n_heads, batch_first=True) + self.norm2 = nn.LayerNorm(dim) + self.mlp = _MLP(dim) + + def forward(self, x, cond, cond_mask): + h = self.norm1(x) + h, _ = self.attn(h, h, h, need_weights=False) + x = x + h + x = x + self.mlp(self.norm2(x)) + return x + + +class CrossAttnBlock(nn.Module): + def __init__(self, dim: int, n_heads: int): + super().__init__() + self.norm1 = nn.LayerNorm(dim) + self.norm_cond = nn.LayerNorm(dim) + self.attn = nn.MultiheadAttention(dim, n_heads, batch_first=True) + self.norm2 = nn.LayerNorm(dim) + self.mlp = _MLP(dim) + + def forward(self, x, cond, cond_mask): + h = self.norm1(x) + c = self.norm_cond(cond) + # key_padding_mask: True = ignore + kpm = ~cond_mask if cond_mask is not None else None + h, _ = self.attn(h, c, c, key_padding_mask=kpm, need_weights=False) + x = x + h + x = x + self.mlp(self.norm2(x)) + return x + + +class FlowMatchingExpert(nn.Module): + def __init__( + self, + action_dim: int = 8, + chunk_size: int = 50, + d_model: int = 512, + n_layers: int = 12, + n_heads: int = 8, + min_period: float = 4e-3, + max_period: float = 4.0, + ): + super().__init__() + self.action_dim = action_dim + self.chunk_size = chunk_size + self.d_model = d_model + self.min_period = min_period + self.max_period = max_period + + self.action_in = nn.Linear(action_dim, d_model) + self.pos_emb = nn.Parameter(torch.zeros(1, chunk_size, d_model)) + self.time_mlp = nn.Sequential( + nn.Linear(d_model, d_model), + nn.GELU(), + nn.Linear(d_model, d_model), + ) + # even blocks cross-attend to the conditioning set, odd blocks self-attend + self.blocks = nn.ModuleList( + [ + (CrossAttnBlock if i % 2 == 0 else SelfAttnBlock)(d_model, n_heads) + for i in range(n_layers) + ] + ) + self.norm_out = nn.LayerNorm(d_model) + self.action_out = nn.Linear(d_model, action_dim) + + nn.init.trunc_normal_(self.pos_emb, std=0.02) + nn.init.zeros_(self.action_out.weight) + nn.init.zeros_(self.action_out.bias) + + def forward( + self, + noisy_actions: torch.Tensor, # (B, T, A) + time: torch.Tensor, # (B,) in [0, 1] + cond: torch.Tensor, # (B, N, d_model) + cond_mask: torch.Tensor | None = None, # (B, N) True = valid + ) -> torch.Tensor: + x = self.action_in(noisy_actions) + self.pos_emb[:, : noisy_actions.shape[1]] + t_emb = self.time_mlp( + sincos_time_embedding(time, self.d_model, self.min_period, self.max_period) + ) + x = x + t_emb[:, None, :] + for block in self.blocks: + x = block(x, cond, cond_mask) + return self.action_out(self.norm_out(x)) diff --git a/src/tinyvla/modules/semantic.py b/src/tinyvla/modules/semantic.py new file mode 100644 index 0000000000000000000000000000000000000000..ad83724f076da105f14e8f7a6800d163f8482adb --- /dev/null +++ b/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 + + # *n instruction *K + ids = torch.cat( + [ + torch.full((b, 1), self.vision_start_id, dtype=torch.long, device=device), + torch.full((b, n_img), self.image_token_id, dtype=torch.long, device=device), + torch.full((b, 1), self.vision_end_id, dtype=torch.long, device=device), + lang_tokens, + torch.full((b, self.num_readout), self.pad_token_id, dtype=torch.long, device=device), + ], + dim=1, + ) + attn = torch.cat( + [ + torch.ones(b, n_img + 2, dtype=torch.bool, device=device), + lang_mask.bool(), + torch.ones(b, self.num_readout, dtype=torch.bool, device=device), + ], + dim=1, + ) + + 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) diff --git a/src/tinyvla/processor_tinyvla.py b/src/tinyvla/processor_tinyvla.py new file mode 100644 index 0000000000000000000000000000000000000000..dadb4fe3c6c04b1373c8b3f190e7a8013a6e8fba --- /dev/null +++ b/src/tinyvla/processor_tinyvla.py @@ -0,0 +1,45 @@ +"""Pre/post-processor pipelines for TinyVLA (normalization + language tokenization).""" + +from __future__ import annotations + +from typing import Any + +import torch + +from lerobot.processor import ( + PolicyAction, + PolicyProcessorPipeline, + TokenizerProcessorStep, + make_default_policy_processor_steps, + make_policy_processor_pipelines, +) + +from .configuration_tinyvla import TinyVLAConfig + + +def make_tinyvla_pre_post_processors( + config: TinyVLAConfig, + dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None, +) -> tuple[ + PolicyProcessorPipeline[dict[str, Any], dict[str, Any]], + PolicyProcessorPipeline[PolicyAction, PolicyAction], +]: + steps = make_default_policy_processor_steps(config, dataset_stats) + + input_steps = [ + steps.rename_observations, + steps.add_batch_dim, + TokenizerProcessorStep( + tokenizer_name=config.lm_model_name, + padding=config.pad_language_to, + padding_side="right", + max_length=config.tokenizer_max_length, + ), + steps.to_device, + steps.normalize, + ] + output_steps = [ + steps.unnormalize, + steps.to_cpu, + ] + return make_policy_processor_pipelines(input_steps=input_steps, output_steps=output_steps) diff --git a/videos/ep00_task0_succ.mp4 b/videos/ep00_task0_succ.mp4 new file mode 100644 index 0000000000000000000000000000000000000000..1eba741f70876d95c5830c64edf78144f9a01487 --- /dev/null +++ b/videos/ep00_task0_succ.mp4 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8b7354454803adf0939cb3f0d5b01bb007f2d3018f0553701c268d4bcc7bd0a7 +size 342396 diff --git a/videos/ep01_task1_succ.mp4 b/videos/ep01_task1_succ.mp4 new file mode 100644 index 0000000000000000000000000000000000000000..10bec621644e334d71e33658f099ce00e11bb760 --- /dev/null +++ b/videos/ep01_task1_succ.mp4 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a01e333b9a1d95cd4b82ac17f3c4ac608d660a9d294fb081e2404c4c377cac55 +size 200616 diff --git a/videos/ep02_task2_succ.mp4 b/videos/ep02_task2_succ.mp4 new file mode 100644 index 0000000000000000000000000000000000000000..29df6733616170e1e20a6f94c3c7795ba8cad429 --- /dev/null +++ b/videos/ep02_task2_succ.mp4 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:4effb03d7d6648f480a19a6d540ad77a778d1b3d3d96e43bdb8b26388fd5d483 +size 221834 diff --git a/videos/ep03_task3_fail.mp4 b/videos/ep03_task3_fail.mp4 new file mode 100644 index 0000000000000000000000000000000000000000..b0a0ab240fb91e5dea9f056b7f2bd260e545e5f0 --- /dev/null +++ b/videos/ep03_task3_fail.mp4 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:16c5794895726fa74b7ccda2c8278e0444db72ab0c616df5a572962c1bd204b1 +size 1241690 diff --git a/videos/ep04_task4_succ.mp4 b/videos/ep04_task4_succ.mp4 new file mode 100644 index 0000000000000000000000000000000000000000..640613e36cb2466ddeceb5af56134b40e1f33deb --- /dev/null +++ b/videos/ep04_task4_succ.mp4 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d96507ac413852cd2f73c68fe5219de984c9bcc92cd3c8ee336d75b759e85841 +size 306058 diff --git a/videos/ep05_task5_fail.mp4 b/videos/ep05_task5_fail.mp4 new file mode 100644 index 0000000000000000000000000000000000000000..0802a4ef8e052027223b3e7162620091fa1fd46b --- /dev/null +++ b/videos/ep05_task5_fail.mp4 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a7965732e8e820938cd35e6dec3ba1cc94d3f75632abe7a560d420ba4cbce8ce +size 1081202