tinyvla_b200: B200 streaming training on nvidia/physical-ai (745 samples/s e2e, 7x ViT fast path, hub episode streaming, full instructions)

tinyvla_b200/B200.md

@@ -0,0 +1,160 @@
+# B200: дообучение tv2_C_scaled на nvidia/physical-ai стримингом с хаба
+
+Полная инструкция запуска. Все числа ниже — **измеренные** на реальной B200-ноде
+RunPod 28.08.2026, не оценки.
+
+## Итог в одну строку
+
+**745 семплов/с end-to-end** (стрим с хаба + обучение): **1M кадров = 0.37 ч = $2.5**;
+бюджет C-scaled (60k шагов × batch 256 = 15.4M семплов) = **5.7 ч ≈ $39**;
+эпоха всего пула (~86M семплов @ 10 Гц) = **32 ч ≈ $220**.
+
+## Что в пуле
+
+`make_stream_specs.py` находит в коллекции **163 LeRobot-v2.x датасета, 214.7M
+исходных кадров**: GR00T-X-Embodiment-Sim (127: GR1 arms/waist, single/bimanual
+panda, sim_behavior R1-Pro), GR00T-Teleop-Sim, GR00T-Tuned-Tasks, Teleop-G1,
+AppleToPlate. Схемы разные (state 32–256, action 12–44, 1–3 камеры, fps 20–30),
+всё гармонизируется паддингом + `action_dim_mask`, как в TinyVLA-2.
+
+НЕ входят (формат v3, нужен `build_shards.py`, оба уже были в претрейне C-scaled):
+`BridgeData2_LeRobot_v3`, `LIBERO_LeRobot_v3`.
+
+## Почему стриминг, а не скачивание
+
+Пул — 1.9 ТБ видео; шарды были бы ~1.3 ТБ. v2-формат хранит **один mp4 + один
+parquet на эпизод** (2–10 МБ): воркер скачивает случайный эпизод в память,
+последовательно декодирует (ноль seek'ов), выдаёт ~150–500 семплов, берёт
+следующий. Диск не используется вообще. Замерено: стрим-смесь держит 1019
+семплов/с на 16 воркерах — модель (745/с) не голодает. ВАЖНО: воркеров именно
+16, на 32–64 пропускная способность ПАДАЕТ (677/584) — конкуренция HTTP.
+
+## Обязательные оптимизации (без них B200 бессмысленна)
+
+Всё уже в этом репозитории; числа — вклад каждой (B200, batch 256, bf16):
+
+| что | было → стало | файл |
+|---|---|---|
+| **fla-core + causal-conv1d** (линейное внимание Qwen3.5; без них torch-фоллбэк) | 63 → 236 с/с (3.7×) | pip install |
+| **дедуп ViT**: cam0 через башню 1 раз/шаг вместо 2 (fwd через `image_embeds` в LM) | +23–30% | `modeling_tinyvla.py`, `modules/semantic.py` |
+| **плотный fixed-grid путь ViT**: HF гоняет B отдельных SDPA на слой (3072 вызова/проход) + пересобирает rotary питоном; у нас сетка всегда [1,16,16] → батчевый SDPA, константы кэшированы, torch.compile | encode_image 182 → 26 мс (**7×**), шаг 398 → **738 с/с** | `modules/semantic.py` (`_encode_image_fixed`) |
+| SDPA в эксперте (имена параметров сохранены — чекпойнт грузится) | в составе | `modules/expert.py` (`_MHA`) |
+| M-RoPE position_ids при `inputs_embeds` (фикс бага stale `rope_deltas`) | корректность | `modules/semantic.py` |
+| fused AdamW, bf16 autocast, TF32, токен-кэш инструкций | в составе | `train_fast.py` |
+
+Эквивалентность проверена: fixed-grid ViT **бит-в-бит** совпадает с HF в fp32;
+SDPA-эксперт бит-в-бит с nn.MultiheadAttention на общих весах.
+
+## Запуск с нуля
+
+### 1. Нода
+
+RunPod GraphQL (или консоль): `NVIDIA B200`, образ
+`runpod/pytorch:2.8.0-py3.11-cuda12.8.1-cudnn-devel-ubuntu22.04` (torch с sm_100
+из коробки), 1 GPU, ≥28 vCPU, ≥128 ГБ RAM, диск 50 ГБ (данные не хранятся!),
+`ports: "22/tcp"`. $6.79/ч с диском.
+
+Альтернатива по деньгам: H100 SXM $3.29/ч даёт ~0.5–0.55 ч/1M кадров с теми же
+патчами — **на доллар почти то же самое**; B200 берёт скоростью на узле.
+
+### 2. Окружение (5 минут)
+
+```bash
+pip install pyarrow pillow "transformers==5.5.4" accelerate av webdataset \
+    "huggingface_hub[hf_transfer]" pyyaml safetensors requests msgpack fla-core
+pip install causal-conv1d --no-build-isolation   # колесо ставится на sm_100
+hf auth login --token hf_...        # доступ к AlexWortega/tinyvla
+```
+
+Проверка, что быстрые ядра живы (иначе всё в 3.7 раза медленнее):
+```bash
+python -c "import fla, causal_conv1d; print('ok')"
+```
+
+### 3. Код и чекпойнт
+
+```bash
+cd /workspace
+hf download AlexWortega/tinyvla --local-dir tinyvla_repo \
+    --include "tinyvla_b200/*" "tinyvla2/pretrain_C_scaled/*"
+mv tinyvla_repo/tinyvla_b200 tinyvla_h100
+mkdir -p ckpt && ln -s /workspace/tinyvla_repo/tinyvla2/pretrain_C_scaled ckpt/tv2_C_scaled_final
+export PYTHONPATH=/workspace/tinyvla_h100/src
+```
+
+`lerobot` НЕ нужен для бенчей (`scripts/bench_model.py` собирает модель из
+модулей); для полного `train_fast.py` нужен `lerobot==0.6.1` (тянет
+`TinyVLAConfig`/`PreTrainedPolicy`).
+
+### 4. Спеки стрима (уже сгенерированы, можно перегенерить)
+
+```bash
+python scripts/make_stream_specs.py --out configs/physical_ai_stream.yaml
+# 163 датасета, веса = sqrt(frames), embodiment_id по семействам роботов (10-15)
+```
+
+### 5. Смоук → полный запуск
+
+```bash
+# 200 шагов: заменяет оценку измерением за ~3 минуты
+python scripts/train_fast.py --config configs/physical_ai_ft.yaml --profile-steps 200
+
+# полный прогон (60k шагов, ~5.7 ч)
+PYTHONUNBUFFERED=1 nohup python scripts/train_fast.py \
+    --config configs/physical_ai_ft.yaml > train.log 2>&1 &
+```
+
+При старте `load_compatible` напечатает, что перенесено из C-scaled, а что
+переинициализировано из-за смены размерностей — **проверь этот список**:
+ожидаемо перенесены ViT + LM + блоки эксперта, переинициализированы
+`state_proj` (16→256), `expert.action_in`/`action_out` (8→64),
+`embodiment_emb` частично.
+
+## Что честно сказать про это дообучение
+
+- `conditioning: id`, `action_space: native` — у сим-роботов пула нет URDF/EE в
+  state, канонизация C недоступна. Переносится **ствол** C-scaled, не голова.
+- Первый вызов ViT-компиляции на новый batch-size — рекомпиляция ~20–40 с; в
+  stderr сыплются нефатальные dynamo-warning'и на HF-коде — это шум, скорость
+  реальная.
+- Стрим бесконечный, эпоха не определена; шаги считает train-loop.
+  `shuffle_buffer` 4096/воркер разбивает внутриэпизодную корреляцию.
+- fps-конвенции пула проверены (30 и 20 — честные); децимация до 10 Гц через
+  stride, action-chunk 50 = 5 с, как в претрейне.
+- Хочешь r1_pro (state 256d) исключить — убери его блоки из
+  `physical_ai_stream.yaml` и опусти `max_state_dim` до 64.
+
+## Замеры (B200, полная таблица)
+
+| конфигурация | семплов/с | ч на 1M кадров |
+|---|---|---|
+| модель, HF ViT, без fla | 63 | 4.4 |
+| + fla/causal-conv1d | 236–305 | 0.9–1.2 |
+| + дедуп cam0 | 398 | 0.70 |
+| + плотный fixed-grid ViT + compile | **738** | **0.38** |
+| **end-to-end со стримом с хаба** | **745** | **0.37** |
+| стрим-загрузчик отдельно (16 воркеров) | 1019 | — |
+| webdataset из локальных таров (для сравнения) | 1600–2400 | — |
+
+Профиль шага после оптимизаций (batch 256): ViT 2×26 мс, LM fwd+bwd ~остальное,
+эксперт 7 мс. Следующий резерв — LM-путь (fwd 170 мс до оптимизаций): плотный
+фикс применён только к `encode_image`; внутренний вызов башни в невыгодном
+не-дедуп пути остался HF-шным (дедуп это обходит).
+
+## Файлы
+
+```
+scripts/train_fast.py           цикл: shard- или hub-источник, fused AdamW, bf16, compile
+scripts/make_stream_specs.py    генератор спеков стрима по коллекции
+scripts/bench_model.py          бенч модели без lerobot
+scripts/build_shards.py         LeRobot v3 -> parquet-шарды (Bridge/LIBERO/trex)
+scripts/build_shards_v2.py      LeRobot v2 -> parquet-шарды (если нужен офлайн)
+scripts/pack_wds.py             шарды -> self-contained webdataset тары
+src/tinyvla/data/streaming_hub.py  HubEpisodeStream: стрим эпизодов с хаба
+src/tinyvla/data/shards.py         ShardSource (локальные parquet-шарды)
+src/tinyvla/data/wds_mix.py        wds-смесь (локальные тары)
+src/tinyvla/{modeling_tinyvla,modules/*}.py  модель со всеми патчами скорости
+configs/physical_ai_stream.yaml    163 датасета пула
+configs/physical_ai_ft.yaml        конфиг дообучения C-scaled
+patches/*.diff                     те же патчи диффами к tinyvla2/src
+```

tinyvla_b200/README.md

@@ -0,0 +1,188 @@
+# Дообучение C-scaled на trex_dataset, H100: оценка и ускорение
+
+Оценка бюджета `tv2_C_scaled` (60k шагов) на [`zekaiwang/trex_dataset`](https://huggingface.co/datasets/zekaiwang/trex_dataset)
+одной H100, плюс код, без которого эта оценка недостижима.
+
+**Короткий ответ: 2.9–4.4 ч чистого обучения (~$8–12 на H100 SXM) плюс ~1 ч одноразовой
+подготовки данных. Если пайплайн данных не менять — 8–40 ч при простаивающем GPU.**
+
+---
+
+## 1. Что за датасет (проверено по метаданным, не по карточке)
+
+| | |
+|---|---|
+| эпизодов / кадров | 5 464 / **5 473 459** (сходится с `info.json`) |
+| частота | **30 Гц, и это правда** — `length/(to_ts−from_ts)` = 30.0000 на всех эпизодах |
+| длина эпизода | медиана 887 кадров (29.6 с), максимум 6744 (225 с) |
+| задач | 5370 уникальных инструкций |
+| `observation.state` | **58** float: 2 × (7 DoF рука + 22 DoF кисть Sharpa) |
+| `action` | **58** float, целевые углы суставов |
+| видеопотоков | **23**: 3 RGB 360×640 + 20 тактильных |
+| размер | **1.53 ТБ**, из них 383 ГБ — 3 RGB-камеры |
+
+Урок LIBERO (объявленный fps ≠ реальный) здесь **не повторяется** — проверил первым
+делом, `build_shards.py` проверяет это же и падает при расхождении.
+
+Эпизоды внутри каждого mp4 идут подряд, по возрастанию времени, без разрывов, и число
+кадров точно совпадает с таймштампами. Поэтому последовательный проход по видео работает.
+
+## 2. Три блокера, которые надо решить до запуска
+
+**2.1. 58 измерений против `max_action_dim: 8`.** `CanonicalSource._pad`:
+
+```python
+def _pad(x, dim):
+    if x.shape[-1] >= dim:
+        return x[..., :dim]     # <- МОЛЧА обрезает 58 -> 8
+```
+
+С конфигом C-scaled вы обучитесь на левой руке плюс одном суставе кисти и не узнаете об
+этом: лосс будет падать, метрики будут выглядеть нормально. Поднято до 64 в
+`configs/trex_ft.yaml`; `ShardSource._pad` вместо обрезки бросает исключение.
+
+Побочный эффект: `state_proj`, `expert.action_in`, `expert.action_out` меняют форму.
+`load_state_dict(strict=False)` при смене формы **не пропускает тензор, а падает**, поэтому
+`train_fast.py::load_compatible` фильтрует по форме и печатает каждый переинициализированный
+тензор.
+
+**2.2. Канонические действия недоступны.** `action_space: canonical` — дельты EE в базовой
+системе координат; их строит `build_ee_sidecar.py` через FK по URDF (для SO-101 — placo) или
+читает готовую EE-позу из состояния. Для Vega-1 + Sharpa в репозитории нет URDF, а
+состояние — чистые углы суставов. Значит `native`. C-scaled обучался в каноническом
+пространстве, поэтому переносится ствол (ViT + LM + трансформер эксперта), но не семантика
+выхода. Это не «дообучение C-scaled», это «новая голова действий на его стволе» — стоит
+называть вещи своими именами при сравнении с 0.81/0.88.
+
+**2.3. Архитектура выбрасывает то, ради чего датасет собран.** TinyVLA берёт cam0 + cam1;
+регексы `CanonicalSource` дают `head_left` + `left_wrist`. `right_wrist` и все 20 тактильных
+потоков (это **тактильный бимануальный** датасет) не используются никак. `tactile_force`
+(60 чисел) тоже некуда положить — вход состояния уже занят 58 углами. Если тактильность
+нужна, это отдельный входной путь, а не вопрос конфига.
+
+## 3. Почему сейчас медленно
+
+Обучение берёт **один случайный кадр на семпл**. В LeRobot v3 это h264-seek внутрь mp4 на
+~195 МБ, где лежит ~10 склеенных эпизодов — на к��ждую камеру, на каждый семпл. Это ровно то,
+что записано в `CLAUDE.md`: «video decode is the training bottleneck». Никакая H100 это не
+чинит: GPU просто ждёт.
+
+| стоимость случайного кадра | семплов/с (12 воркеров, 2 камеры) | 60k шагов |
+|---|---|---|
+| 20 мс | 300 | 8.0 ч |
+| 50 мс | 120 | **20.0 ч** |
+| 100 мс | 60 | 40.0 ч |
+
+## 4. Что сделано
+
+### 4.1. Шарды вместо случайного seek (главное)
+
+`scripts/build_shards.py` — одноразовый проход: декодирует каждое видео **последовательно**
+(дешёвое направление), ресайзит до 256×256 ровно так же, как `CanonicalSource._img`
+(без сохранения пропорций — совпадение с обучающей конвенцией важнее геометрии), кодирует в
+JPEG и пишет в parquet-шарды. Умеет `--part i --num-parts N` для распараллеливания по
+машинам и удаляет mp4 после обработки, так что диск не растёт.
+
+`src/tinyvla/data/shards.py` — `ShardSource`: `IterableDataset`, последовательное чтение
+row-групп + shuffle-буфер, раздача шардов по воркерам и рангам. Выдаёт **ровно ту же схему**,
+что `CanonicalSource`, поэтому дальше по коду ничего не меняется. State/action целиком
+живут в RAM (58 float × 2 × 1.8M кадров < 1 ГБ), так что случайных обращений в горячем пути
+не остаётся вовсе.
+
+Стоимость: 47 ГБ шардов при 10 Гц, ~1 ч на подготовку (скачать 245 ГБ 20 мин при 200 МБ/с +
+последовательный декод 10.9M кадров ~16 мин на 16 ядрах).
+
+Дальше пайплайну нужно 1100–1650 JPEG-декодов 256×256 в секунду — это **меньше одного ядра**.
+Данные перестают быть узким местом.
+
+### 4.2. Один проход vision tower вместо двух
+
+`_encode_spatial` гоняет ViT по cam0, а затем `_semantic_latent` вызывает VLM, который
+внутри гоняет ViT **по тому же cam0 ещё раз**. Это 45 из 180 GFLOPs forward — четверть,
+выбрасываемая каждый шаг. Патч считает фичи один раз и передаёт их в LM через
+`image_embeds` (тот же `masked_scatter`, который уже был написан для `morph_to_slow`).
+При включённой staleness-аугментации кадр латента другой, и дедупликация корректно
+отключается. — `patches/modeling_tinyvla.py.diff`, `patches/modules_semantic.py.diff`
+
+### 4.3. SDPA в эксперте
+
+`nn.MultiheadAttention` → `F.scaled_dot_product_attention` с сохранением имён параметров
+(`in_proj_weight`/`in_proj_bias`/`out_proj.*`), поэтому чекпойнт грузится без изменений.
+Проверено численно: на общем `state_dict` выход совпадает **бит в бит** (max|Δ| = 0.00e+00)
+для self-attn, cross-attn и cross-attn с маской. — `patches/modules_expert.py.diff`
+
+### 4.4. Остальное в `scripts/train_fast.py`
+
+Предварительная токенизация всех 5370 инструкций один раз вместо токенизации каждого батча;
+fused AdamW; bf16 autocast + TF32; `grad_accum: 1` при batch 144 (80 ГБ вместо 48);
+`torch.compile` на эксперте (12 блоков из мелких операций — там компиляция даёт больше
+всего); лог с samples/s, оценкой MFU и ETA.
+
+## 5. Оценка
+
+Разложение forward из `ablation.md` §6: ViT 45 GFLOPs на камеру, 12 слоёв LM 46, эксперт 44.
+Backward ≈ 2× forward для того, что учится; замороженный ViT — только forward.
+
+| конфигурация | GFLOPs/семпл (fwd+bwd) |
+|---|---|
+| ViT заморожен, дедуп cam0 | **360** |
+| ViT заморожен, как сейчас | 405 |
+| ViT разморожен, дедуп cam0 | 540 |
+
+Бюджет: 60 000 шагов × eff batch 144 = **8.64M семплов** = 3.11 EFLOPs (при 360).
+На trex при 10 Гц это 1.82M кадров → **4.7 эпохи**.
+
+H100 SXM, пик bf16 989 TFLOPS:
+
+| конфигурация | MFU 15% | 20% | 25% | 30% | 40% |
+|---|---|---|---|---|---|
+| **ViT заморожен, дедуп** | 5.8 ч | 4.4 ч | 3.5 ч | **2.9 ч** | 2.2 ч |
+| ViT замо��ожен, как сейчас | 6.6 ч | 4.9 ч | 3.9 ч | 3.3 ч | 2.5 ч |
+| ViT разморожен, дедуп | 8.7 ч | 6.6 ч | 5.2 ч | 4.4 ч | 3.3 ч |
+
+**Рабочая вилка: 2.9–4.4 ч** (MFU 20–30% — реалистично для модели из ViT + короткого LM +
+мелкого эксперта после `torch.compile` и SDPA).
+
+Перекрёстная проверка: тот же расчёт для A6000, на котором это и обучали, даёт 31–42 ч
+чистого счёта без дедупликации. Отношение пиков H100/A6000 = 6.4×, оценки согласуются.
+Что исходный запуск шёл заметно дольше — ожидаемо, он был data-bound.
+
+Цены RunPod (взяты из API вашим ключом, аккаунт `alexdragannikolich@gmail.com`, баланс $240.62):
+
+| GPU | $/ч | обучение 2.9–4.4 ч | + подготовка |
+|---|---|---|---|
+| **H100 NVL** | 2.59 | $8–11 | +$3 |
+| H100 SXM | 2.69 | $8–12 | +$3 |
+| H200 SXM | 3.59 | $10–16 | +$4 |
+| B200 | 5.98 | $17–26 | +$6 |
+
+## 6. Порядок запуска
+
+```bash
+# 1. Шарды (можно на дешёвом CPU-поде; 4 части параллельно)
+for p in 0 1 2 3; do
+  python scripts/build_shards.py --repo-id zekaiwang/trex_dataset \
+    --out /data/trex_shards --hz 10 --image-size 256 \
+    --part $p --num-parts 4 &
+done; wait
+
+# 2. Смоук: 200 шагов, померить samples/s вместо того чтобы верить оценке
+python scripts/train_fast.py --config configs/trex_ft.yaml --profile-steps 200
+
+# 3. Полный прогон
+python scripts/train_fast.py --config configs/trex_ft.yaml
+```
+
+Шаг 2 — важный: он превращает вилку 2.9–4.4 ч в измеренное число за пять минут.
+
+## 7. Чего я не делал
+
+- **Ничего не запускал на GPU** — здесь нет ни CUDA, ни `lerobot`. Проверено: все файлы
+  компилируются; `ShardSource` прогнан end-to-end на синтетических шардах (схема, формы,
+  паддинг хвоста эпизода, отказ от усечения 58→8); замена attention сверена численно с
+  `nn.MultiheadAttention`. Патч дедупликации ViT проверен только чтением и компиляцией.
+- **Не трогал MFU-предположение.** 20–30% — оценка, а не замер. Шаг 2 выше её заменяет.
+- **Дескриптор морфологии для Vega-1 не сверен со спекой** — значения в
+  `patches/descriptors_dexmate.yaml` правдоподобны, но помечены как требующие проверки.
+  Модель читает дескриптор зеро-шот, так что ошибка стоит качества.
+- **Не делал ничего с тактильностью** — см. 2.3, это отдельная работа.

tinyvla_b200/configs/physical_ai_ft.yaml

@@ -0,0 +1,47 @@
+# Дообучение tv2_C_scaled на LeRobot-части коллекции nvidia/physical-ai,
+# СТРИМИНГОМ с хаба (никаких локальных данных). B200 x1.
+#
+# Что переносится из C-scaled: ViT, 12 слоёв LM, трансформер эксперта.
+# Что переинициализируется (train_fast печатает список): state_proj (16->256),
+# expert.action_in / action_out (8->64) — размерности пула больше.
+# conditioning=id: у 163 сим-датасетов нет ни URDF, ни осмысленных дескрипторов
+# морфологии на семейство; id по семействам роботов (6 семейств, ids 10-15).
+# action_space=native: канонизация требует FK/EE — для этих роботов не построена.
+
+output_dir: /workspace/outputs/physical_ai_ft
+wandb: null
+seed: 42
+resume_from: /workspace/ckpt/tv2_C_scaled_final
+
+source:
+  type: hub
+  specs: /workspace/tinyvla_h100/configs/physical_ai_stream.yaml
+
+policy:
+  chunk_size: 50            # 5 с @ 10 Гц
+  n_action_steps: 50
+  image_size: 256
+  max_state_dim: 256        # sim_behavior_r1_pro
+  max_action_dim: 64        # реальный максимум пула 44, запас до 64
+  num_embodiments: 16
+  conditioning: id
+  action_space: native
+  freeze_lm: false          # главный рычаг (v1: 60%->86%)
+  freeze_vision_encoder: true
+
+backbone_lr_mult: 0.1
+batch_size: 256             # B200 183 ГБ; замерено 46.5 ГБ на 256
+grad_accum: 1
+num_workers: 16             # НЕ 32: конкуренция HTTP-потоков душит стрим (замерено)
+prefetch_factor: 4
+shuffle_buffer: 4096        # на воркер; интерливинг эпизодов
+lr: 1.0e-4
+warmup_steps: 1000
+steps: 60000                # 60k x 256 = 15.4M семплов = ~0.18 эпохи пула
+grad_clip: 10.0
+log_freq: 100
+save_freq: 5000
+compile: true
+
+gflops_per_sample: 360.0
+peak_flops: 2250.0e12       # B200 bf16 dense

tinyvla_b200/configs/physical_ai_stream.yaml

@@ -0,0 +1,820 @@
+# Autogenerated by make_stream_specs.py — streaming mixture over the
+# LeRobot-v2.x part of the nvidia/physical-ai collection.
+# 163 datasets, 214,725,127 source frames.
+# max_state_dim needed: 256; max_action_dim needed: 44.
+datasets:
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "sim_behavior_r1_pro.task-0026_assembling_gift_baskets"
+    weight: 2283.0   # sqrt(5,212,011 fr); fps=30.0 state=256 act=23 cams=3
+    embodiment_id: 14
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "sim_behavior_r1_pro.task-0023_boxing_books_up_for_storage"
+    weight: 2201.2   # sqrt(4,845,500 fr); fps=30.0 state=256 act=23 cams=3
+    embodiment_id: 14
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "gr1_arms_waist.WineToCabinet"
+    weight: 2181.7   # sqrt(4,759,865 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "sim_behavior_r1_pro.task-0048_canning_food"
+    weight: 2143.6   # sqrt(4,595,098 fr); fps=30.0 state=256 act=23 cams=3
+    embodiment_id: 14
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "gr1_arms_waist.PotatoToMicrowave"
+    weight: 2106.4   # sqrt(4,436,785 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "sim_behavior_r1_pro.task-0029_clean_up_your_desk"
+    weight: 2070.2   # sqrt(4,285,704 fr); fps=30.0 state=256 act=23 cams=3
+    embodiment_id: 14
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "sim_behavior_r1_pro.task-0024_storing_food"
+    weight: 1993.4   # sqrt(3,973,824 fr); fps=30.0 state=256 act=23 cams=3
+    embodiment_id: 14
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "sim_behavior_r1_pro.task-0013_loading_the_car"
+    weight: 1960.9   # sqrt(3,845,274 fr); fps=30.0 state=256 act=23 cams=3
+    embodiment_id: 14
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "sim_behavior_r1_pro.task-0049_make_pizza"
+    weight: 1958.9   # sqrt(3,837,385 fr); fps=30.0 state=256 act=23 cams=3
+    embodiment_id: 14
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "sim_behavior_r1_pro.task-0021_collecting_childrens_toys"
+    weight: 1958.9   # sqrt(3,837,265 fr); fps=30.0 state=256 act=23 cams=3
+    embodiment_id: 14
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "gr1_arms_waist.PlaceMilkToMicrowave"
+    weight: 1957.9   # sqrt(3,833,290 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "sim_behavior_r1_pro.task-0007_picking_up_toys"
+    weight: 1943.7   # sqrt(3,778,110 fr); fps=30.0 state=256 act=23 cams=3
+    embodiment_id: 14
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "gr1_arms_waist.PlaceBottleToCabinet"
+    weight: 1878.1   # sqrt(3,527,122 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "gr1_arms_waist.CanToDrawer"
+    weight: 1791.8   # sqrt(3,210,426 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "sim_behavior_r1_pro.task-0027_sorting_household_items"
+    weight: 1778.1   # sqrt(3,161,551 fr); fps=30.0 state=256 act=23 cams=3
+    embodiment_id: 14
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "sim_behavior_r1_pro.task-0028_getting_organized_for_work"
+    weight: 1770.4   # sqrt(3,134,202 fr); fps=30.0 state=256 act=23 cams=3
+    embodiment_id: 14
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "sim_behavior_r1_pro.task-0043_slicing_vegetables"
+    weight: 1723.1   # sqrt(2,968,988 fr); fps=30.0 state=256 act=23 cams=3
+    embodiment_id: 14
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "sim_behavior_r1_pro.task-0016_moving_boxes_to_storage"
+    weight: 1708.6   # sqrt(2,919,245 fr); fps=30.0 state=256 act=23 cams=3
+    embodiment_id: 14
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "sim_behavior_r1_pro.task-0014_carrying_in_groceries"
+    weight: 1689.7   # sqrt(2,854,961 fr); fps=30.0 state=256 act=23 cams=3
+    embodiment_id: 14
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "sim_behavior_r1_pro.task-0041_cook_cabbage"
+    weight: 1680.4   # sqrt(2,823,597 fr); fps=30.0 state=256 act=23 cams=3
+    embodiment_id: 14
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "sim_behavior_r1_pro.task-0002_putting_away_Halloween_decorations"
+    weight: 1663.3   # sqrt(2,766,429 fr); fps=30.0 state=256 act=23 cams=3
+    embodiment_id: 14
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "sim_behavior_r1_pro.task-0009_putting_up_Christmas_decorations_inside"
+    weight: 1656.4   # sqrt(2,743,752 fr); fps=30.0 state=256 act=23 cams=3
+    embodiment_id: 14
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "sim_behavior_r1_pro.task-0003_cleaning_up_plates_and_food"
+    weight: 1655.1   # sqrt(2,739,238 fr); fps=30.0 state=256 act=23 cams=3
+    embodiment_id: 14
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "sim_behavior_r1_pro.task-0015_bringing_in_wood"
+    weight: 1645.3   # sqrt(2,707,131 fr); fps=30.0 state=256 act=23 cams=3
+    embodiment_id: 14
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "gr1_arms_waist.CupToDrawer"
+    weight: 1632.4   # sqrt(2,664,870 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "sim_behavior_r1_pro.task-0025_clearing_food_from_table_into_fridge"
+    weight: 1616.7   # sqrt(2,613,632 fr); fps=30.0 state=256 act=23 cams=3
+    embodiment_id: 14
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "sim_behavior_r1_pro.task-0047_freeze_pies"
+    weight: 1578.3   # sqrt(2,491,046 fr); fps=30.0 state=256 act=23 cams=3
+    embodiment_id: 14
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "sim_behavior_r1_pro.task-0036_clean_a_patio"
+    weight: 1553.8   # sqrt(2,414,162 fr); fps=30.0 state=256 act=23 cams=3
+    embodiment_id: 14
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "gr1_arms_waist.CuttingboardToBasket"
+    weight: 1553.0   # sqrt(2,411,674 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "gr1_arms_waist.PlacematToBasket"
+    weight: 1552.8   # sqrt(2,411,290 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "sim_behavior_r1_pro.task-0020_sorting_vegetables"
+    weight: 1543.0   # sqrt(2,380,740 fr); fps=30.0 state=256 act=23 cams=3
+    embodiment_id: 14
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "sim_behavior_r1_pro.task-0004_can_meat"
+    weight: 1539.3   # sqrt(2,369,415 fr); fps=30.0 state=256 act=23 cams=3
+    embodiment_id: 14
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "gr1_arms_waist.TrayToTieredShelf"
+    weight: 1524.6   # sqrt(2,324,367 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "gr1_arms_waist.PlacematToTieredShelf"
+    weight: 1522.5   # sqrt(2,318,114 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "sim_behavior_r1_pro.task-0033_wash_dog_toys"
+    weight: 1498.2   # sqrt(2,244,560 fr); fps=30.0 state=256 act=23 cams=3
+    embodiment_id: 14
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "gr1_arms_waist.TrayToTieredBasket"
+    weight: 1495.2   # sqrt(2,235,593 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "gr1_arms_waist.CuttingboardToTieredBasket"
+    weight: 1488.5   # sqrt(2,215,695 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "sim_behavior_r1_pro.task-0018_tidying_bedroom"
+    weight: 1485.8   # sqrt(2,207,489 fr); fps=30.0 state=256 act=23 cams=3
+    embodiment_id: 14
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "sim_behavior_r1_pro.task-0011_putting_dishes_away_after_cleaning"
+    weight: 1480.1   # sqrt(2,190,686 fr); fps=30.0 state=256 act=23 cams=3
+    embodiment_id: 14
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "sim_behavior_r1_pro.task-0044_chopping_wood"
+    weight: 1466.4   # sqrt(2,150,397 fr); fps=30.0 state=256 act=23 cams=3
+    embodiment_id: 14
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "sim_behavior_r1_pro.task-0019_outfit_a_basic_toolbox"
+    weight: 1458.6   # sqrt(2,127,594 fr); fps=30.0 state=256 act=23 cams=3
+    embodiment_id: 14
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "gr1_arms_waist.CuttingboardToPan"
+    weight: 1454.0   # sqrt(2,114,010 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "gr1_arms_waist.PlateToPan"
+    weight: 1453.0   # sqrt(2,111,187 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "gr1_arms_waist.PlacematToBowl"
+    weight: 1443.4   # sqrt(2,083,466 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "gr1_arms_waist.PlateToBowl"
+    weight: 1442.7   # sqrt(2,081,304 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "sim_behavior_r1_pro.task-0005_setting_mousetraps"
+    weight: 1426.3   # sqrt(2,034,344 fr); fps=30.0 state=256 act=23 cams=3
+    embodiment_id: 14
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "gr1_arms_waist.TrayToCardboardBox"
+    weight: 1424.1   # sqrt(2,027,980 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "gr1_arms_waist.PlateToCardboardBox"
+    weight: 1417.1   # sqrt(2,008,179 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "gr1_arms_waist.CuttingboardToCardboardBox"
+    weight: 1414.9   # sqrt(2,002,060 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "gr1_arms_waist.TrayToPlate"
+    weight: 1404.1   # sqrt(1,971,615 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "gr1_arms_waist.PlacematToPlate"
+    weight: 1401.1   # sqrt(1,963,149 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "gr1_arms_waist.PlateToPlate"
+    weight: 1401.1   # sqrt(1,963,078 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "gr1_arms_waist.TrayToPot"
+    weight: 1396.6   # sqrt(1,950,435 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "gr1_arms_waist.CuttingboardToPot"
+    weight: 1394.7   # sqrt(1,945,224 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "sim_behavior_r1_pro.task-0017_bringing_water"
+    weight: 1373.9   # sqrt(1,887,709 fr); fps=30.0 state=256 act=23 cams=3
+    embodiment_id: 14
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "sim_behavior_r1_pro.task-0045_cook_hot_dogs"
+    weight: 1352.4   # sqrt(1,828,971 fr); fps=30.0 state=256 act=23 cams=3
+    embodiment_id: 14
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "sim_behavior_r1_pro.task-0030_setting_the_fire"
+    weight: 1350.4   # sqrt(1,823,626 fr); fps=30.0 state=256 act=23 cams=3
+    embodiment_id: 14
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "sim_behavior_r1_pro.task-0008_rearranging_kitchen_furniture"
+    weight: 1337.4   # sqrt(1,788,603 fr); fps=30.0 state=256 act=23 cams=3
+    embodiment_id: 14
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "sim_behavior_r1_pro.task-0032_wash_a_baseball_cap"
+    weight: 1292.2   # sqrt(1,669,882 fr); fps=30.0 state=256 act=23 cams=3
+    embodiment_id: 14
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "sim_behavior_r1_pro.task-0039_spraying_fruit_trees"
+    weight: 1292.0   # sqrt(1,669,174 fr); fps=30.0 state=256 act=23 cams=3
+    embodiment_id: 14
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "sim_behavior_r1_pro.task-0012_preparing_lunch_box"
+    weight: 1284.2   # sqrt(1,649,060 fr); fps=30.0 state=256 act=23 cams=3
+    embodiment_id: 14
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "sim_behavior_r1_pro.task-0031_clean_boxing_gloves"
+    weight: 1283.4   # sqrt(1,647,072 fr); fps=30.0 state=256 act=23 cams=3
+    embodiment_id: 14
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "sim_behavior_r1_pro.task-0022_putting_shoes_on_rack"
+    weight: 1240.4   # sqrt(1,538,490 fr); fps=30.0 state=256 act=23 cams=3
+    embodiment_id: 14
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "sim_behavior_r1_pro.task-0046_cook_bacon"
+    weight: 1239.3   # sqrt(1,535,909 fr); fps=30.0 state=256 act=23 cams=3
+    embodiment_id: 14
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "sim_behavior_r1_pro.task-0006_hiding_Easter_eggs"
+    weight: 1234.5   # sqrt(1,523,945 fr); fps=30.0 state=256 act=23 cams=3
+    embodiment_id: 14
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "single_panda_gripper.PnPCounterToSink"
+    weight: 1149.3   # sqrt(1,320,967 fr); fps=20.0 state=53 act=12 cams=3
+    embodiment_id: 13
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "sim_behavior_r1_pro.task-0038_spraying_for_bugs"
+    weight: 1138.4   # sqrt(1,295,895 fr); fps=30.0 state=256 act=23 cams=3
+    embodiment_id: 14
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "single_panda_gripper.PnPCounterToMicrowave"
+    weight: 1131.9   # sqrt(1,281,129 fr); fps=20.0 state=53 act=12 cams=3
+    embodiment_id: 13
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "sim_behavior_r1_pro.task-0042_chop_an_onion"
+    weight: 1131.4   # sqrt(1,279,960 fr); fps=30.0 state=256 act=23 cams=3
+    embodiment_id: 14
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "sim_behavior_r1_pro.task-0010_set_up_a_coffee_station_in_your_kitchen"
+    weight: 1119.5   # sqrt(1,253,243 fr); fps=30.0 state=256 act=23 cams=3
+    embodiment_id: 14
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "single_panda_gripper.CloseDoubleDoor"
+    weight: 1099.8   # sqrt(1,209,559 fr); fps=20.0 state=53 act=12 cams=3
+    embodiment_id: 13
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "single_panda_gripper.OpenDoubleDoor"
+    weight: 1038.1   # sqrt(1,077,708 fr); fps=20.0 state=53 act=12 cams=3
+    embodiment_id: 13
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "sim_behavior_r1_pro.task-0037_clean_a_trumpet"
+    weight: 1028.8   # sqrt(1,058,409 fr); fps=30.0 state=256 act=23 cams=3
+    embodiment_id: 14
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "sim_behavior_r1_pro.task-0001_picking_up_trash"
+    weight: 1026.4   # sqrt(1,053,550 fr); fps=30.0 state=256 act=23 cams=3
+    embodiment_id: 14
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "single_panda_gripper.PnPCabToCounter"
+    weight: 1007.8   # sqrt(1,015,757 fr); fps=20.0 state=53 act=12 cams=3
+    embodiment_id: 13
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "single_panda_gripper.CoffeeServeMug"
+    weight: 1005.3   # sqrt(1,010,648 fr); fps=20.0 state=53 act=12 cams=3
+    embodiment_id: 13
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "single_panda_gripper.PnPCounterToStove"
+    weight: 978.4   # sqrt(957,229 fr); fps=20.0 state=53 act=12 cams=3
+    embodiment_id: 13
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "single_panda_gripper.PnPMicrowaveToCounter"
+    weight: 971.8   # sqrt(944,476 fr); fps=20.0 state=53 act=12 cams=3
+    embodiment_id: 13
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "single_panda_gripper.CoffeeSetupMug"
+    weight: 913.0   # sqrt(833,482 fr); fps=20.0 state=53 act=12 cams=3
+    embodiment_id: 13
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "single_panda_gripper.PnPSinkToCounter"
+    weight: 906.9   # sqrt(822,477 fr); fps=20.0 state=53 act=12 cams=3
+    embodiment_id: 13
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "single_panda_gripper.PnPStoveToCounter"
+    weight: 903.5   # sqrt(816,366 fr); fps=20.0 state=53 act=12 cams=3
+    embodiment_id: 13
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "single_panda_gripper.OpenSingleDoor"
+    weight: 899.0   # sqrt(808,202 fr); fps=20.0 state=53 act=12 cams=3
+    embodiment_id: 13
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "single_panda_gripper.CloseSingleDoor"
+    weight: 886.5   # sqrt(785,853 fr); fps=20.0 state=53 act=12 cams=3
+    embodiment_id: 13
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "sim_behavior_r1_pro.task-0035_attach_a_camera_to_a_tripod"
+    weight: 884.5   # sqrt(782,315 fr); fps=30.0 state=256 act=23 cams=3
+    embodiment_id: 14
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "single_panda_gripper.PnPCounterToCab"
+    weight: 881.9   # sqrt(777,674 fr); fps=20.0 state=53 act=12 cams=3
+    embodiment_id: 13
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "single_panda_gripper.TurnOnSinkFaucet"
+    weight: 847.9   # sqrt(718,880 fr); fps=20.0 state=53 act=12 cams=3
+    embodiment_id: 13
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "single_panda_gripper.TurnOffMicrowave"
+    weight: 838.8   # sqrt(703,625 fr); fps=20.0 state=53 act=12 cams=3
+    embodiment_id: 13
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "single_panda_gripper.TurnOnMicrowave"
+    weight: 838.1   # sqrt(702,401 fr); fps=20.0 state=53 act=12 cams=3
+    embodiment_id: 13
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "single_panda_gripper.OpenDrawer"
+    weight: 821.2   # sqrt(674,374 fr); fps=20.0 state=53 act=12 cams=3
+    embodiment_id: 13
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "sim_behavior_r1_pro.task-0040_make_microwave_popcorn"
+    weight: 804.7   # sqrt(647,555 fr); fps=30.0 state=256 act=23 cams=3
+    embodiment_id: 14
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "single_panda_gripper.TurnOnStove"
+    weight: 798.5   # sqrt(637,617 fr); fps=20.0 state=53 act=12 cams=3
+    embodiment_id: 13
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "single_panda_gripper.CloseDrawer"
+    weight: 787.0   # sqrt(619,411 fr); fps=20.0 state=53 act=12 cams=3
+    embodiment_id: 13
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "single_panda_gripper.TurnOffSinkFaucet"
+    weight: 773.1   # sqrt(597,645 fr); fps=20.0 state=53 act=12 cams=3
+    embodiment_id: 13
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "single_panda_gripper.TurnOffStove"
+    weight: 738.1   # sqrt(544,815 fr); fps=20.0 state=53 act=12 cams=3
+    embodiment_id: 13
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "bimanual_panda_hand.LiftTray"
+    weight: 707.2   # sqrt(500,131 fr); fps=20.0 state=40 act=24 cams=3
+    embodiment_id: 12
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "sim_behavior_r1_pro.task-0034_hanging_pictures"
+    weight: 691.4   # sqrt(478,034 fr); fps=30.0 state=256 act=23 cams=3
+    embodiment_id: 14
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "gr1_unified.PnPWineToCabinetClose_GR1ArmsAndWaistFourierHands_1000"
+    weight: 687.8   # sqrt(473,076 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-Teleop-Sim
+    prefix: "LeRobot/gr1_unified.PnPWineToCabinetClose"
+    weight: 664.4   # sqrt(441,387 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "gr1_unified.PnPPotatoToMicrowaveClose_GR1ArmsAndWaistFourierHands_1000"
+    weight: 662.9   # sqrt(439,435 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "bimanual_panda_gripper.Transport"
+    weight: 639.8   # sqrt(409,386 fr); fps=20.0 state=32 act=14 cams=3
+    embodiment_id: 11
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "single_panda_gripper.TurnSinkSpout"
+    weight: 627.4   # sqrt(393,622 fr); fps=20.0 state=53 act=12 cams=3
+    embodiment_id: 13
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-Teleop-Sim
+    prefix: "LeRobot/gr1_unified.PnPMilkToMicrowaveClose"
+    weight: 622.3   # sqrt(387,316 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "gr1_unified.PnPMilkToMicrowaveClose_GR1ArmsAndWaistFourierHands_1000"
+    weight: 616.4   # sqrt(379,992 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-Teleop-Sim
+    prefix: "LeRobot/gr1_unified.PnPPotatoToMicrowaveClose"
+    weight: 611.0   # sqrt(373,353 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "single_panda_gripper.CoffeePressButton"
+    weight: 610.3   # sqrt(372,446 fr); fps=20.0 state=53 act=12 cams=3
+    embodiment_id: 13
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "gr1_unified.PnPBottleToCabinetClose_GR1ArmsAndWaistFourierHands_1000"
+    weight: 593.9   # sqrt(352,726 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-Teleop-Sim
+    prefix: "LeRobot/gr1_unified.PnPBottleToCabinetClose"
+    weight: 591.2   # sqrt(349,476 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "sim_behavior_r1_pro.task-0000_turning_on_radio"
+    weight: 570.1   # sqrt(324,964 fr); fps=30.0 state=256 act=23 cams=3
+    embodiment_id: 14
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "gr1_full_upper_body.Coffee"
+    weight: 569.7   # sqrt(324,517 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "gr1_unified.PnPCanToDrawerClose_GR1ArmsAndWaistFourierHands_1000"
+    weight: 565.4   # sqrt(319,656 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-Teleop-Sim
+    prefix: "LeRobot/gr1_unified.PnPCanToDrawerClose"
+    weight: 564.4   # sqrt(318,546 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "gr1_arms_only.CanSort"
+    weight: 562.9   # sqrt(316,878 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-Teleop-Sim
+    prefix: "LeRobot/gr1_unified.PnPCupToDrawerClose"
+    weight: 559.8   # sqrt(313,404 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-GR00T-Tuned-Tasks
+    prefix: "Nut-Pouring-task"
+    weight: 555.3   # sqrt(308,383 fr); fps=20 state=26 act=26 cams=1
+    embodiment_id: 15
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "bimanual_panda_hand.DrawerCleanup"
+    weight: 539.2   # sqrt(290,716 fr); fps=20.0 state=40 act=24 cams=3
+    embodiment_id: 12
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-Teleop-Sim
+    prefix: "LeRobot/gr1_unified.PosttrainPnPNovelFromTrayToTieredshelfSplitA"
+    weight: 523.8   # sqrt(274,326 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "gr1_full_upper_body.Pouring"
+    weight: 517.5   # sqrt(267,780 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "gr1_unified.PnPCupToDrawerClose_GR1ArmsAndWaistFourierHands_1000"
+    weight: 515.8   # sqrt(266,040 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-Teleop-Sim
+    prefix: "LeRobot/gr1_unified.PosttrainPnPNovelFromPlacematToTieredshelfSplitA"
+    weight: 508.8   # sqrt(258,920 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "gr1_unified.PosttrainPnPNovelFromCuttingboardToBasketSplitA_GR1ArmsAndWaistFourierHands_1000"
+    weight: 489.2   # sqrt(239,307 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "gr1_unified.PosttrainPnPNovelFromPlacematToBasketSplitA_GR1ArmsAndWaistFourierHands_1000"
+    weight: 488.4   # sqrt(238,539 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "bimanual_panda_gripper.ThreePieceAssembly"
+    weight: 488.2   # sqrt(238,372 fr); fps=20.0 state=32 act=14 cams=3
+    embodiment_id: 11
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-Teleop-Sim
+    prefix: "LeRobot/gr1_unified.PosttrainPnPNovelFromTrayToTieredbasketSplitA"
+    weight: 484.2   # sqrt(234,452 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "gr1_unified.PosttrainPnPNovelFromTrayToTieredshelfSplitA_GR1ArmsAndWaistFourierHands_1000"
+    weight: 481.7   # sqrt(232,039 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "gr1_unified.PosttrainPnPNovelFromPlacematToTieredshelfSplitA_GR1ArmsAndWaistFourierHands_1000"
+    weight: 480.5   # sqrt(230,861 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "bimanual_panda_hand.BoxCleanup"
+    weight: 480.3   # sqrt(230,689 fr); fps=20.0 state=40 act=24 cams=3
+    embodiment_id: 12
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-Teleop-Sim
+    prefix: "LeRobot/gr1_unified.PosttrainPnPNovelFromPlateToCardboardboxSplitA"
+    weight: 474.9   # sqrt(225,540 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-GR00T-Tuned-Tasks
+    prefix: "Exhaust-Pipe-Sorting-task"
+    weight: 473.1   # sqrt(223,823 fr); fps=20 state=26 act=26 cams=1
+    embodiment_id: 15
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-Teleop-Sim
+    prefix: "LeRobot/gr1_unified.PosttrainPnPNovelFromPlateToPlateSplitA"
+    weight: 472.0   # sqrt(222,740 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "gr1_unified.PosttrainPnPNovelFromTrayToTieredbasketSplitA_GR1ArmsAndWaistFourierHands_1000"
+    weight: 471.9   # sqrt(222,664 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "gr1_unified.PosttrainPnPNovelFromCuttingboardToTieredbasketSplitA_GR1ArmsAndWaistFourierHands_1000"
+    weight: 469.5   # sqrt(220,405 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-Teleop-Sim
+    prefix: "LeRobot/gr1_unified.PosttrainPnPNovelFromTrayToPlateSplitA"
+    weight: 468.1   # sqrt(219,082 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-Teleop-Sim
+    prefix: "LeRobot/gr1_unified.PosttrainPnPNovelFromPlateToPanSplitA"
+    weight: 466.7   # sqrt(217,848 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-Teleop-Sim
+    prefix: "LeRobot/gr1_unified.PosttrainPnPNovelFromCuttingboardToTieredbasketSplitA"
+    weight: 464.7   # sqrt(215,965 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "bimanual_panda_gripper.Threading"
+    weight: 462.1   # sqrt(213,493 fr); fps=20.0 state=32 act=14 cams=3
+    embodiment_id: 11
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-Teleop-Sim
+    prefix: "LeRobot/gr1_unified.PosttrainPnPNovelFromCuttingboardToPotSplitA"
+    weight: 462.0   # sqrt(213,457 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "gr1_unified.PosttrainPnPNovelFromPlateToPanSplitA_GR1ArmsAndWaistFourierHands_1000"
+    weight: 458.2   # sqrt(209,968 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "gr1_unified.PosttrainPnPNovelFromCuttingboardToPanSplitA_GR1ArmsAndWaistFourierHands_1000"
+    weight: 458.1   # sqrt(209,896 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "gr1_unified.PosttrainPnPNovelFromPlacematToBowlSplitA_GR1ArmsAndWaistFourierHands_1000"
+    weight: 454.8   # sqrt(206,826 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "gr1_unified.PosttrainPnPNovelFromPlateToBowlSplitA_GR1ArmsAndWaistFourierHands_1000"
+    weight: 454.6   # sqrt(206,643 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "gr1_unified.PosttrainPnPNovelFromPlateToCardboardboxSplitA_GR1ArmsAndWaistFourierHands_1000"
+    weight: 446.9   # sqrt(199,723 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "gr1_unified.PosttrainPnPNovelFromTrayToCardboardboxSplitA_GR1ArmsAndWaistFourierHands_1000"
+    weight: 446.7   # sqrt(199,505 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-Teleop-Sim
+    prefix: "LeRobot/gr1_unified.PosttrainPnPNovelFromTrayToCardboardboxSplitA"
+    weight: 446.5   # sqrt(199,399 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "gr1_unified.PosttrainPnPNovelFromCuttingboardToCardboardboxSplitA_GR1ArmsAndWaistFourierHands_1000"
+    weight: 446.5   # sqrt(199,398 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "gr1_unified.PosttrainPnPNovelFromTrayToPlateSplitA_GR1ArmsAndWaistFourierHands_1000"
+    weight: 442.4   # sqrt(195,734 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "gr1_unified.PosttrainPnPNovelFromPlateToPlateSplitA_GR1ArmsAndWaistFourierHands_1000"
+    weight: 442.1   # sqrt(195,446 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "gr1_unified.PosttrainPnPNovelFromPlacematToPlateSplitA_GR1ArmsAndWaistFourierHands_1000"
+    weight: 441.3   # sqrt(194,706 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "gr1_unified.PosttrainPnPNovelFromTrayToPotSplitA_GR1ArmsAndWaistFourierHands_1000"
+    weight: 440.6   # sqrt(194,126 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "gr1_unified.PosttrainPnPNovelFromCuttingboardToPotSplitA_GR1ArmsAndWaistFourierHands_1000"
+    weight: 439.7   # sqrt(193,347 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-Teleop-Sim
+    prefix: "LeRobot/gr1_unified.PosttrainPnPNovelFromTrayToPotSplitA"
+    weight: 433.4   # sqrt(187,852 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-Teleop-Sim
+    prefix: "LeRobot/gr1_unified.PosttrainPnPNovelFromCuttingboardToCardboardboxSplitA"
+    weight: 427.1   # sqrt(182,379 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-Teleop-Sim
+    prefix: "LeRobot/gr1_unified.PosttrainPnPNovelFromCuttingboardToBasketSplitA"
+    weight: 421.4   # sqrt(177,569 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/GR00T-N1.7-AppleToPlate
+    prefix: ""
+    weight: 414.3   # sqrt(171,625 fr); fps=30 state=43 act=43 cams=1
+    embodiment_id: 15
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-Teleop-Sim
+    prefix: "LeRobot/gr1_unified.PosttrainPnPNovelFromPlateToBowlSplitA"
+    weight: 414.1   # sqrt(171,460 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-Teleop-Sim
+    prefix: "LeRobot/gr1_unified.PosttrainPnPNovelFromPlacematToPlateSplitA"
+    weight: 412.6   # sqrt(170,207 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-Teleop-Sim
+    prefix: "LeRobot/gr1_unified.PosttrainPnPNovelFromCuttingboardToPanSplitA"
+    weight: 407.8   # sqrt(166,293 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-Teleop-Sim
+    prefix: "LeRobot/gr1_unified.PosttrainPnPNovelFromPlacematToBowlSplitA"
+    weight: 397.9   # sqrt(158,324 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-Teleop-Sim
+    prefix: "LeRobot/gr1_unified.PosttrainPnPNovelFromPlacematToBasketSplitA"
+    weight: 375.5   # sqrt(140,982 fr); fps=20.0 state=44 act=44 cams=1
+    embodiment_id: 10
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+    prefix: "unitree_g1.LMPnPAppleToPlateDC"
+    weight: 279.9   # sqrt(78,369 fr); fps=50.0 state=43 act=43 cams=1
+    embodiment_id: 15
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-Teleop-G1
+    prefix: "g1-pick-starfruit"
+    weight: 186.1   # sqrt(34,640 fr); fps=20 state=43 act=43 cams=1
+    embodiment_id: 15
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-Teleop-G1
+    prefix: "g1-pick-pear"
+    weight: 179.8   # sqrt(32,319 fr); fps=20 state=43 act=43 cams=1
+    embodiment_id: 15
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-Teleop-G1
+    prefix: "g1-pick-apple"
+    weight: 172.8   # sqrt(29,864 fr); fps=20 state=43 act=43 cams=1
+    embodiment_id: 15
+    hz: 10.0
+  - repo_id: nvidia/PhysicalAI-Robotics-GR00T-Teleop-G1
+    prefix: "g1-pick-grapes"
+    weight: 164.2   # sqrt(26,970 fr); fps=20 state=43 act=43 cams=1
+    embodiment_id: 15
+    hz: 10.0

tinyvla_b200/configs/trex_ft.yaml

@@ -0,0 +1,62 @@
+# Дообучение C-scaled на zekaiwang/trex_dataset (Dexmate Vega-1 + Sharpa Wave).
+#
+# ВНИМАНИЕ, три вещи отличаются от всех предыдущих запусков TinyVLA-2:
+#
+# 1. state/action = 58 измерений (2 x (7 DoF рука + 22 DoF кисть)), а не 8/16.
+#    max_action_dim/max_state_dim подняты до 64. Без этого CanonicalSource._pad
+#    МОЛЧА обрежет действие до первых 8 чисел (левая рука + 1 сустав кисти).
+#    Побочный эффект: state_proj, expert.action_in и expert.action_out меняют
+#    форму и переинициализируются -- train_fast.py печатает это явно.
+#
+# 2. action_space: native, а не canonical. Канонические действия -- дельты EE в
+#    базовой системе координат, их строит build_ee_sidecar.py через FK по URDF.
+#    Для Vega-1 + Sharpa в репозитории нет ни URDF, ни EE-позы в состоянии
+#    (58 чисел -- это чистые углы суставов). Значит смысл выходного пространства
+#    меняется, и от C-scaled переносится ствол (ViT, LM, трансформер эксперта),
+#    но не голова действий.
+#
+# 3. Данные читаются из заранее построенных шардов (scripts/build_shards.py),
+#    а не из LeRobotDataset. 1.53 ТБ h264 со случайным seek -- это то, что
+#    держало A6000 на ~10% загрузки.
+
+output_dir: /workspace/outputs/trex_ft
+wandb: null
+seed: 42
+morphology_descriptors: /workspace/tinyvla/configs/morphology/descriptors.yaml
+resume_from: /workspace/ckpt/tv2_C_scaled_final
+
+source:
+  root: /workspace/data/trex_shards
+  embodiment_id: 10          # свободный id: C-scaled использует 0..9
+  morph_key: dexmate_vega1   # см. patches/descriptors_dexmate.yaml
+  robot_prompt: null         # текстовое описание робота ВРЕДИТ (RESULTS.md)
+
+policy:
+  chunk_size: 50             # 50 шагов @ 10 Гц = 5 с, как в претрейне
+  n_action_steps: 50
+  image_size: 256
+  max_state_dim: 64          # 58 -> паддинг до 64
+  max_action_dim: 64
+  num_embodiments: 16
+  conditioning: morph        # единственный канал, который реально работает
+  action_space: native       # см. пункт 2 выше
+  freeze_lm: false           # самый крупный рычаг в обоих поколениях
+  freeze_vision_encoder: true
+
+backbone_lr_mult: 0.1
+batch_size: 144              # 80 ГБ H100: grad_accum больше не нужен
+grad_accum: 1
+num_workers: 12
+prefetch_factor: 6
+shuffle_buffer: 16384
+lr: 1.0e-4
+warmup_steps: 1000
+steps: 60000
+grad_clip: 10.0
+log_freq: 50
+save_freq: 10000
+compile: true
+
+# для строки MFU в логе (ablation.md §6: ViT 45 + LM 46 + expert 44 GFLOPs)
+gflops_per_sample: 360.0
+peak_flops: 989.0e12

tinyvla_b200/patches/descriptors_dexmate.yaml

@@ -0,0 +1,24 @@
+# Дописать в configs/morphology/descriptors.yaml.
+# Порядок полей ОБЯЗАН совпадать с MORPH_FIELDS в modules/embodiment.py.
+#
+# ЗНАЧЕНИЯ НУЖНО СВЕРИТЬ СО СПЕКОЙ Dexmate Vega-1 / Sharpa Wave -- ниже
+# правдоподобные, но не проверенные по документации числа. Дескриптор читается
+# моделью зеро-шот (RESULTS.md: правильный дескриптор бьёт неправильный на
+# 5-45%), так что ошибка здесь стоит качества, а не падения.
+dexmate_vega1:      # бимануальный, 2 x (7-DoF рука + 22-DoF кисть Sharpa)
+  arm_dof: 7             # на руку; 16-полевая схема не выражает бимануальность
+  reach_m: 0.70
+  gripper_width_m: 0.10  # многопалая кисть, не параллельный захват
+  num_cameras: 3         # head_left + 2 запястных (обучаем на 2 из них)
+  is_mobile: 1           # ПРОВЕРИТЬ: есть ли у Vega-1 подвижное основание
+  control_hz: 30
+  joint_lo_mean: -2.6
+  joint_hi_mean: 2.6
+  workspace_x: 0.50
+  workspace_y: 0.60
+  workspace_z: 0.70
+  payload_kg: 3.0
+  ee_type_parallel: -1
+  ee_type_multi: 1       # единственное поле, кодирующее многопалую кисть
+  base_holonomic: -1     # ПРОВЕРИТЬ
+  reserved: 0

tinyvla_b200/patches/modeling_tinyvla.py.diff

@@ -0,0 +1,75 @@
+--- ../tinyvla/tinyvla2/src/tinyvla/modeling_tinyvla.py	2026-08-28 13:41:44
++++ src/tinyvla/modeling_tinyvla.py	2026-08-28 15:59:41
+@@ -121,18 +121,28 @@
+         """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)."""
++    def _encode_spatial(self, batch: dict, feats: dict | None = None) -> torch.Tensor:
++        """Per-camera spatial tokens for the fast path: (B, num_cams*64, d).
++
++        `feats` lets the caller pass vision-tower outputs computed once per step
++        (see `_vision_features`) instead of re-running the tower here.
++        """
+         cond = []
+         for i, key in enumerate(self.image_keys):
+-            tokens = self.semantic.encode_image(self._last_obs(batch[key], 4))  # (B, 64, C)
++            tokens = (feats or {}).get(key)
++            if tokens is None:
++                tokens = self.semantic.encode_image(self._last_obs(batch[key], 4))  # (B, 64, C)
+             proj = self.spatial_proj(tokens) + self.camera_emb.weight[i][None, None, :]
+             if i == 0:
+                 self._cam0_proj = proj  # consumed by the spatial-distillation head
+             cond.append(proj)
+         return torch.cat(cond, dim=1)
+ 
+-    def _semantic_latent(self, batch: dict, image=None) -> torch.Tensor:
++    def _vision_features(self, batch: dict, keys) -> dict:
++        """One vision-tower pass per camera, shared by the fast and slow paths."""
++        return {k: self.semantic.encode_image(self._last_obs(batch[k], 4)) for k in keys}
++
++    def _semantic_latent(self, batch: dict, image=None, image_embeds=None) -> torch.Tensor:
+         """Slow-path latent: (B, K, d), cacheable, conditions the fast path.
+ 
+         vlm_native=True makes this latent RICH — the LM sequence also carries the
+@@ -163,23 +173,35 @@
+             batch[OBS_LANGUAGE_TOKENS],
+             batch[OBS_LANGUAGE_ATTENTION_MASK],
+             morphology=morph,
++            image_embeds=image_embeds,
+         )
+ 
+     def _conditioning(self, batch: dict) -> torch.Tensor:
+         b = batch[self.image_keys[0]].shape[0]
+         device = batch[self.image_keys[0]].device
++
++        # Run the shared Qwen vision tower ONCE per camera per step. Previously
++        # cam0 went through it twice: here for the spatial tokens, and again
++        # inside the LM call in `_semantic_latent`. The tower is ~45 GFLOPs per
++        # camera, i.e. a quarter of the forward pass thrown away every step.
++        feats = None
++        if self.config.use_spatial_tokens:
++            feats = self._vision_features(batch, self.image_keys)
++        elif self.config.use_semantic_latent and not self.config.vlm_native:
++            feats = self._vision_features(batch, self.image_keys[:1])
+ 
+         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
++                if "latent_image" in batch:  # staleness aug: latent from an older frame,
+                     latent = self._semantic_latent(batch, image=batch["latent_image"])
+-                else:
+-                    latent = self._semantic_latent(batch)
++                else:                        # ...so the cached cam0 features do NOT apply
++                    cam0 = None if (feats is None or self.config.vlm_native) else feats[self.image_keys[0]]
++                    latent = self._semantic_latent(batch, image_embeds=cam0)
+             parts.append(latent)
+         if self.config.use_spatial_tokens:
+-            parts.append(self._encode_spatial(batch))
++            parts.append(self._encode_spatial(batch, feats))
+         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, :])

tinyvla_b200/patches/modules_expert.py.diff

@@ -0,0 +1,76 @@
+--- ../tinyvla/tinyvla2/src/tinyvla/modules/expert.py	2026-08-28 13:41:44
++++ src/tinyvla/modules/expert.py	2026-08-28 15:58:45
+@@ -11,6 +11,7 @@
+ 
+ import torch
+ import torch.nn as nn
++import torch.nn.functional as F
+ 
+ 
+ def sincos_time_embedding(t: torch.Tensor, dim: int, min_period: float = 4e-3, max_period: float = 4.0) -> torch.Tensor:
+@@ -23,6 +24,47 @@
+     return torch.cat([angles.sin(), angles.cos()], dim=-1)
+ 
+ 
++class _MHA(nn.Module):
++    """Drop-in replacement for torch's batch-first multi-head attention, routed
++    through F.scaled_dot_product_attention (flash / mem-efficient kernels, and
++    far friendlier to torch.compile than the eager fallback).
++
++    Parameter names are kept identical (in_proj_weight / in_proj_bias /
++    out_proj.weight / out_proj.bias) so existing checkpoints load unchanged.
++    `key_padding_mask` keeps torch semantics: True = ignore that key.
++    """
++
++    def __init__(self, dim: int, n_heads: int):
++        super().__init__()
++        assert dim % n_heads == 0
++        self.dim, self.n_heads, self.head_dim = dim, n_heads, dim // n_heads
++        self.in_proj_weight = nn.Parameter(torch.empty(3 * dim, dim))
++        self.in_proj_bias = nn.Parameter(torch.zeros(3 * dim))
++        self.out_proj = nn.Linear(dim, dim)
++        nn.init.xavier_uniform_(self.in_proj_weight)
++        nn.init.zeros_(self.out_proj.bias)
++
++    def forward(self, query, key, value, key_padding_mask=None, need_weights=False):
++        b, lq, d = query.shape
++        w_q, w_k, w_v = self.in_proj_weight.split(d, dim=0)
++        b_q, b_k, b_v = self.in_proj_bias.split(d, dim=0)
++        q = F.linear(query, w_q, b_q)
++        k = F.linear(key, w_k, b_k)
++        v = F.linear(value, w_v, b_v)
++
++        def split(t):
++            return t.view(b, -1, self.n_heads, self.head_dim).transpose(1, 2)
++
++        q, k, v = split(q), split(k), split(v)
++        attn_mask = None
++        if key_padding_mask is not None:
++            # torch convention: True = ignore. SDPA bool mask: True = attend.
++            attn_mask = (~key_padding_mask)[:, None, None, :]
++        o = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask)
++        o = o.transpose(1, 2).reshape(b, lq, d)
++        return self.out_proj(o), None
++
++
+ class _MLP(nn.Module):
+     def __init__(self, dim: int, ratio: int = 4):
+         super().__init__()
+@@ -40,7 +82,7 @@
+     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.attn = _MHA(dim, n_heads)
+         self.norm2 = nn.LayerNorm(dim)
+         self.mlp = _MLP(dim)
+ 
+@@ -57,7 +99,7 @@
+         super().__init__()
+         self.norm1 = nn.LayerNorm(dim)
+         self.norm_cond = nn.LayerNorm(dim)
+-        self.attn = nn.MultiheadAttention(dim, n_heads, batch_first=True)
++        self.attn = _MHA(dim, n_heads)
+         self.norm2 = nn.LayerNorm(dim)
+         self.mlp = _MLP(dim)
+ 

tinyvla_b200/patches/modules_semantic.py.diff

@@ -0,0 +1,180 @@
+--- ../tinyvla/tinyvla2/src/tinyvla/modules/semantic.py	2026-08-28 13:41:44
++++ src/tinyvla/modules/semantic.py	2026-08-28 17:23:51
+@@ -69,6 +69,7 @@
+ 
+         self.out_proj = nn.Linear(self.hidden_size, out_dim)
+         self.num_readout = num_readout
++        self._fixed_fn = None  # lazily compiled fixed-grid tower path
+         self.morph_to_slow = morph_to_slow
+         if morph_to_slow:
+             # inject robot morphology into the slow LM prefix so the semantic
+@@ -112,14 +113,95 @@
+         """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.
++        directly, bypassing the language stack. The SAME tensor can be handed to
++        `forward(..., image_embeds=...)` so cam0 goes through the tower once per
++        step instead of twice (once here, once inside the LM call).
++
++        Uses a dense fixed-grid fast path: every image here is the same
++        (1, g, g) grid, so the HF packed-varlen tower (which runs B separate
++        SDPA calls PER LAYER and rebuilds pos/rotary tables in Python on every
++        call) is replaced by one batched SDPA per layer over (B, L, D), with the
++        per-image positional constants computed once and cached. Numerically
++        equivalent — per-image attention with equal lengths IS batched attention.
+         """
+         b = images.shape[0]
+         pixel_values, grid_thw = self._patchify(images)
++        fn = self._fixed_fn
++        if fn is None:
++            import os
++
++            if os.environ.get("TINYVLA_NO_COMPILE_VIT"):
++                fn = self._encode_image_fixed
++            else:
++                # static shapes -> compile is safe and worth ~2x on top of the
++                # dense rewrite; recompiles once per new batch size
++                fn = torch.compile(self._encode_image_fixed, dynamic=False)
++            self._fixed_fn = fn
++        try:
++            return fn(b, pixel_values)
++        except Exception:
++            if fn is not self._encode_image_fixed:  # compile blew up -> eager forever
++                self._fixed_fn = self._encode_image_fixed
++                return self._encode_image_fixed(b, pixel_values)
++            raise
++
++    def _encode_image_hf(self, images: torch.Tensor) -> torch.Tensor:
++        """Reference path through the stock HF tower (kept for verification)."""
++        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 _fixed_grid_ctx(self, device):
++        """Per-image positional constants for the fixed (1, g, g) grid, computed
++        ONCE via the tower's own functions (so the layout can't drift) and cached."""
++        g = self.image_size // self.patch_size
++        ctx = getattr(self, "_fixed_ctx", None)
++        if ctx is not None and ctx[0] == (g, str(device)):
++            return ctx[1]
++        visual = self.vlm.visual
++        grid1 = torch.tensor([[1, g, g]], device=device)
++        with torch.no_grad():
++            pos = visual.fast_pos_embed_interpolate(grid1)  # (g*g, hidden)
++            rot = visual.rot_pos_emb(grid1)                 # (g*g, head_dim/2)
++        emb = torch.cat((rot, rot), dim=-1)
++        val = (pos.detach(), emb.cos().detach(), emb.sin().detach())
++        self._fixed_ctx = ((g, str(device)), val)
++        return val
++
++    def _encode_image_fixed(self, b: int, pixel_values: torch.Tensor) -> torch.Tensor:
++        from transformers.models.qwen3_5.modeling_qwen3_5 import rotate_half
++
++        visual = self.vlm.visual
++        g = self.image_size // self.patch_size
++        L = g * g
++        x = visual.patch_embed(pixel_values.to(self.pixel_mean.dtype))  # (B*L, D)
++        pos, cos, sin = self._fixed_grid_ctx(x.device)
++        d = x.shape[-1]
++        x = x.view(b, L, d) + pos[None].to(x.dtype)
++        # (1, 1, L, hd) fp32 — same upcast-to-fp32 rotary as the HF path
++        cosf = cos[None, None].float()
++        sinf = sin[None, None].float()
++        for blk in visual.blocks:
++            attn = blk.attn
++            h = blk.norm1(x)
++            qkv = attn.qkv(h).view(b, L, 3, attn.num_heads, attn.head_dim)
++            q, k, v = qkv.permute(2, 0, 3, 1, 4).unbind(0)  # (B, nh, L, hd)
++            qf, kf = q.float(), k.float()
++            q = ((qf * cosf) + (rotate_half(qf) * sinf)).to(q.dtype)
++            k = ((kf * cosf) + (rotate_half(kf) * sinf)).to(k.dtype)
++            o = F.scaled_dot_product_attention(q, k, v, scale=attn.scaling)
++            o = o.transpose(1, 2).reshape(b, L, d)
++            x = x + attn.proj(o)
++            x = x + blk.mlp(blk.norm2(x))
++        merged = visual.merger(x.reshape(b * L, d))  # (B*L/m^2, out_hidden)
++        return merged.reshape(b, self.num_image_tokens, -1)
++
++    def _grid_thw(self, b: int, device) -> torch.Tensor:
++        g = self.image_size // self.patch_size
++        return torch.tensor([[1, g, g]], device=device).expand(b, -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).
+ 
+@@ -147,16 +229,24 @@
+ 
+     def forward(
+         self,
+-        images: torch.Tensor,  # (B, 3, H, W) in [0, 1]
++        images: torch.Tensor | None,  # (B, 3, H, W) in [0, 1]; may be None with image_embeds
+         lang_tokens: torch.Tensor,  # (B, L)
+         lang_mask: torch.Tensor,  # (B, L) bool, True = valid
+         morphology: torch.Tensor | None = None,  # (B, morph_dim) for morph_to_slow
++        image_embeds: torch.Tensor | None = None,  # (B, n_img, C) already through the tower
+     ) -> torch.Tensor:
+-        """Returns semantic latent (B, num_readout, out_dim)."""
+-        b = images.shape[0]
+-        device = images.device
++        """Returns semantic latent (B, num_readout, out_dim).
+ 
+-        pixel_values, grid_thw = self._patchify(images)
++        `image_embeds` lets the caller reuse the vision-tower output the fast path
++        already computed for the same frame — the tower is ~45 GFLOPs/camera and
++        cam0 was otherwise encoded twice per training step.
++        """
++        if image_embeds is not None:
++            b, device = image_embeds.shape[0], image_embeds.device
++            pixel_values, grid_thw = None, self._grid_thw(b, device)
++        else:
++            b, device = images.shape[0], images.device
++            pixel_values, grid_thw = self._patchify(images)
+         n_img = self.num_image_tokens
+ 
+         # <vision_start> <image>*n <vision_end> instruction <pad>*K
+@@ -180,19 +270,34 @@
+         )
+         mm_type = (ids == self.image_token_id).int()  # text 0, image 1
+ 
+-        if self.morph_to_slow and morphology is not None:
+-            # embed tokens, scatter image features, add morphology to readout slots,
+-            # then run the LM on inputs_embeds (keeps M-RoPE via mm_token_type_ids)
++        use_embeds = image_embeds is not None or (self.morph_to_slow and morphology is not None)
++        if use_embeds:
++            # embed tokens, scatter image features, optionally add morphology to the
++            # readout slots, then run the LM on inputs_embeds (keeps M-RoPE via
++            # mm_token_type_ids). This is also the path that consumes precomputed
++            # `image_embeds` — identical maths, one fewer vision-tower pass.
+             inputs_embeds = self.vlm.get_input_embeddings()(ids)
+-            img = self.vlm.get_image_features(pixel_values.to(self.pixel_mean.dtype), grid_thw)
+-            img = torch.cat(img.pooler_output, dim=0).to(inputs_embeds.dtype)  # (B*n_img, hidden)
++            if image_embeds is not None:
++                img = image_embeds.reshape(-1, image_embeds.shape[-1])
++            else:
++                img = self.vlm.get_image_features(pixel_values.to(self.pixel_mean.dtype), grid_thw)
++                img = torch.cat(img.pooler_output, dim=0)  # (B*n_img, hidden)
++            img = img.to(inputs_embeds.dtype)
+             imask = (ids == self.image_token_id).unsqueeze(-1).expand_as(inputs_embeds)
+             inputs_embeds = inputs_embeds.masked_scatter(imask, img)
+-            morph_tok = self.morph_to_lm(morphology.to(inputs_embeds.dtype))  # (B, hidden)
+-            inputs_embeds[:, -self.num_readout:] = inputs_embeds[:, -self.num_readout:] + morph_tok[:, None, :]
++            if self.morph_to_slow and morphology is not None:
++                morph_tok = self.morph_to_lm(morphology.to(inputs_embeds.dtype))  # (B, hidden)
++                inputs_embeds[:, -self.num_readout:] = inputs_embeds[:, -self.num_readout:] + morph_tok[:, None, :]
++            # With inputs_embeds and no input_ids the model cannot recompute
++            # M-RoPE and falls back to a CACHED rope_deltas from a previous call
++            # (crashes on a different batch size; silently wrong on the same one).
++            # Compute the 3D position ids explicitly from the token layout.
++            position_ids, _ = self.vlm.get_rope_index(
++                ids, image_grid_thw=grid_thw, attention_mask=attn, mm_token_type_ids=mm_type,
++            )
+             out = self.vlm(
+                 inputs_embeds=inputs_embeds, image_grid_thw=grid_thw, attention_mask=attn,
+-                mm_token_type_ids=mm_type, use_cache=False,
++                mm_token_type_ids=mm_type, position_ids=position_ids, use_cache=False,
+             )
+         else:
+             out = self.vlm(

tinyvla_b200/scripts/bench_model.py

@@ -0,0 +1,139 @@
+#!/usr/bin/env python
+"""Measure the TinyVLA training step on real hardware, synthetic data.
+
+Builds the model out of the raw modules (SemanticPath / FlowMatchingExpert /
+projections) so lerobot isn't needed, and reproduces _conditioning + the flow
+loss exactly as modeling_tinyvla.py does. Sweeps: precision, dedup of the cam0
+vision-tower pass, torch.compile, batch size.
+
+Prints samples/s -> hours per 1M frames and per the 8.64M-sample C-scaled budget.
+"""
+
+from __future__ import annotations
+
+import argparse
+import time
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from tinyvla.modules.expert import FlowMatchingExpert
+from tinyvla.modules.semantic import SemanticPath
+
+
+class Bench(nn.Module):
+    """_conditioning + forward of TinyVLAPolicy, minus the lerobot wrapper."""
+
+    def __init__(self, d=512, max_state=64, max_action=64, chunk=50,
+                 lm_layers=12, readout=8, freeze_lm=False, freeze_vision=True,
+                 model_name="Qwen/Qwen3.5-0.8B"):
+        super().__init__()
+        self.semantic = SemanticPath(
+            model_name=model_name, num_layers=lm_layers, num_readout=readout,
+            out_dim=d, image_size=256, freeze_lm=freeze_lm, freeze_vision=freeze_vision,
+        )
+        self.spatial_proj = nn.Linear(self.semantic.visual_hidden_size, d)
+        self.camera_emb = nn.Embedding(3, d)
+        self.state_proj = nn.Linear(max_state, d)
+        self.embodiment_emb = nn.Embedding(16, d)
+        self.expert = FlowMatchingExpert(action_dim=max_action, chunk_size=chunk, d_model=d)
+        self.chunk, self.max_action = chunk, max_action
+
+    def forward(self, batch, dedup=True):
+        cam0, cam1 = batch["cam0"], batch["cam1"]
+        f0 = self.semantic.encode_image(cam0)
+        f1 = self.semantic.encode_image(cam1)
+        if dedup:
+            latent = self.semantic(None, batch["tok"], batch["mask"], image_embeds=f0)
+        else:
+            latent = self.semantic(cam0, batch["tok"], batch["mask"])
+        spatial = torch.cat([
+            self.spatial_proj(f0) + self.camera_emb.weight[0][None, None],
+            self.spatial_proj(f1) + self.camera_emb.weight[1][None, None],
+        ], dim=1)
+        state_tok = self.state_proj(batch["state"])[:, None]
+        emb = self.embodiment_emb(batch["emb_id"])[:, None]
+        cond = torch.cat([latent, spatial, state_tok, emb], dim=1)
+
+        actions = batch["actions"]
+        b = actions.shape[0]
+        t = torch.rand(b, device=actions.device) * 0.999 + 0.001
+        noise = torch.randn_like(actions)
+        x_t = t[:, None, None] * noise + (1 - t[:, None, None]) * actions
+        pred = self.expert(x_t, t, cond)
+        return F.mse_loss(pred, noise - actions)
+
+
+def make_batch(b, dev, max_state=64, max_action=64, chunk=50, lm_len=48):
+    return {
+        "cam0": torch.rand(b, 3, 256, 256, device=dev),
+        "cam1": torch.rand(b, 3, 256, 256, device=dev),
+        "tok": torch.randint(1000, 5000, (b, lm_len), device=dev),
+        "mask": torch.ones(b, lm_len, dtype=torch.bool, device=dev),
+        "state": torch.randn(b, max_state, device=dev),
+        "actions": torch.randn(b, chunk, max_action, device=dev),
+        "emb_id": torch.randint(0, 16, (b,), device=dev),
+    }
+
+
+def run(model, opt, batch, dedup, steps, warmup, autocast):
+    for i in range(warmup + steps):
+        if i == warmup:
+            torch.cuda.synchronize()
+            t0 = time.time()
+        opt.zero_grad(set_to_none=True)
+        with torch.autocast("cuda", dtype=torch.bfloat16, enabled=autocast):
+            loss = model(batch, dedup=dedup)
+        loss.backward()
+        opt.step()
+    torch.cuda.synchronize()
+    return (time.time() - t0) / steps
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--model-name", default="Qwen/Qwen3.5-0.8B")
+    ap.add_argument("--steps", type=int, default=30)
+    ap.add_argument("--warmup", type=int, default=10)
+    ap.add_argument("--batches", type=int, nargs="+", default=[48, 96, 144, 192])
+    ap.add_argument("--compile", action="store_true")
+    ap.add_argument("--freeze-lm", action="store_true")
+    args = ap.parse_args()
+
+    torch.backends.cuda.matmul.allow_tf32 = True
+    torch.backends.cudnn.allow_tf32 = True
+    dev = torch.device("cuda")
+
+    model = Bench(freeze_lm=args.freeze_lm, model_name=args.model_name).to(dev)
+    n_train = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    n_all = sum(p.numel() for p in model.parameters())
+    print(f"params: {n_all/1e6:.1f}M total, {n_train/1e6:.1f}M trainable "
+          f"(freeze_lm={args.freeze_lm}, vision frozen)")
+
+    opt = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad],
+                            lr=1e-4, betas=(0.9, 0.95), weight_decay=1e-10, fused=True)
+    if args.compile:
+        model.expert = torch.compile(model.expert, dynamic=False)
+
+    BUDGET = 8_640_000  # C-scaled: 60k steps x eff batch 144
+    print(f"\n{'batch':>6} {'dedup':>6} {'ms/step':>9} {'samples/s':>10} "
+          f"{'ч на 1M кадров':>15} {'ч на 8.64M':>11} {'VRAM GB':>8}")
+    for b in args.batches:
+        batch = make_batch(b, dev)
+        for dedup in (True, False):
+            torch.cuda.reset_peak_memory_stats()
+            try:
+                dt = run(model, opt, batch, dedup, args.steps, args.warmup, autocast=True)
+            except torch.cuda.OutOfMemoryError:
+                print(f"{b:>6} {str(dedup):>6}      OOM")
+                torch.cuda.empty_cache()
+                continue
+            sps = b / dt
+            mem = torch.cuda.max_memory_allocated() / 1e9
+            print(f"{b:>6} {str(dedup):>6} {dt*1000:>9.1f} {sps:>10.0f} "
+                  f"{1e6/sps/3600:>15.2f} {BUDGET/sps/3600:>11.2f} {mem:>8.1f}", flush=True)
+
+
+if __name__ == "__main__":
+    main()

tinyvla_b200/scripts/build_shards.py

@@ -0,0 +1,392 @@
+#!/usr/bin/env python
+"""One-time shard builder: LeRobot v3 video dataset -> sequential JPEG shards.
+
+Why this exists
+---------------
+Training reads ONE random frame per sample. On a LeRobot v3 repo that means an
+h264 seek into a ~195 MB mp4 holding ~10 concatenated episodes, per camera, per
+sample. That is the bottleneck the project notes already call out ("video decode
+is the training bottleneck, use 12+ dataloader workers"), and no amount of GPU
+fixes it.
+
+This pass decodes every video ONCE, sequentially (the cheap direction), at a
+target rate, resizes to the training resolution, and writes JPEG bytes into
+parquet shards. Training then does sequential reads + JPEG decode, which is
+~100x cheaper per frame and trivially parallel.
+
+Layout produced in --out:
+    shards/shard-{i:05d}.parquet   episode_index, frame_index, cam0, cam1 (JPEG bytes)
+    frames.parquet                 episode_index, frame_index, state[D], action[D], task_index
+    tasks.json                     task_index -> instruction string
+    stats.json                     per-source mean/std for state and action
+    manifest.json                  fps, image_size, dims, camera keys, counts
+
+Usage (parallel across boxes: each takes a disjoint episode range):
+    python build_shards.py --repo-id zekaiwang/trex_dataset --out /data/trex_shards \
+        --hz 10 --image-size 256 --part 0 --num-parts 4
+
+Verified against zekaiwang/trex_dataset metadata: episodes are contiguous and
+ordered inside every mp4, frame counts match the timestamps exactly, and the
+declared fps (30) is the real one.
+"""
+
+from __future__ import annotations
+
+import argparse
+import io
+import json
+import os
+from pathlib import Path
+
+import numpy as np
+import pyarrow as pa
+import pyarrow.parquet as pq
+
+
+# --------------------------------------------------------------------- metadata
+
+
+def load_metadata(repo_id: str, revision: str | None):
+    from huggingface_hub import hf_hub_download
+
+    def grab(path):
+        return hf_hub_download(repo_id, path, repo_type="dataset", revision=revision)
+
+    info = json.loads(Path(grab("meta/info.json")).read_text())
+
+    from huggingface_hub import HfApi
+
+    api = HfApi()
+    files = api.list_repo_files(repo_id, repo_type="dataset", revision=revision)
+    ep_files = sorted(f for f in files if f.startswith("meta/episodes/") and f.endswith(".parquet"))
+    tables = [pq.read_table(grab(f)) for f in ep_files]
+    episodes = pa.concat_tables(tables) if len(tables) > 1 else tables[0]
+    return info, episodes, files
+
+
+def pick_cameras(info: dict, cam0: str | None, cam1: str | None):
+    """Mirror CanonicalSource's cam0/cam1 choice unless the caller overrides it.
+
+    CanonicalSource sorts the image keys, takes PRIMARY_RE (or the first key) as
+    cam0 and the first WRIST_RE match as cam1. On trex that resolves to
+    head_left / left_wrist — right_wrist and all 20 tactile streams are dropped.
+    Made explicit here rather than left to a regex.
+    """
+    import re
+
+    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$")
+    keys = sorted(k for k, v in info["features"].items() if v.get("dtype") == "video")
+    if cam0 is None:
+        cam0 = next((k for k in keys if PRIMARY_RE.search(k)), keys[0])
+    if cam1 is None:
+        cam1 = next((k for k in keys if k != cam0 and WRIST_RE.search(k)), None)
+    assert cam0 in keys, f"{cam0} not among {keys}"
+    assert cam1 is None or cam1 in keys, f"{cam1} not among {keys}"
+    return cam0, cam1
+
+
+# ------------------------------------------------------------------ video decode
+
+
+class VideoReader:
+    """Opens one mp4 at a time, seeks to an episode, decodes its frames forward.
+
+    Episodes inside a file are contiguous and time-ordered (verified), so walking
+    episodes in order keeps this effectively sequential: one seek per episode,
+    costing at most one extra GOP.
+    """
+
+    def __init__(self, repo_id, revision, key, cache_dir, evict: bool):
+        self.repo_id, self.revision, self.key = repo_id, revision, key
+        self.cache_dir, self.evict = cache_dir, evict
+        self._path = None
+        self._ident = None
+        self._container = None
+        self._stream = None
+
+    def _open(self, chunk_index: int, file_index: int):
+        ident = (chunk_index, file_index)
+        if ident == self._ident:
+            return
+        self._close(drop=True)
+        from huggingface_hub import hf_hub_download
+
+        rel = f"videos/{self.key}/chunk-{chunk_index:03d}/file-{file_index:03d}.mp4"
+        self._path = hf_hub_download(
+            self.repo_id, rel, repo_type="dataset", revision=self.revision, cache_dir=self.cache_dir
+        )
+        import av
+
+        self._container = av.open(self._path)
+        self._stream = self._container.streams.video[0]
+        self._stream.thread_type = "AUTO"  # frame+slice threading; the whole point
+        self._ident = ident
+
+    def _close(self, drop: bool):
+        if self._container is not None:
+            self._container.close()
+            self._container = None
+        if drop and self.evict and self._path and os.path.exists(self._path):
+            real = os.path.realpath(self._path)  # hf cache stores a symlink to the blob
+            for p in {self._path, real}:
+                try:
+                    os.remove(p)
+                except OSError:
+                    pass
+        self._path = self._ident = self._stream = None
+
+    def close(self):
+        self._close(drop=True)
+
+    def read_episode(self, chunk_index, file_index, from_ts, n_frames, fps, stride):
+        """Yield (frame_index_in_episode, PIL.Image) for every stride-th frame."""
+        self._open(chunk_index, file_index)
+        tb = self._stream.time_base
+        target = int(from_ts / tb)
+        self._container.seek(target, stream=self._stream, backward=True, any_frame=False)
+        wanted = set(range(0, n_frames, stride))
+        eps = 0.5 / fps  # half a frame of slack against float timestamps
+        got = 0
+        for frame in self._container.decode(self._stream):
+            if frame.pts is None:
+                continue
+            t = float(frame.pts * tb)
+            if t < from_ts - eps:
+                continue
+            idx = int(round((t - from_ts) * fps))
+            if idx >= n_frames:
+                break
+            if idx in wanted:
+                wanted.discard(idx)
+                got += 1
+                yield idx, frame.to_image()
+            if not wanted:
+                break
+        if wanted:
+            raise RuntimeError(
+                f"{self.key} chunk{chunk_index}/file{file_index} @{from_ts:.3f}: "
+                f"decoded {got}, still missing {len(wanted)} frames"
+            )
+
+
+def encode_jpeg(img, size: int, quality: int) -> bytes:
+    """Resize to (size, size) and JPEG-encode.
+
+    Plain resize, no crop: this reproduces CanonicalSource._img, which squashes
+    the aspect ratio with F.interpolate. Matching it matters more than being
+    geometrically nicer, otherwise the shards are off-distribution for the
+    pretrained checkpoint.
+    """
+    from PIL import Image
+
+    if img.size != (size, size):
+        img = img.resize((size, size), Image.BILINEAR)
+    buf = io.BytesIO()
+    img.save(buf, format="JPEG", quality=quality, subsampling=0)
+    return buf.getvalue()
+
+
+# ------------------------------------------------------------------------ state
+
+
+def load_frame_table(repo_id, revision, files, episodes, keep_mask, stride, cache_dir):
+    """Pull state/action/task for the kept frames out of the data parquets."""
+    from huggingface_hub import hf_hub_download
+
+    ep = {c: np.asarray(episodes[c]) for c in
+          ["episode_index", "length", "dataset_from_index", "data/chunk_index", "data/file_index"]}
+    want = {}  # (chunk, file) -> list of episode rows
+    for i in np.nonzero(keep_mask)[0]:
+        want.setdefault((int(ep["data/chunk_index"][i]), int(ep["data/file_index"][i])), []).append(i)
+
+    out_ep, out_fr, out_state, out_action, out_task = [], [], [], [], []
+    tasks: dict[str, int] = {}
+    for (ci, fi), rows in sorted(want.items()):
+        rel = f"data/chunk-{ci:03d}/file-{fi:03d}.parquet"
+        if rel not in files:
+            raise FileNotFoundError(rel)
+        path = hf_hub_download(repo_id, rel, repo_type="dataset", revision=revision, cache_dir=cache_dir)
+        t = pq.read_table(path, columns=["episode_index", "frame_index", "observation.state", "action", "task_index"])
+        d = {c: t[c].to_numpy(zero_copy_only=False) for c in t.column_names}
+        for r in rows:
+            e = int(ep["episode_index"][r])
+            sel = np.nonzero(d["episode_index"] == e)[0]
+            if len(sel) != int(ep["length"][r]):
+                raise RuntimeError(f"episode {e}: parquet has {len(sel)} rows, meta says {ep['length'][r]}")
+            order = np.argsort(d["frame_index"][sel])
+            sel = sel[order][::stride]
+            out_ep.append(np.full(len(sel), e, dtype=np.int32))
+            out_fr.append(np.arange(len(sel), dtype=np.int32))
+            out_state.append(np.stack([np.asarray(x, dtype=np.float32) for x in d["observation.state"][sel]]))
+            out_action.append(np.stack([np.asarray(x, dtype=np.float32) for x in d["action"][sel]]))
+            task = episodes["tasks"][r].as_py()
+            task = task[0] if isinstance(task, list) and task else ""
+            out_task.append(np.full(len(sel), tasks.setdefault(task, len(tasks)), dtype=np.int32))
+        del t, d
+    return (
+        np.concatenate(out_ep), np.concatenate(out_fr),
+        np.concatenate(out_state), np.concatenate(out_action),
+        np.concatenate(out_task), tasks,
+    )
+
+
+# ------------------------------------------------------------------------- main
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--repo-id", default="zekaiwang/trex_dataset")
+    ap.add_argument("--revision", default=None)
+    ap.add_argument("--out", type=Path, required=True)
+    ap.add_argument("--hz", type=float, default=10.0, help="target rate; 0 = keep native fps")
+    ap.add_argument("--image-size", type=int, default=256)
+    ap.add_argument("--jpeg-quality", type=int, default=92)
+    ap.add_argument("--rows-per-shard", type=int, default=20000)
+    ap.add_argument("--cam0", default=None, help="override the primary camera key")
+    ap.add_argument("--cam1", default=None, help="override the wrist camera key")
+    ap.add_argument("--part", type=int, default=0)
+    ap.add_argument("--num-parts", type=int, default=1)
+    ap.add_argument("--max-episodes", type=int, default=None, help="smoke test on the first N")
+    ap.add_argument("--cache-dir", default=None)
+    ap.add_argument("--keep-videos", action="store_true", help="do NOT delete mp4s after use")
+    args = ap.parse_args()
+
+    info, episodes, files = load_metadata(args.repo_id, args.revision)
+    fps = float(info["fps"])
+    stride = 1 if args.hz <= 0 else max(1, int(round(fps / args.hz)))
+    eff_hz = fps / stride
+    cam0, cam1 = pick_cameras(info, args.cam0, args.cam1)
+
+    n_ep = episodes.num_rows
+    lengths = np.asarray(episodes["length"])
+    # sanity: the LIBERO lesson — never trust declared fps
+    for key in (cam0, cam1):
+        if key is None:
+            continue
+        f = np.asarray(episodes[f"videos/{key}/from_timestamp"])
+        t = np.asarray(episodes[f"videos/{key}/to_timestamp"])
+        measured = lengths / (t - f)
+        if not np.allclose(measured, fps, atol=1e-3):
+            raise SystemExit(
+                f"declared fps={fps} but {key} measures {measured.min():.3f}..{measured.max():.3f} — "
+                "fix the rate before building shards"
+            )
+
+    order = np.arange(n_ep) if args.max_episodes is None else np.arange(min(n_ep, args.max_episodes))
+    mine = order[args.part :: args.num_parts]
+    keep = np.zeros(n_ep, dtype=bool)
+    keep[mine] = True
+
+    print(f"repo {args.repo_id}: {n_ep} episodes, {lengths.sum():,} frames @ {fps} Hz")
+    print(f"cam0={cam0}\ncam1={cam1}")
+    print(f"stride {stride} -> {eff_hz:g} Hz, part {args.part}/{args.num_parts}: {len(mine)} episodes")
+
+    args.out.mkdir(parents=True, exist_ok=True)
+    (args.out / "shards").mkdir(exist_ok=True)
+
+    print("collecting state/action ...")
+    ep_i, fr_i, state, action, task_i, tasks = load_frame_table(
+        args.repo_id, args.revision, files, episodes, keep, stride, args.cache_dir
+    )
+    print(f"  {len(ep_i):,} frames kept, state {state.shape[1]}d, action {action.shape[1]}d, {len(tasks)} tasks")
+
+    readers = {
+        "cam0": VideoReader(args.repo_id, args.revision, cam0, args.cache_dir, not args.keep_videos),
+        "cam1": None if cam1 is None else VideoReader(args.repo_id, args.revision, cam1, args.cache_dir, not args.keep_videos),
+    }
+
+    schema = pa.schema([
+        ("episode_index", pa.int32()), ("frame_index", pa.int32()),
+        ("cam0", pa.binary()), ("cam1", pa.binary()),
+    ])
+    buf = {k: [] for k in ["episode_index", "frame_index", "cam0", "cam1"]}
+    shard_id, written, in_shard = args.part * 100000, 0, 0
+
+    writer = None
+
+    def flush(final=False):
+        nonlocal writer, shard_id, buf, written, in_shard
+        if buf["episode_index"]:
+            if writer is None:
+                path = args.out / "shards" / f"shard-{shard_id:05d}.parquet"
+                writer = pq.ParquetWriter(path, schema, compression="zstd", compression_level=1)
+            writer.write_table(pa.table(buf, schema=schema))
+            written += len(buf["episode_index"])
+            in_shard += len(buf["episode_index"])
+            buf = {k: [] for k in buf}
+        if writer is not None and (final or in_shard >= args.rows_per_shard):
+            writer.close()
+            writer = None
+            shard_id += 1
+            in_shard = 0
+
+    import time
+
+    t0 = time.time()
+    for n, r in enumerate(mine):
+        e = int(episodes["episode_index"][r].as_py())
+        L = int(lengths[r])
+        frames = {}
+        for slot, key in (("cam0", cam0), ("cam1", cam1)):
+            if key is None:
+                continue
+            rd = readers[slot]
+            ci = int(episodes[f"videos/{key}/chunk_index"][r].as_py())
+            fi = int(episodes[f"videos/{key}/file_index"][r].as_py())
+            ts = float(episodes[f"videos/{key}/from_timestamp"][r].as_py())
+            frames[slot] = {
+                i // stride: encode_jpeg(img, args.image_size, args.jpeg_quality)
+                for i, img in rd.read_episode(ci, fi, ts, L, fps, stride)
+            }
+        n_kept = len(range(0, L, stride))
+        blank = b""
+        for j in range(n_kept):
+            buf["episode_index"].append(e)
+            buf["frame_index"].append(j)
+            buf["cam0"].append(frames["cam0"][j])
+            buf["cam1"].append(frames.get("cam1", {}).get(j, blank))
+        if len(buf["episode_index"]) >= 2048:
+            flush()
+        if (n + 1) % 50 == 0:
+            el = time.time() - t0
+            print(f"  {n+1}/{len(mine)} episodes, {written:,} frames, "
+                  f"{written/el:.0f} fr/s, eta {(len(mine)-n-1)*el/(n+1)/60:.1f} min", flush=True)
+    flush(final=True)
+    for rd in readers.values():
+        if rd is not None:
+            rd.close()
+
+    # frames.parquet holds only the kept rows, in shard order
+    pq.write_table(
+        pa.table({
+            "episode_index": pa.array(ep_i, pa.int32()),
+            "frame_index": pa.array(fr_i, pa.int32()),
+            "state": pa.array(list(state), pa.list_(pa.float32(), state.shape[1])),
+            "action": pa.array(list(action), pa.list_(pa.float32(), action.shape[1])),
+            "task_index": pa.array(task_i, pa.int32()),
+        }),
+        args.out / f"frames-part{args.part:02d}.parquet",
+        compression="zstd",
+    )
+    (args.out / f"tasks-part{args.part:02d}.json").write_text(
+        json.dumps({v: k for k, v in tasks.items()}, ensure_ascii=False)
+    )
+    (args.out / f"stats-part{args.part:02d}.json").write_text(json.dumps({
+        "observation.state": {"mean": state.mean(0).tolist(), "std": state.std(0).tolist()},
+        "action": {"mean": action.mean(0).tolist(), "std": action.std(0).tolist()},
+        "count": int(len(ep_i)),
+    }))
+    (args.out / f"manifest-part{args.part:02d}.json").write_text(json.dumps({
+        "repo_id": args.repo_id, "revision": args.revision,
+        "source_fps": fps, "stride": stride, "fps": eff_hz,
+        "image_size": args.image_size, "jpeg_quality": args.jpeg_quality,
+        "cam0": cam0, "cam1": cam1,
+        "state_dim": int(state.shape[1]), "action_dim": int(action.shape[1]),
+        "episodes": int(len(mine)), "frames": int(written),
+    }, indent=2))
+    print(f"done: {written:,} frames in {(time.time()-t0)/60:.1f} min -> {args.out}")
+
+
+if __name__ == "__main__":
+    main()

tinyvla_b200/scripts/build_shards_v2.py

@@ -0,0 +1,246 @@
+#!/usr/bin/env python
+"""Shard builder for LeRobot v2.0/v2.1 repos (one parquet + one mp4 PER EPISODE),
+including nested repos like nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim
+where each subdataset lives under <prefix>/.
+
+v2 layout (paths come from meta/info.json templates):
+    data_path : data/chunk-{episode_chunk:03d}/episode_{episode_index:06d}.parquet
+    video_path: videos/chunk-{episode_chunk:03d}/{video_key}/episode_{episode_index:06d}.mp4
+
+No seeking at all: every episode is its own mp4, decoded start to finish.
+Output layout matches build_shards.py, so ShardSource reads it unchanged.
+
+Usage:
+    python build_shards_v2.py \
+        --repo-id nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim \
+        --prefix gr1_arms_waist.CanToDrawer \
+        --out /data/shards/CanToDrawer --hz 10 \
+        --max-episodes 100 --workers 16
+"""
+
+from __future__ import annotations
+
+import argparse
+import io
+import json
+import os
+import time
+from concurrent.futures import ProcessPoolExecutor
+from pathlib import Path
+
+import numpy as np
+import pyarrow as pa
+import pyarrow.parquet as pq
+
+# module-level so ProcessPoolExecutor workers inherit them post-fork
+ARGS = None
+INFO = None
+CAMS = None
+
+
+def hf_get(repo_id, path, revision=None):
+    from huggingface_hub import hf_hub_download
+
+    return hf_hub_download(repo_id, path, repo_type="dataset", revision=revision)
+
+
+def pick_cameras(features: dict, cam0, cam1, cam_keys_re=None):
+    import re
+
+    PRIMARY_RE = re.compile(r"ego_view|top|front|base|main|scene|image_0$|images\.image$|image2$")
+    WRIST_RE = re.compile(r"wrist|gripper|hand|image_1$")
+    keys = sorted(k for k, v in features.items() if v.get("dtype") == "video")
+    if cam0 is None:
+        cam0 = next((k for k in keys if PRIMARY_RE.search(k)), keys[0])
+    if cam1 is None:
+        cam1 = next((k for k in keys if k != cam0 and WRIST_RE.search(k)), None)
+        if cam1 is None and len(keys) > 1:
+            cam1 = next(k for k in keys if k != cam0)
+    return cam0, cam1
+
+
+def encode_jpeg(arr: np.ndarray, size: int, quality: int) -> bytes:
+    from PIL import Image
+
+    img = Image.fromarray(arr)
+    if img.size != (size, size):
+        img = img.resize((size, size), Image.BILINEAR)
+    buf = io.BytesIO()
+    img.save(buf, format="JPEG", quality=quality, subsampling=0)
+    return buf.getvalue()
+
+
+def process_episode(job):
+    """One episode -> (rows for the shard, state/action arrays, task string).
+
+    Runs in a worker process. Downloads are hf-cache-deduplicated; the decode is
+    a plain forward pass over a per-episode mp4 (~200-2000 frames).
+    """
+    ep_idx, ep_chunk, task = job
+    repo, rev, prefix = ARGS.repo_id, ARGS.revision, ARGS.prefix
+    pfx = f"{prefix}/" if prefix else ""
+    stride = STRIDE
+
+    dp = INFO["data_path"].format(episode_chunk=ep_chunk, episode_index=ep_idx)
+    t = pq.read_table(hf_get(repo, pfx + dp, rev))
+    cols = {c: t[c].to_numpy(zero_copy_only=False) for c in t.column_names}
+    n = t.num_rows
+    keep = np.arange(0, n, stride)
+
+    state = np.stack([np.asarray(x, np.float32) for x in cols["observation.state"][keep]])
+    action = np.stack([np.asarray(x, np.float32) for x in cols["action"][keep]])
+
+    import av
+
+    jpgs = {}
+    for slot, key in (("cam0", CAMS[0]), ("cam1", CAMS[1])):
+        if key is None:
+            jpgs[slot] = None
+            continue
+        vp = INFO["video_path"].format(episode_chunk=ep_chunk, video_key=key, episode_index=ep_idx)
+        path = hf_get(repo, pfx + vp, rev)
+        out = []
+        with av.open(path) as c:
+            st = c.streams.video[0]
+            st.thread_type = "AUTO"
+            i = 0
+            for fr in c.decode(st):
+                if i % stride == 0:
+                    out.append(encode_jpeg(fr.to_rgb().to_ndarray(), ARGS.image_size, ARGS.jpeg_quality))
+                i += 1
+        if ARGS.evict:
+            real = os.path.realpath(path)
+            for p in {path, real}:
+                try:
+                    os.remove(p)
+                except OSError:
+                    pass
+        # A LONGER video is benign (recorders often append tail frames after the
+        # last logged step) — truncate to the parquet row count. A SHORTER one
+        # means missing observations: hard-fail beyond 1 frame of slack.
+        if len(out) + 1 < len(keep):
+            raise RuntimeError(f"ep {ep_idx} {key}: {len(out)} decoded vs {len(keep)} parquet rows")
+        jpgs[slot] = out[: len(keep)]
+
+    m = min(len(keep), *(len(v) for v in jpgs.values() if v is not None))
+    return ep_idx, jpgs["cam0"][:m], (jpgs["cam1"][:m] if jpgs["cam1"] else None), state[:m], action[:m], task
+
+
+def main():
+    global ARGS, INFO, CAMS, STRIDE
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--repo-id", required=True)
+    ap.add_argument("--prefix", default=None, help="subdataset dir inside the repo (nested repos)")
+    ap.add_argument("--revision", default=None)
+    ap.add_argument("--out", type=Path, required=True)
+    ap.add_argument("--hz", type=float, default=10.0)
+    ap.add_argument("--image-size", type=int, default=256)
+    ap.add_argument("--jpeg-quality", type=int, default=92)
+    ap.add_argument("--rows-per-shard", type=int, default=5000)
+    ap.add_argument("--cam0", default=None)
+    ap.add_argument("--cam1", default=None)
+    ap.add_argument("--max-episodes", type=int, default=None)
+    ap.add_argument("--workers", type=int, default=16)
+    ap.add_argument("--evict", action="store_true", help="delete each mp4 after decoding")
+    args = ap.parse_args()
+    ARGS = args
+
+    pfx = f"{args.prefix}/" if args.prefix else ""
+    INFO = json.loads(Path(hf_get(args.repo_id, pfx + "meta/info.json", args.revision)).read_text())
+    fps = float(INFO["fps"])
+    STRIDE = max(1, int(round(fps / args.hz))) if args.hz > 0 else 1
+    CAMS = pick_cameras(INFO["features"], args.cam0, args.cam1)
+
+    # episodes + tasks: jsonl in v2
+    episodes = []
+    with open(hf_get(args.repo_id, pfx + "meta/episodes.jsonl", args.revision)) as f:
+        for line in f:
+            episodes.append(json.loads(line))
+    if args.max_episodes:
+        episodes = episodes[: args.max_episodes]
+    chunk_size = INFO.get("chunks_size", 1000)
+
+    jobs = []
+    for e in episodes:
+        i = e["episode_index"]
+        task = e.get("tasks", [""])
+        task = task[0] if isinstance(task, list) and task else str(task)
+        jobs.append((i, i // chunk_size, task))
+
+    sd = INFO["features"]["observation.state"]["shape"][0]
+    ad = INFO["features"]["action"]["shape"][0]
+    print(f"{args.repo_id}{'/' + args.prefix if args.prefix else ''}: {len(jobs)} эпизодов, "
+          f"fps={fps} stride={STRIDE} -> {fps/STRIDE:g} Hz | state={sd} action={ad} | cam0={CAMS[0]} cam1={CAMS[1]}")
+
+    args.out.mkdir(parents=True, exist_ok=True)
+    (args.out / "shards").mkdir(exist_ok=True)
+    schema = pa.schema([("episode_index", pa.int32()), ("frame_index", pa.int32()),
+                        ("cam0", pa.binary()), ("cam1", pa.binary())])
+    buf = {k: [] for k in ("episode_index", "frame_index", "cam0", "cam1")}
+    shard_id, in_shard, written = 0, 0, 0
+    writer = None
+
+    def flush(final=False):
+        nonlocal writer, shard_id, buf, in_shard, written
+        if buf["episode_index"]:
+            if writer is None:
+                writer = pq.ParquetWriter(args.out / "shards" / f"shard-{shard_id:05d}.parquet",
+                                          schema, compression="zstd", compression_level=1)
+            writer.write_table(pa.table(buf, schema=schema))
+            written += len(buf["episode_index"])
+            in_shard += len(buf["episode_index"])
+            buf = {k: [] for k in buf}
+        if writer is not None and (final or in_shard >= args.rows_per_shard):
+            writer.close()
+            writer, in_shard = None, 0
+            shard_id += 1
+
+    all_ep, all_fr, all_state, all_action, all_ti = [], [], [], [], []
+    tasks: dict[str, int] = {}
+    t0 = time.time()
+    done = 0
+    with ProcessPoolExecutor(args.workers) as ex:
+        for ep_idx, c0, c1, state, action, task in ex.map(process_episode, jobs, chunksize=1):
+            m = len(c0)
+            for j in range(m):
+                buf["episode_index"].append(ep_idx)
+                buf["frame_index"].append(j)
+                buf["cam0"].append(c0[j])
+                buf["cam1"].append(c1[j] if c1 else b"")
+            if len(buf["episode_index"]) >= 2048:
+                flush()
+            all_ep.append(np.full(m, ep_idx, np.int32))
+            all_fr.append(np.arange(m, dtype=np.int32))
+            all_state.append(state[:m])
+            all_action.append(action[:m])
+            all_ti.append(np.full(m, tasks.setdefault(task, len(tasks)), np.int32))
+            done += 1
+            if done % 20 == 0:
+                el = time.time() - t0
+                print(f"  {done}/{len(jobs)} eps, {written:,} fr, {written/el:.0f} fr/s, "
+                      f"eta {(len(jobs)-done)*el/done/60:.1f} min", flush=True)
+    flush(final=True)
+
+    state = np.concatenate(all_state); action = np.concatenate(all_action)
+    ep_i = np.concatenate(all_ep); fr_i = np.concatenate(all_fr); ti = np.concatenate(all_ti)
+    pq.write_table(pa.table({
+        "episode_index": pa.array(ep_i, pa.int32()), "frame_index": pa.array(fr_i, pa.int32()),
+        "state": pa.array(list(state), pa.list_(pa.float32(), state.shape[1])),
+        "action": pa.array(list(action), pa.list_(pa.float32(), action.shape[1])),
+        "task_index": pa.array(ti, pa.int32())}), args.out / "frames-part00.parquet", compression="zstd")
+    (args.out / "tasks-part00.json").write_text(json.dumps({v: k for k, v in tasks.items()}, ensure_ascii=False))
+    (args.out / "stats-part00.json").write_text(json.dumps({
+        "observation.state": {"mean": state.mean(0).tolist(), "std": state.std(0).tolist()},
+        "action": {"mean": action.mean(0).tolist(), "std": action.std(0).tolist()},
+        "count": int(len(ep_i))}))
+    (args.out / "manifest-part00.json").write_text(json.dumps({
+        "repo_id": args.repo_id, "prefix": args.prefix, "source_fps": fps, "stride": STRIDE,
+        "fps": fps / STRIDE, "image_size": args.image_size, "jpeg_quality": args.jpeg_quality,
+        "cam0": CAMS[0], "cam1": CAMS[1], "state_dim": int(state.shape[1]),
+        "action_dim": int(action.shape[1]), "episodes": len(jobs), "frames": int(written)}, indent=2))
+    el = time.time() - t0
+    print(f"done: {written:,} кадров за {el/60:.1f} мин = {written/el:.0f} fr/s -> {args.out}")
+
+
+if __name__ == "__main__":
+    main()

tinyvla_b200/scripts/make_stream_specs.py

@@ -0,0 +1,119 @@
+#!/usr/bin/env python
+"""Probe the nvidia/physical-ai collection and emit a streaming specs YAML.
+
+Walks the robotics repos, finds every LeRobot v2.x (sub)dataset that
+HubEpisodeStream can consume, assigns embodiment ids BY ROBOT FAMILY (not per
+dataset — 127 ids would defeat the embedding) and weights proportional to
+sqrt(frames) (plain proportional lets the two biggest GR1 tasks dominate the
+mixture; sqrt keeps diversity without starving the big sources).
+
+v3.0 repos (BridgeData2_LeRobot_v3, LIBERO_LeRobot_v3) are NOT emitted — they
+need the v3 shard path (build_shards.py); both were already in the C-scaled
+pretraining mixture anyway.
+
+Usage:
+    python make_stream_specs.py --out configs/physical_ai_stream.yaml
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import math
+import os
+import urllib.request
+from concurrent.futures import ThreadPoolExecutor
+
+FAMILIES = [  # ordered: first match wins; ids continue after C-scaled's 0..9
+    ("gr1", 10), ("bimanual_panda_gripper", 11), ("bimanual_panda_hand", 12),
+    ("single_panda_gripper", 13), ("sim_behavior_r1_pro", 14), ("unitree_g1", 15),
+    ("g1", 15),
+]
+
+ROOTS = [
+    ("nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim", None),   # nested
+    ("nvidia/PhysicalAI-Robotics-GR00T-Teleop-Sim", "LeRobot"),    # nested under LeRobot/
+    ("nvidia/PhysicalAI-GR00T-Tuned-Tasks", None),                 # nested
+    ("nvidia/PhysicalAI-Robotics-GR00T-Teleop-G1", None),          # nested
+    ("nvidia/GR00T-N1.7-AppleToPlate", ""),                        # root-level v2.1
+]
+
+
+def hf_json(url: str, token: str | None):
+    req = urllib.request.Request(url, headers={"Authorization": f"Bearer {token}"} if token else {})
+    try:
+        return json.load(urllib.request.urlopen(req, timeout=60))
+    except Exception:
+        return None
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--out", default="configs/physical_ai_stream.yaml")
+    ap.add_argument("--hz", type=float, default=10.0)
+    args = ap.parse_args()
+    token = None
+    tp = os.path.expanduser("~/.cache/huggingface/token")
+    if os.path.exists(tp):
+        token = open(tp).read().strip()
+
+    jobs = []  # (repo_id, prefix)
+    for repo, sub in ROOTS:
+        if sub == "":  # root-level dataset
+            jobs.append((repo, ""))
+            continue
+        base = f"https://huggingface.co/api/datasets/{repo}/tree/main" + (f"/{sub}" if sub else "")
+        tree = hf_json(base, token) or []
+        for e in tree:
+            if e.get("type") == "directory":
+                jobs.append((repo, e["path"]))
+
+    def probe(job):
+        repo, pfx = job
+        p = f"{pfx}/" if pfx else ""
+        info = hf_json(f"https://huggingface.co/datasets/{repo}/resolve/main/{p}meta/info.json", token)
+        if not info or "total_frames" not in info:
+            return None
+        ver = str(info.get("codebase_version", ""))
+        if not ver.startswith("v2"):
+            return None  # v3 needs the shard path, v1 untested
+        return dict(
+            repo_id=repo, prefix=pfx, frames=info["total_frames"], fps=info["fps"],
+            state_dim=info["features"]["observation.state"]["shape"][0],
+            action_dim=info["features"]["action"]["shape"][0],
+            ncam=sum(1 for v in info["features"].values() if v.get("dtype") == "video"),
+        )
+
+    with ThreadPoolExecutor(16) as ex:
+        found = [r for r in ex.map(probe, jobs) if r]
+
+    max_state = max(r["state_dim"] for r in found)
+    max_action = max(r["action_dim"] for r in found)
+    total = sum(r["frames"] for r in found)
+    lines = [
+        "# Autogenerated by make_stream_specs.py — streaming mixture over the",
+        "# LeRobot-v2.x part of the nvidia/physical-ai collection.",
+        f"# {len(found)} datasets, {total:,} source frames.",
+        f"# max_state_dim needed: {max_state}; max_action_dim needed: {max_action}.",
+        "datasets:",
+    ]
+    for r in sorted(found, key=lambda x: (-x["frames"])):
+        name = r["prefix"] or r["repo_id"].split("/")[-1]
+        emb = next((i for k, i in FAMILIES if name.startswith(k) or f"/{k}" in name.lower() or k in name.lower()), 15)
+        w = math.sqrt(r["frames"])
+        lines += [
+            f"  - repo_id: {r['repo_id']}",
+            f"    prefix: \"{r['prefix']}\"",
+            f"    weight: {w:.1f}   # sqrt({r['frames']:,} fr); fps={r['fps']} state={r['state_dim']} act={r['action_dim']} cams={r['ncam']}",
+            f"    embodiment_id: {emb}",
+            f"    hz: {args.hz}",
+        ]
+    os.makedirs(os.path.dirname(args.out) or ".", exist_ok=True)
+    with open(args.out, "w") as f:
+        f.write("\n".join(lines) + "\n")
+    print(f"{len(found)} датасетов, {total:,} кадров -> {args.out}")
+    print(f"max_state_dim={max_state} max_action_dim={max_action}")
+
+
+if __name__ == "__main__":
+    main()

tinyvla_b200/scripts/pack_wds.py

@@ -0,0 +1,121 @@
+#!/usr/bin/env python
+"""Pack one shard-builder output dir into self-contained webdataset tars.
+
+Each wds sample carries EVERYTHING training needs — no sidecar lookups in the
+loader, which is what makes the mixture loader trivial and robust:
+
+    __key__     "{dsname}_{ep:06d}_{fr:05d}"
+    cam0.jpg    primary camera JPEG (as produced by build_shards*)
+    cam1.jpg    wrist camera JPEG, or b"" for single-camera datasets
+    meta.npz    state        float32[S]   normalized (dataset mean/std)
+                action_chunk float32[50,A] normalized FUTURE actions from this
+                             step; episode tail padded by repeating the last row
+                action_is_pad bool[50]
+                action_dim / state_dim / task_index / embodiment_id  int32
+
+Normalization happens HERE, at pack time, with the per-dataset stats the shard
+builder wrote — the loader never touches stats. The stats used are copied into
+the output manifest so unnormalization at eval time uses the same numbers.
+
+JPEG bytes are copied verbatim (no re-encode); tars are uncompressed because
+the payload already is.
+
+Usage:
+    python pack_wds.py --src /workspace/data/CanToDrawer \
+        --out /workspace/data/wds/CanToDrawer --chunk 50 --maxcount 3000
+"""
+
+from __future__ import annotations
+
+import argparse
+import io
+import json
+from pathlib import Path
+
+import numpy as np
+import pyarrow.parquet as pq
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--src", type=Path, required=True, help="shard-builder output dir")
+    ap.add_argument("--out", type=Path, required=True)
+    ap.add_argument("--name", default=None, help="dataset name for keys (default: src dirname)")
+    ap.add_argument("--chunk", type=int, default=50)
+    ap.add_argument("--maxcount", type=int, default=3000)
+    ap.add_argument("--embodiment-id", type=int, default=0)
+    args = ap.parse_args()
+    name = args.name or args.src.name
+
+    import webdataset as wds
+
+    manifest = json.loads(next(args.src.glob("manifest-part*.json")).read_text())
+    stats = json.loads(next(args.src.glob("stats-part*.json")).read_text())
+    s_mean = np.asarray(stats["observation.state"]["mean"], np.float32)
+    s_std = np.clip(np.asarray(stats["observation.state"]["std"], np.float32), 1e-6, None)
+    a_mean = np.asarray(stats["action"]["mean"], np.float32)
+    a_std = np.clip(np.asarray(stats["action"]["std"], np.float32), 1e-6, None)
+
+    # frames table -> per-episode contiguous state/action, RAM-resident
+    tabs = [pq.read_table(p) for p in sorted(args.src.glob("frames-part*.parquet"))]
+    import pyarrow as pa
+
+    t = pa.concat_tables(tabs) if len(tabs) > 1 else tabs[0]
+    ep = np.asarray(t["episode_index"], np.int64)
+    fr = np.asarray(t["frame_index"], np.int64)
+    order = np.lexsort((fr, ep))
+    state = np.asarray(t["state"].to_numpy(zero_copy_only=False).tolist(), np.float32)[order]
+    action = np.asarray(t["action"].to_numpy(zero_copy_only=False).tolist(), np.float32)[order]
+    task_i = np.asarray(t["task_index"], np.int32)[order]
+    ep, fr = ep[order], fr[order]
+    uniq, starts = np.unique(ep, return_index=True)
+    ep_start = dict(zip(uniq.tolist(), starts.tolist()))
+    ep_len = dict(zip(uniq.tolist(), np.diff(np.append(starts, len(ep))).tolist()))
+
+    state = (state - s_mean) / s_std
+    action = (action - a_mean) / a_std
+
+    S, A, CH = state.shape[1], action.shape[1], args.chunk
+    args.out.mkdir(parents=True, exist_ok=True)
+    sink = wds.ShardWriter(str(args.out / "shard-%05d.tar"), maxcount=args.maxcount, verbose=0)
+    n = 0
+    for shard in sorted((args.src / "shards").glob("shard-*.parquet")):
+        pf = pq.ParquetFile(shard)
+        for batch in pf.iter_batches(batch_size=512):
+            d = batch.to_pydict()
+            for e, f, c0, c1 in zip(d["episode_index"], d["frame_index"], d["cam0"], d["cam1"]):
+                e, f = int(e), int(f)
+                base, ln = ep_start[e], ep_len[e]
+                i = base + f
+                avail = min(CH, ln - f)
+                chunk = np.empty((CH, A), np.float32)
+                chunk[:avail] = action[i : i + avail]
+                if avail < CH:
+                    chunk[avail:] = action[i + avail - 1]
+                is_pad = np.zeros(CH, bool)
+                is_pad[avail:] = True
+                buf = io.BytesIO()
+                np.savez(buf, state=state[i], action_chunk=chunk, action_is_pad=is_pad,
+                         action_dim=np.int32(A), state_dim=np.int32(S),
+                         task_index=task_i[i], embodiment_id=np.int32(args.embodiment_id))
+                sink.write({"__key__": f"{name}_{e:06d}_{f:05d}",
+                            "cam0.jpg": c0, "cam1.jpg": bytes(c1 or b""),
+                            "meta.npz": buf.getvalue()})
+                n += 1
+    sink.close()
+
+    tasks = {}
+    for p in sorted(args.src.glob("tasks-part*.json")):
+        tasks.update(json.loads(p.read_text()))
+    (args.out / "tasks.json").write_text(json.dumps(tasks, ensure_ascii=False))
+    (args.out / "manifest.json").write_text(json.dumps({
+        **manifest, "name": name, "samples": n, "chunk": CH,
+        "embodiment_id": args.embodiment_id,
+        "norm_stats": stats,  # exactly what was applied, for unnormalization
+        "shards": sorted(p.name for p in args.out.glob("shard-*.tar")),
+    }, indent=2))
+    print(f"{name}: {n:,} семплов -> {len(list(args.out.glob('shard-*.tar')))} тарболов в {args.out}")
+
+
+if __name__ == "__main__":
+    main()

tinyvla_b200/scripts/train_fast.py

@@ -0,0 +1,309 @@
+#!/usr/bin/env python
+"""Training loop for shard-streamed data, tuned for a single H100.
+
+Differences from scripts/train.py (all of them are speed or correctness, none
+change the objective):
+
+  data      ShardSource (sequential parquet + JPEG) instead of LeRobotDataset
+            random access into h264. This is the change that matters.
+  attention expert routed through SDPA (see patches/modules_expert.py.diff)
+  vision    cam0 goes through the Qwen tower ONCE per step, not twice
+            (patches/modeling_tinyvla.py.diff)
+  tokens    task strings pre-tokenized once into a lookup, not per step
+  optimizer fused AdamW, foreach off, set_to_none
+  precision bf16 autocast + TF32 matmuls, grad_accum 1 at a large batch
+  compile   torch.compile on the expert (many small ops; biggest compile win)
+  resume    shape-aware load: tensors whose shape changed (state/action
+            projections when the action dim changes) are re-initialised and
+            REPORTED, instead of raising or silently loading garbage
+
+Usage:
+    python train_fast.py --config configs/trex_ft.yaml
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import math
+import time
+from pathlib import Path
+
+import torch
+import yaml
+
+
+# --------------------------------------------------------------------- policy
+
+
+def make_policy(cfg: dict, state_dim: int, action_dim: int):
+    from lerobot.configs import FeatureType, PolicyFeature
+    from tinyvla.configuration_tinyvla import TinyVLAConfig
+    from tinyvla.modeling_tinyvla import TinyVLAPolicy
+
+    pcfg = TinyVLAConfig(**cfg.get("policy", {}))
+    if action_dim > pcfg.max_action_dim or state_dim > pcfg.max_state_dim:
+        raise SystemExit(
+            f"dataset has state {state_dim}d / action {action_dim}d but the config caps them at "
+            f"{pcfg.max_state_dim} / {pcfg.max_action_dim}. Raise max_state_dim/max_action_dim — "
+            f"CanonicalSource._pad would TRUNCATE silently."
+        )
+    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=(action_dim,))}
+    pcfg.validate_features()
+    return TinyVLAPolicy(pcfg), pcfg
+
+
+def load_compatible(model, path: Path):
+    """Load a checkpoint, keeping only tensors whose shape still matches.
+
+    Continuing C-scaled on a 58-DoF bimanual robot changes the action/state
+    projections. torch's strict=False does NOT tolerate a shape change (it
+    raises), so filter explicitly and say out loud what was dropped — a silently
+    re-initialised action head is the difference between "fine-tuning" and
+    "training a new head on a frozen trunk".
+    """
+    from safetensors.torch import load_file
+
+    sd = load_file(path / "model.safetensors")
+    own = model.state_dict()
+    keep, reshaped, unexpected = {}, [], []
+    for k, v in sd.items():
+        if k not in own:
+            unexpected.append(k)
+        elif own[k].shape != v.shape:
+            reshaped.append((k, tuple(v.shape), tuple(own[k].shape)))
+        else:
+            keep[k] = v
+    missing = [k for k in own if k not in keep]
+    model.load_state_dict(keep, strict=False)
+    print(f"resume {path}: loaded {len(keep)}/{len(own)} tensors")
+    for k, a, b in reshaped:
+        print(f"  RE-INIT (shape changed) {k}: {a} -> {b}")
+    if unexpected:
+        print(f"  ignored {len(unexpected)} unexpected keys, e.g. {unexpected[:3]}")
+    left = [k for k in missing if all(k != r[0] for r in reshaped)]
+    if left:
+        print(f"  {len(left)} tensors kept at init, e.g. {left[:3]}")
+    return len(keep)
+
+
+# ----------------------------------------------------------------------- main
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--config", type=Path, required=True)
+    ap.add_argument("--profile-steps", type=int, default=0,
+                    help="run N steps, print throughput, exit (use to replace the estimate with a measurement)")
+    args = ap.parse_args()
+    cfg = yaml.safe_load(args.config.read_text())
+
+    torch.backends.cuda.matmul.allow_tf32 = True
+    torch.backends.cudnn.allow_tf32 = True
+    torch.backends.cudnn.benchmark = True
+
+    dev = torch.device("cuda")
+
+    src_cfg = cfg["source"]
+    morph = None
+    if cfg.get("morphology_descriptors") and cfg["policy"].get("conditioning") == "morph":
+        from tinyvla.modules.embodiment import MORPH_FIELDS
+
+        raw = yaml.safe_load(Path(cfg["morphology_descriptors"]).read_text())
+        sc = {"arm_dof": 0.1, "reach_m": 2, "gripper_width_m": 10, "num_cameras": 1 / 3,
+              "control_hz": 1 / 30, "joint_lo_mean": 1 / 3.1416, "joint_hi_mean": 1 / 3.1416,
+              "workspace_x": 2, "workspace_y": 2, "workspace_z": 2, "payload_kg": 0.2}
+        d = raw[src_cfg["morph_key"]]
+        morph = torch.tensor([d.get(f, 0) * sc.get(f, 1) for f in MORPH_FIELDS], dtype=torch.float32)
+
+    chunk = cfg["policy"]["chunk_size"]
+    if src_cfg.get("type") == "hub":
+        # stream episodes straight from HF hub (LeRobot v2.x), no local data at all
+        import os as _os
+
+        from tinyvla.data.streaming_hub import HubEpisodeStream
+
+        specs = yaml.safe_load(Path(src_cfg["specs"]).read_text())["datasets"]
+        tok_path = _os.path.expanduser("~/.cache/huggingface/token")
+        source = HubEpisodeStream(
+            specs,
+            token=open(tok_path).read().strip() if _os.path.exists(tok_path) else None,
+            chunk=chunk,
+            image_size=cfg["policy"].get("image_size", 256),
+            max_state_dim=cfg["policy"]["max_state_dim"],
+            max_action_dim=cfg["policy"]["max_action_dim"],
+            shuffle_buffer=cfg.get("shuffle_buffer", 4096),
+            seed=cfg.get("seed", 42),
+        )
+        print(f"source: HUB STREAM, {len(specs)} datasets (infinite mixture)")
+    else:
+        from tinyvla.data.shards import ShardSource
+
+        source = ShardSource(
+            root=src_cfg["root"],
+            embodiment_id=src_cfg.get("embodiment_id", 0),
+            chunk=chunk,
+            image_size=cfg["policy"].get("image_size", 256),
+            max_state_dim=cfg["policy"]["max_state_dim"],
+            max_action_dim=cfg["policy"]["max_action_dim"],
+            morphology=morph,
+            robot_prompt=src_cfg.get("robot_prompt"),
+            shuffle_buffer=cfg.get("shuffle_buffer", 8192),
+            seed=cfg.get("seed", 42),
+        )
+        print(f"source: {source.num_frames:,} frames @ {source.fps} Hz, "
+              f"state {source.state_dim}d action {source.action_dim}d, {len(source.shards)} shards")
+
+    if src_cfg.get("type") == "hub":
+        # per-dataset dims vary; the source pads everything to the config caps
+        policy, pcfg = make_policy(cfg, cfg["policy"]["max_state_dim"], cfg["policy"]["max_action_dim"])
+    else:
+        policy, pcfg = make_policy(cfg, source.state_dim, source.action_dim)
+    policy = policy.to(dev)
+
+    if cfg.get("resume_from"):
+        load_compatible(policy, Path(cfg["resume_from"]))
+
+    # ---- pre-tokenize every distinct instruction once -----------------------
+    from transformers import AutoTokenizer
+
+    tokenizer = AutoTokenizer.from_pretrained(pcfg.lm_model_name)
+    if src_cfg.get("type") == "hub":
+        _tok_cache: dict = {}
+
+        def _tok(tasks):
+            new = [t for t in set(tasks) if t not in _tok_cache]
+            if new:
+                e = tokenizer(new, padding="max_length", truncation=True,
+                              max_length=pcfg.tokenizer_max_length, return_tensors="pt")
+                for i, t in enumerate(new):
+                    _tok_cache[t] = (e["input_ids"][i], e["attention_mask"][i].bool())
+            return (torch.stack([_tok_cache[t][0] for t in tasks]),
+                    torch.stack([_tok_cache[t][1] for t in tasks]))
+    else:
+        texts = sorted(set(source._tasks.values()))
+        enc = tokenizer(texts, padding="max_length", truncation=True,
+                        max_length=pcfg.tokenizer_max_length, return_tensors="pt")
+        tok_ids, tok_mask = enc["input_ids"], enc["attention_mask"].bool()
+        tok_lookup = {t: i for i, t in enumerate(texts)}
+        print(f"pre-tokenized {len(texts)} instructions at fixed length {pcfg.tokenizer_max_length}")
+
+        def _tok(tasks):
+            idx = torch.tensor([tok_lookup.get(t, 0) for t in tasks])
+            return tok_ids[idx], tok_mask[idx]
+
+    def collate(items):
+        out = {}
+        for k in items[0]:
+            if k == "task":
+                ids, mask = _tok([it["task"] for it in items])
+                out["observation.language.tokens"] = ids
+                out["observation.language.attention_mask"] = mask
+            else:
+                out[k] = torch.stack([it[k] for it in items])
+        return out
+
+    loader = torch.utils.data.DataLoader(
+        source,
+        batch_size=cfg["batch_size"],
+        num_workers=cfg.get("num_workers", 12),
+        pin_memory=True,
+        persistent_workers=cfg.get("num_workers", 12) > 0,
+        prefetch_factor=cfg.get("prefetch_factor", 6) if cfg.get("num_workers", 12) > 0 else None,
+        drop_last=True,
+        collate_fn=collate,
+    )
+
+    # ---- optimizer: backbone at a lower lr, exactly as train.py does --------
+    backbone = [p for n, p in policy.named_parameters() if p.requires_grad and "semantic.vlm" in n]
+    head = [p for n, p in policy.named_parameters() if p.requires_grad and "semantic.vlm" not in n]
+    groups = [{"params": head, "lr": cfg["lr"]}]
+    if backbone:
+        groups.append({"params": backbone, "lr": cfg["lr"] * cfg.get("backbone_lr_mult", 0.1)})
+    opt = torch.optim.AdamW(groups, betas=(0.9, 0.95), weight_decay=1e-10, fused=True)
+    print(f"trainable: head {sum(p.numel() for p in head)/1e6:.1f}M, "
+          f"backbone {sum(p.numel() for p in backbone)/1e6:.1f}M at {cfg.get('backbone_lr_mult',0.1)}x lr")
+
+    steps, warmup = cfg["steps"], cfg.get("warmup_steps", 1000)
+
+    def lr_lambda(s):
+        if s < warmup:
+            return s / max(1, 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)
+
+    if cfg.get("compile", True):
+        policy.expert = torch.compile(policy.expert, dynamic=False)
+        print("torch.compile: expert")
+
+    out_dir = Path(cfg["output_dir"])
+    out_dir.mkdir(parents=True, exist_ok=True)
+    if hasattr(source, "manifest"):
+        (out_dir / "source_manifest.json").write_text(json.dumps(source.manifest, indent=2))
+    if cfg.get("wandb"):
+        import wandb
+
+        wandb.init(project=cfg["wandb"], config=cfg)
+
+    grad_accum = cfg.get("grad_accum", 1)
+    eff_batch = cfg["batch_size"] * grad_accum
+    log_freq = cfg.get("log_freq", 50)
+    target = args.profile_steps or steps
+
+    step, seen, t0 = 0, 0, time.time()
+    data_iter = iter(loader)
+    while step < target:
+        opt.zero_grad(set_to_none=True)
+        for _ in range(grad_accum):
+            try:
+                batch = next(data_iter)
+            except StopIteration:
+                source.set_epoch(source.epoch + 1)
+                data_iter = iter(loader)
+                batch = next(data_iter)
+            batch = {k: (v.to(dev, non_blocking=True) if torch.is_tensor(v) else v)
+                     for k, v in batch.items()}
+            with torch.autocast("cuda", dtype=torch.bfloat16):
+                loss, info = policy(batch)
+            (loss / grad_accum).backward()
+            seen += cfg["batch_size"]
+        torch.nn.utils.clip_grad_norm_(policy.parameters(), cfg.get("grad_clip", 10.0))
+        opt.step()
+        sched.step()
+        step += 1
+
+        if step % log_freq == 0:
+            torch.cuda.synchronize()
+            el = time.time() - t0
+            sps = seen / el
+            gflops = cfg.get("gflops_per_sample", 360.0)
+            mfu = sps * gflops * 1e9 / (cfg.get("peak_flops", 989e12)) * 100
+            eta = (steps - step) * eff_batch / sps / 3600
+            print(f"step {step}/{steps} loss {info['loss']:.4f} | {sps:.0f} samples/s "
+                  f"| {step/el:.2f} it/s | MFU~{mfu:.0f}% | eta {eta:.2f} h "
+                  f"| mem {torch.cuda.max_memory_allocated()/1e9:.1f} GB", flush=True)
+            if cfg.get("wandb"):
+                wandb.log({"loss": info["loss"], "lr": sched.get_last_lr()[0],
+                           "samples_per_s": sps}, step=step)
+            seen, t0 = 0, time.time()
+
+        if args.profile_steps == 0 and step % cfg.get("save_freq", 10000) == 0:
+            policy.save_pretrained(out_dir / f"step_{step}")
+
+    if args.profile_steps:
+        print(f"\nprofile done. Multiply the samples/s above by {steps * eff_batch:,} "
+              f"total samples to get the real wall clock.")
+        return
+    policy.save_pretrained(out_dir / "final")
+
+
+if __name__ == "__main__":
+    main()

tinyvla_b200/src/tinyvla/data/shards.py

@@ -0,0 +1,222 @@
+"""Streaming source over prebuilt JPEG shards (see scripts/build_shards.py).
+
+Emits exactly the schema CanonicalSource emits, so it drops into the existing
+WeightedMixtureDataset / train loop unchanged:
+
+  observation.images.cam0   (3, S, S) float [0,1]
+  observation.images.cam1   (3, S, S) float [0,1]  (zeros if the source has one camera)
+  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
+  action_is_pad             (chunk,) bool
+  embodiment_id             () long
+  task                      str
+  morphology                (16,) float           (variant C)
+
+Difference from CanonicalSource: this is an IterableDataset doing SEQUENTIAL
+reads with a shuffle buffer, instead of random row lookups that each trigger an
+h264 seek. That is the whole speedup — the model side is untouched.
+"""
+
+from __future__ import annotations
+
+import io
+import json
+import random
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch.utils.data import IterableDataset, get_worker_info
+
+
+def _pad(x: torch.Tensor, dim: int) -> torch.Tensor:
+    """Zero-pad the last dim. Unlike CanonicalSource._pad this REFUSES to
+    truncate: silently cutting a 58-dim bimanual action down to 8 is exactly the
+    kind of convention bug that costs weeks."""
+    d = x.shape[-1]
+    if d == dim:
+        return x
+    if d > dim:
+        raise ValueError(
+            f"tensor has {d} dims but max is {dim} — raise max_action_dim/max_state_dim "
+            f"in the policy config instead of truncating"
+        )
+    return F.pad(x, (0, dim - d))
+
+
+class ShardSource(IterableDataset):
+    def __init__(
+        self,
+        root: str | Path,
+        embodiment_id: int,
+        chunk: int = 50,
+        image_size: int = 256,
+        max_state_dim: int = 64,
+        max_action_dim: int = 64,
+        morphology: torch.Tensor | None = None,
+        robot_prompt: str | None = None,
+        shuffle_buffer: int = 8192,
+        seed: int = 42,
+        normalize: bool = True,
+        rank: int = 0,
+        world_size: int = 1,
+    ):
+        super().__init__()
+        self.root = Path(root)
+        self.embodiment_id = embodiment_id
+        self.chunk = chunk
+        self.image_size = image_size
+        self.max_state_dim = max_state_dim
+        self.max_action_dim = max_action_dim
+        self.morphology = morphology
+        self.robot_prompt = robot_prompt
+        self.shuffle_buffer = shuffle_buffer
+        self.seed = seed
+        self.normalize = normalize
+        self.rank, self.world_size = rank, world_size
+        self.epoch = 0
+
+        self.shards = sorted(str(p) for p in (self.root / "shards").glob("shard-*.parquet"))
+        if not self.shards:
+            raise FileNotFoundError(f"no shards under {self.root/'shards'}")
+
+        parts = sorted(self.root.glob("manifest-part*.json"))
+        self.manifest = json.loads(parts[0].read_text())
+        self.fps = self.manifest["fps"]
+        self.action_dim = self.manifest["action_dim"]
+        self.state_dim = self.manifest["state_dim"]
+
+        # per-episode state/action, resident in RAM: 58 floats x 2 x ~1.8M frames
+        # is well under 1 GB, and it removes every random lookup from the hot path
+        import pyarrow.parquet as pq
+
+        tabs = [pq.read_table(p) for p in sorted(self.root.glob("frames-part*.parquet"))]
+        import pyarrow as pa
+
+        t = pa.concat_tables(tabs) if len(tabs) > 1 else tabs[0]
+        ep = np.asarray(t["episode_index"]).astype(np.int64)
+        fr = np.asarray(t["frame_index"]).astype(np.int64)
+        order = np.lexsort((fr, ep))
+        self._state = np.asarray(t["state"].to_numpy(zero_copy_only=False).tolist(), dtype=np.float32)[order]
+        self._action = np.asarray(t["action"].to_numpy(zero_copy_only=False).tolist(), dtype=np.float32)[order]
+        self._task_index = np.asarray(t["task_index"]).astype(np.int64)[order]
+        ep, fr = ep[order], fr[order]
+        uniq, starts = np.unique(ep, return_index=True)
+        self._ep_start = dict(zip(uniq.tolist(), starts.tolist()))
+        self._ep_len = dict(zip(uniq.tolist(), np.diff(np.append(starts, len(ep))).tolist()))
+
+        tasks = {}
+        for p in sorted(self.root.glob("tasks-part*.json")):
+            tasks.update({int(k): v for k, v in json.loads(p.read_text()).items()})
+        self._tasks = tasks
+
+        stats = {}
+        for p in sorted(self.root.glob("stats-part*.json")):
+            s = json.loads(p.read_text())
+            for k in ("observation.state", "action"):
+                stats.setdefault(k, []).append((np.array(s[k]["mean"]), np.array(s[k]["std"]), s["count"]))
+        self.stats = {}
+        for k, entries in stats.items():  # pooled mean/std across parts
+            n = sum(e[2] for e in entries)
+            mean = sum(e[0] * e[2] for e in entries) / n
+            var = sum((e[1] ** 2 + (e[0] - mean) ** 2) * e[2] for e in entries) / n
+            self.stats[k] = (
+                torch.tensor(mean, dtype=torch.float32),
+                torch.tensor(np.sqrt(var), dtype=torch.float32).clamp(min=1e-6),
+            )
+
+        self.num_frames = int(len(ep))
+
+    def set_epoch(self, epoch: int):
+        self.epoch = epoch
+
+    # ------------------------------------------------------------------ helpers
+
+    def _decode(self, blob: bytes) -> torch.Tensor:
+        from PIL import Image
+
+        img = Image.open(io.BytesIO(blob))
+        img = img.convert("RGB")
+        x = torch.from_numpy(np.asarray(img, dtype=np.uint8).copy())
+        x = x.permute(2, 0, 1).float().div_(255.0)
+        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 _sample(self, e: int, f: int, cam0: bytes, cam1: bytes) -> dict:
+        out = {"observation.images.cam0": self._decode(cam0)}
+        if cam1:
+            out["observation.images.cam1"] = self._decode(cam1)
+            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
+
+        base, n = self._ep_start[e], self._ep_len[e]
+        i = base + f
+        state = torch.from_numpy(self._state[i].copy())
+        avail = min(self.chunk, n - f)
+        act = torch.from_numpy(self._action[i : i + avail].copy())
+        is_pad = torch.zeros(self.chunk, dtype=torch.bool)
+        if avail < self.chunk:
+            # hold the last action, mark the tail as padding — the loss masks it
+            act = torch.cat([act, act[-1:].expand(self.chunk - avail, -1)], dim=0)
+            is_pad[avail:] = True
+
+        if self.normalize:
+            m, s = self.stats["observation.state"]
+            state = (state - m) / s
+            m, s = self.stats["action"]
+            act = (act - m) / s
+
+        out["observation.state"] = _pad(state, self.max_state_dim)
+        out["action"] = _pad(act, self.max_action_dim)
+        mask = torch.zeros(self.max_action_dim, dtype=torch.bool)
+        mask[: self.action_dim] = True
+        out["action_dim_mask"] = mask
+        out["action_is_pad"] = is_pad
+        out["embodiment_id"] = torch.tensor(self.embodiment_id, dtype=torch.long)
+        task = self._tasks.get(int(self._task_index[i]), "")
+        out["task"] = f"{self.robot_prompt} {task}" if self.robot_prompt else task
+        if self.morphology is not None:
+            out["morphology"] = self.morphology
+        return out
+
+    # -------------------------------------------------------------------- iter
+
+    def _my_shards(self):
+        wi = get_worker_info()
+        wid, nw = (wi.id, wi.num_workers) if wi is not None else (0, 1)
+        gid = self.rank * nw + wid
+        gn = self.world_size * nw
+        shards = list(self.shards)
+        random.Random(self.seed + self.epoch).shuffle(shards)
+        return shards[gid::gn], gid
+
+    def __iter__(self):
+        import pyarrow.parquet as pq
+
+        shards, gid = self._my_shards()
+        rng = random.Random(self.seed * 1000 + self.epoch * 97 + gid)
+        buf: list[dict] = []
+        for path in shards:
+            pf = pq.ParquetFile(path)
+            for batch in pf.iter_batches(batch_size=512):
+                d = batch.to_pydict()
+                for e, f, c0, c1 in zip(d["episode_index"], d["frame_index"], d["cam0"], d["cam1"]):
+                    if e not in self._ep_start:
+                        continue  # frame table and shards built from different parts
+                    item = self._sample(int(e), int(f), c0, c1)
+                    if len(buf) < self.shuffle_buffer:
+                        buf.append(item)
+                        continue
+                    j = rng.randrange(len(buf))
+                    buf[j], item = item, buf[j]
+                    yield item
+        rng.shuffle(buf)
+        yield from buf

tinyvla_b200/src/tinyvla/data/streaming_hub.py

@@ -0,0 +1,277 @@
+"""Stream training samples straight from HF hub LeRobot v2.x repos. No disk.
+
+Why this works: v2.0/v2.1 stores ONE parquet + ONE mp4 per camera PER EPISODE
+(2-10 MB each). A DataLoader worker samples a (dataset, episode) pair by weight,
+HTTP-GETs the files into memory, decodes the episode SEQUENTIALLY (the cheap
+direction — no seeks at all), yields every stride-th frame as a full training
+sample, then moves to the next episode. A downstream shuffle buffer breaks the
+intra-episode correlation. Nothing ever touches disk, so the 1.3 TB
+materialization step disappears entirely.
+
+Bandwidth math (nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim): source h264
+is ~10 KB/frame over all cams; at 400 samples/s consumption and stride 2-3 the
+stream needs ~10-25 MB/s and ~4 HTTP requests/s — trivial for the HF CDN.
+
+Emits the exact CanonicalSource/ShardSource schema, so it drops into the
+training loop unchanged.
+"""
+
+from __future__ import annotations
+
+import io
+import json
+import random
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch.utils.data import IterableDataset, get_worker_info
+
+
+def _resolve(repo_id: str, path: str) -> str:
+    return f"https://huggingface.co/datasets/{repo_id}/resolve/main/{path}"
+
+
+def _pad(x: torch.Tensor, dim: int, what: str) -> torch.Tensor:
+    d = x.shape[-1]
+    if d > dim:
+        raise ValueError(f"{what} has {d} dims > max {dim}; raise the cap, never truncate")
+    return F.pad(x, (0, dim - d)) if d < dim else x
+
+
+class HubEpisodeStream(IterableDataset):
+    """Infinite weighted mixture over LeRobot v2.x datasets, streamed from HF.
+
+    spec per dataset:
+      repo_id     HF dataset repo
+      prefix      subdir for nested repos ("" for root-level)
+      weight      mixture sampling weight
+      embodiment_id
+      hz          target rate (stride = round(fps/hz))
+    """
+
+    def __init__(
+        self,
+        specs: list[dict],
+        token: str | None = None,
+        chunk: int = 50,
+        image_size: int = 256,
+        max_state_dim: int = 64,
+        max_action_dim: int = 64,
+        shuffle_buffer: int = 4096,
+        seed: int = 42,
+        rank: int = 0,
+        world_size: int = 1,
+    ):
+        super().__init__()
+        self.specs = specs
+        self.token = token
+        self.chunk = chunk
+        self.image_size = image_size
+        self.max_state_dim = max_state_dim
+        self.max_action_dim = max_action_dim
+        self.shuffle_buffer = shuffle_buffer
+        self.seed = seed
+        self.rank, self.world_size = rank, world_size
+        self.epoch = 0
+        self._session = None  # per-process requests.Session (created lazily post-fork)
+        self._meta: dict[int, dict] = {}
+
+    # ------------------------------------------------------------------- http
+
+    def _http(self):
+        if self._session is None:
+            import requests
+
+            s = requests.Session()
+            if self.token:
+                s.headers["Authorization"] = f"Bearer {self.token}"
+            a = requests.adapters.HTTPAdapter(pool_connections=8, pool_maxsize=8, max_retries=3)
+            s.mount("https://", a)
+            self._session = s
+        return self._session
+
+    def _fetch(self, repo_id: str, path: str) -> bytes:
+        r = self._http().get(_resolve(repo_id, path), timeout=120)
+        r.raise_for_status()
+        return r.content
+
+    # ------------------------------------------------------------------- meta
+
+    def _load_meta(self, si: int) -> dict:
+        """info.json + episodes list + normalization stats, once per worker."""
+        if si in self._meta:
+            return self._meta[si]
+        spec = self.specs[si]
+        rid, pfx = spec["repo_id"], spec.get("prefix", "")
+        p = f"{pfx}/" if pfx else ""
+        info = json.loads(self._fetch(rid, p + "meta/info.json"))
+        fps = float(info["fps"])
+        stride = max(1, int(round(fps / spec.get("hz", 10.0))))
+
+        episodes = []
+        for line in self._fetch(rid, p + "meta/episodes.jsonl").decode().splitlines():
+            if line.strip():
+                e = json.loads(line)
+                episodes.append((e["episode_index"], e.get("length", 0),
+                                 (e.get("tasks") or [""])[0] if isinstance(e.get("tasks"), list) else ""))
+
+        stats = self._load_stats(rid, p, info)
+        cam0, cam1 = self._pick_cams(info["features"])
+        meta = dict(info=info, fps=fps, stride=stride, episodes=episodes, stats=stats,
+                    cams=(cam0, cam1), chunks_size=info.get("chunks_size", 1000),
+                    state_dim=info["features"]["observation.state"]["shape"][0],
+                    action_dim=info["features"]["action"]["shape"][0])
+        self._meta[si] = meta
+        return meta
+
+    def _load_stats(self, rid: str, p: str, info: dict) -> dict:
+        """v2.0: meta/stats.json. v2.1: aggregate meta/episodes_stats.jsonl."""
+        try:
+            raw = json.loads(self._fetch(rid, p + "meta/stats.json"))
+            return {k: (np.asarray(raw[k]["mean"], np.float32).reshape(-1),
+                        np.asarray(raw[k]["std"], np.float32).reshape(-1))
+                    for k in ("observation.state", "action")}
+        except Exception:
+            pass
+        acc = {k: [0.0, 0.0, 0] for k in ("observation.state", "action")}  # sum, sumsq, n
+        for line in self._fetch(rid, p + "meta/episodes_stats.jsonl").decode().splitlines():
+            if not line.strip():
+                continue
+            st = json.loads(line)["stats"]
+            for k in acc:
+                m = np.asarray(st[k]["mean"], np.float32).reshape(-1)
+                s = np.asarray(st[k]["std"], np.float32).reshape(-1)
+                n = int(np.asarray(st[k]["count"]).reshape(-1)[0])
+                acc[k][0] = acc[k][0] + m * n
+                acc[k][1] = acc[k][1] + (s ** 2 + m ** 2) * n
+                acc[k][2] += n
+        out = {}
+        for k, (sm, sq, n) in acc.items():
+            mean = sm / n
+            out[k] = (mean.astype(np.float32),
+                      np.sqrt(np.maximum(sq / n - mean ** 2, 1e-12)).astype(np.float32))
+        return out
+
+    @staticmethod
+    def _pick_cams(features: dict):
+        import re
+
+        PRIMARY = re.compile(r"ego_view|top|front|base|main|scene|image_0$|images\.image$|image2$")
+        WRIST = re.compile(r"wrist|gripper|hand|image_1$")
+        keys = sorted(k for k, v in features.items() if v.get("dtype") == "video")
+        cam0 = next((k for k in keys if PRIMARY.search(k)), keys[0])
+        cam1 = next((k for k in keys if k != cam0 and WRIST.search(k)), None)
+        if cam1 is None and len(keys) > 1:
+            cam1 = next(k for k in keys if k != cam0)
+        return cam0, cam1
+
+    # ---------------------------------------------------------------- episode
+
+    def _decode_video(self, blob: bytes, stride: int) -> list[torch.Tensor]:
+        import av
+
+        out = []
+        with av.open(io.BytesIO(blob)) as c:
+            st = c.streams.video[0]
+            st.thread_type = "AUTO"
+            i = 0
+            for fr in c.decode(st):
+                if i % stride == 0:
+                    a = fr.to_rgb().to_ndarray()
+                    x = torch.from_numpy(np.ascontiguousarray(a)).permute(2, 0, 1).float().div_(255.0)
+                    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]
+                    out.append(x)
+                i += 1
+        return out
+
+    def _stream_episode(self, si: int, ep_idx: int, task: str):
+        spec = self.specs[si]
+        meta = self._load_meta(si)
+        rid, pfx = spec["repo_id"], spec.get("prefix", "")
+        p = f"{pfx}/" if pfx else ""
+        info, stride = meta["info"], meta["stride"]
+        ch = ep_idx // meta["chunks_size"]
+
+        t = __import__("pyarrow.parquet", fromlist=["read_table"]).read_table(
+            io.BytesIO(self._fetch(rid, p + info["data_path"].format(
+                episode_chunk=ch, episode_index=ep_idx))))
+        keep = np.arange(0, t.num_rows, stride)
+        cols = {c: t[c].to_numpy(zero_copy_only=False) for c in ("observation.state", "action")}
+        sm, ss = meta["stats"]["observation.state"]
+        am, as_ = meta["stats"]["action"]
+        state = (np.stack([np.asarray(x, np.float32) for x in cols["observation.state"][keep]]) - sm) / np.maximum(ss, 1e-6)
+        action = (np.stack([np.asarray(x, np.float32) for x in cols["action"][keep]]) - am) / np.maximum(as_, 1e-6)
+
+        cams = {}
+        for slot, key in zip(("cam0", "cam1"), meta["cams"]):
+            if key is None:
+                cams[slot] = None
+                continue
+            blob = self._fetch(rid, p + info["video_path"].format(
+                episode_chunk=ch, video_key=key, episode_index=ep_idx))
+            frames = self._decode_video(blob, stride)
+            if len(frames) + 1 < len(keep):  # video may run a frame or two long, never short
+                raise RuntimeError(f"{rid}/{pfx} ep{ep_idx} {key}: {len(frames)} frames vs {len(keep)} rows")
+            cams[slot] = frames[: len(keep)]
+
+        m = min(len(keep), *(len(v) for v in cams.values() if v is not None))
+        state_t = torch.from_numpy(state[:m])
+        action_t = torch.from_numpy(action[:m])
+        ad = meta["action_dim"]
+        dim_mask = torch.zeros(self.max_action_dim, dtype=torch.bool)
+        dim_mask[:ad] = True
+        cam_mask = torch.tensor([True, cams["cam1"] is not None])
+        zero_img = None
+
+        for j in range(m):
+            avail = m - j
+            if avail >= self.chunk:
+                chunk_a = action_t[j : j + self.chunk]
+                is_pad = torch.zeros(self.chunk, dtype=torch.bool)
+            else:
+                chunk_a = torch.cat([action_t[j:], action_t[-1:].expand(self.chunk - avail, -1)])
+                is_pad = torch.zeros(self.chunk, dtype=torch.bool)
+                is_pad[avail:] = True
+            if cams["cam1"] is None and zero_img is None:
+                zero_img = torch.zeros_like(cams["cam0"][j])
+            yield {
+                "observation.images.cam0": cams["cam0"][j],
+                "observation.images.cam1": cams["cam1"][j] if cams["cam1"] is not None else zero_img,
+                "camera_mask": cam_mask,
+                "observation.state": _pad(state_t[j], self.max_state_dim, "state"),
+                "action": _pad(chunk_a, self.max_action_dim, "action"),
+                "action_dim_mask": dim_mask,
+                "action_is_pad": is_pad,
+                "embodiment_id": torch.tensor(spec.get("embodiment_id", si), dtype=torch.long),
+                "task": task,
+            }
+
+    # ------------------------------------------------------------------- iter
+
+    def __iter__(self):
+        wi = get_worker_info()
+        wid, nw = (wi.id, wi.num_workers) if wi is not None else (0, 1)
+        gid = self.rank * nw + wid
+        rng = random.Random(self.seed * 7919 + self.epoch * 104729 + gid)
+        weights = [s["weight"] for s in self.specs]
+
+        buf: list[dict] = []
+        while True:  # infinite stream; the training loop counts steps, not epochs
+            si = rng.choices(range(len(self.specs)), weights=weights, k=1)[0]
+            try:
+                meta = self._load_meta(si)
+                ep_idx, _, task = meta["episodes"][rng.randrange(len(meta["episodes"]))]
+                for item in self._stream_episode(si, ep_idx, task):
+                    if len(buf) < self.shuffle_buffer:
+                        buf.append(item)
+                        continue
+                    k = rng.randrange(len(buf))
+                    buf[k], item = item, buf[k]
+                    yield item
+            except Exception as e:  # network hiccup / bad episode: log and resample
+                import logging
+
+                logging.warning(f"stream: skipping {self.specs[si]['repo_id']} ep after {type(e).__name__}: {e}")

tinyvla_b200/src/tinyvla/data/wds_mix.py

@@ -0,0 +1,173 @@
+"""Weighted mixture loader over self-contained webdataset tars (see pack_wds.py).
+
+Yields the same batch schema as ShardSource + collate, so it drops into the
+training loop unchanged:
+
+    observation.images.cam0 / cam1   (B, 3, S, S) float [0,1]
+    camera_mask                      (B, 2) bool
+    observation.state                (B, max_state_dim)  already normalized
+    action                           (B, chunk, max_action_dim) already normalized
+    action_dim_mask                  (B, max_action_dim) bool
+    action_is_pad                    (B, chunk) bool
+    embodiment_id                    (B,) long
+    task_index                       (B,) long   -- per-dataset index; map via tasks.json
+
+webdataset is easy to hold wrong. The choices here, and why:
+
+  * ``resampled=True`` on every per-dataset stream: each worker samples SHARDS
+    WITH REPLACEMENT, forever. This is the one mode where none of the classic
+    failure modes exist — no "fewer shards than workers" (empty_check), no
+    silent worker starvation when shards % workers != 0, no need for
+    nodesplitter/workersplitter at all (independent draws are already i.i.d.
+    across workers and nodes), and a weighted mixture needs an infinite stream
+    per source anyway because "epoch" is undefined for a mixture.
+  * ``wds.RandomMix(..., probs=weights)`` composes the per-dataset streams; the
+    mixture is sampled per-sample, so batches interleave datasets (important:
+    per-dataset normalization is baked into the tars, the model must see mixed
+    batches, not dataset-blocked ones).
+  * ``.shuffle(buffer)`` per stream AFTER resampling: shard-level sampling alone
+    leaves frames of one episode adjacent; the buffer breaks that correlation.
+  * ``with_epoch(steps_per_epoch)`` bounds the infinite stream so DataLoader
+    epoch semantics (and persistent_workers restarts) still work.
+  * empty_check=False everywhere: with resampling a worker can never prove a
+    stream empty, the check only produces false positives.
+
+Deterministic-ish: pass seed; resampling uses it per worker.
+"""
+
+from __future__ import annotations
+
+import io
+import json
+from pathlib import Path
+
+import numpy as np
+import torch
+
+
+def _decode_and_pad(sample, image_size, max_state_dim, max_action_dim):
+    from PIL import Image
+
+    meta = np.load(io.BytesIO(sample["meta.npz"]))
+    S = int(meta["state_dim"])
+    A = int(meta["action_dim"])
+    if S > max_state_dim or A > max_action_dim:
+        raise ValueError(
+            f"sample has state {S}d / action {A}d, caps are {max_state_dim}/{max_action_dim} — "
+            f"raise the caps instead of truncating"
+        )
+
+    def img(key):
+        raw = sample.get(key, b"")
+        if not raw:
+            return None
+        im = Image.open(io.BytesIO(raw)).convert("RGB")
+        x = torch.from_numpy(np.asarray(im, np.uint8).copy()).permute(2, 0, 1).float().div_(255.0)
+        return x
+
+    cam0 = img("cam0.jpg")
+    cam1 = img("cam1.jpg")
+    if cam1 is None:
+        cam1 = torch.zeros_like(cam0)
+        cam_mask = torch.tensor([True, False])
+    else:
+        cam_mask = torch.tensor([True, True])
+
+    state = torch.zeros(max_state_dim)
+    state[:S] = torch.from_numpy(meta["state"])
+    chunk = meta["action_chunk"]
+    act = torch.zeros(chunk.shape[0], max_action_dim)
+    act[:, :A] = torch.from_numpy(chunk)
+    adm = torch.zeros(max_action_dim, dtype=torch.bool)
+    adm[:A] = True
+
+    return {
+        "observation.images.cam0": cam0,
+        "observation.images.cam1": cam1,
+        "camera_mask": cam_mask,
+        "observation.state": state,
+        "action": act,
+        "action_dim_mask": adm,
+        "action_is_pad": torch.from_numpy(meta["action_is_pad"]),
+        "embodiment_id": torch.tensor(int(meta["embodiment_id"]), dtype=torch.long),
+        "task_index": torch.tensor(int(meta["task_index"]), dtype=torch.long),
+    }
+
+
+def _stream(spec, image_size, max_state_dim, max_action_dim, shuffle_buffer, seed):
+    import webdataset as wds
+
+    d = Path(spec["dir"])
+    manifest = json.loads((d / "manifest.json").read_text())
+    tars = [str(d / s) for s in manifest["shards"]]
+    ds = wds.WebDataset(
+        tars,
+        resampled=True,          # shards sampled with replacement, infinite
+        shardshuffle=False,      # resampling already randomizes shard order
+        empty_check=False,
+        seed=seed,
+    )
+    if shuffle_buffer > 1:
+        ds = ds.shuffle(shuffle_buffer, initial=min(shuffle_buffer, 1000))
+    emb = spec.get("embodiment_id")
+    def mapper(sample):
+        out = _decode_and_pad(sample, image_size, max_state_dim, max_action_dim)
+        if emb is not None:  # override what was packed, if the mixture says so
+            out["embodiment_id"] = torch.tensor(int(emb), dtype=torch.long)
+        return out
+    return ds.map(mapper), manifest
+
+
+def make_mixture_loader(
+    specs: list[dict],            # [{"dir": ..., "weight": ..., "embodiment_id": ...}]
+    batch_size: int,
+    num_workers: int = 16,
+    image_size: int = 256,
+    max_state_dim: int = 64,
+    max_action_dim: int = 64,
+    shuffle_buffer: int = 4096,
+    steps_per_epoch: int = 1000,  # bounds the infinite stream per epoch
+    seed: int = 42,
+    prefetch_factor: int = 4,
+):
+    import webdataset as wds
+
+    # bound each infinite stream per worker; RandomMix (longest=False) then ends
+    # the epoch when the first stream is exhausted, keeping DataLoader epoch
+    # semantics intact. RandomMix itself takes no seed in wds 1.x — per-stream
+    # seeds provide the determinism.
+    per_worker = max(1, steps_per_epoch * batch_size // max(1, num_workers))
+    streams, manifests, weights = [], [], []
+    for i, spec in enumerate(specs):
+        ds, mf = _stream(spec, image_size, max_state_dim, max_action_dim,
+                         shuffle_buffer, seed + 1000 * i)
+        streams.append(ds.with_epoch(per_worker))
+        manifests.append(mf)
+        weights.append(float(spec.get("weight", mf["samples"])))
+
+    if len(streams) == 1:
+        mix = streams[0]
+    else:
+        total = sum(weights)
+        mix = wds.RandomMix(streams, probs=[w / total for w in weights])
+
+    # wds re-injects __key__/__url__ after EVERY pipeline stage, so stripping
+    # them inside the pipeline is futile — drop them at collate instead
+    from torch.utils.data import default_collate
+
+    def collate(items):
+        return default_collate([
+            {k: v for k, v in it.items() if not k.startswith("__")} for it in items
+        ])
+
+    loader = torch.utils.data.DataLoader(
+        mix,
+        batch_size=batch_size,
+        num_workers=num_workers,
+        pin_memory=True,
+        persistent_workers=num_workers > 0,
+        prefetch_factor=prefetch_factor if num_workers > 0 else None,
+        drop_last=True,
+        collate_fn=collate,
+    )
+    return loader, manifests

tinyvla_b200/src/tinyvla/modeling_tinyvla.py

@@ -0,0 +1,308 @@
+"""TinyVLA policy: dual-rate VLA on a single shared Qwen3.5 vision encoder.
+
+Slow path: image -> Qwen3.5 vision tower -> language stack (first 12 layers) +
+instruction -> 8-token semantic latent (cached at deployment, refreshed at
+0.5-2 Hz).
+
+Fast path: Qwen3.5 vision tower spatial tokens (64/camera, no language stack)
++ proprio + embodiment token + cached semantic latent -> ~40M flow-matching
+action expert -> 50-step action chunk.
+
+During training both paths run jointly (single-rate); staleness augmentation
+is added by the Stage-2 training loop.
+"""
+
+from __future__ import annotations
+
+from collections import deque
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from lerobot.configs import FeatureType
+from lerobot.policies.pretrained import PreTrainedPolicy
+from lerobot.utils.constants import (
+    ACTION,
+    OBS_LANGUAGE_ATTENTION_MASK,
+    OBS_LANGUAGE_TOKENS,
+    OBS_STATE,
+)
+
+from .configuration_tinyvla import TinyVLAConfig
+from .modules.embodiment import EmbodimentEmbedding, pad_vector
+from .modules.expert import FlowMatchingExpert
+from .modules.semantic import SemanticPath
+
+
+class TinyVLAPolicy(PreTrainedPolicy):
+    config_class = TinyVLAConfig
+    name = "tinyvla"
+
+    def __init__(self, config: TinyVLAConfig, dataset_stats=None, dataset_meta=None, **kwargs):
+        super().__init__(config)
+        self.config = config
+
+        self.image_keys = sorted(
+            k for k, v in config.input_features.items() if v.type is FeatureType.VISUAL
+        )
+        state_ft = config.input_features.get(OBS_STATE)
+        self.state_dim = state_ft.shape[0] if state_ft is not None else 0
+        self.action_dim = config.output_features[ACTION].shape[0]
+
+        d = config.expert_dim
+
+        self.semantic = SemanticPath(
+            model_name=config.lm_model_name,
+            num_layers=config.lm_num_layers,
+            num_readout=config.num_readout_tokens,
+            out_dim=d,
+            image_size=config.image_size,
+            freeze_lm=config.freeze_lm,
+            freeze_vision=config.freeze_vision_encoder,
+            morph_to_slow=config.morph_to_slow,
+            morph_dim=16,
+            morph_qwen=(config.conditioning == "morph_qwen"),
+            num_morph_readout=config.num_morph_readout,
+            vlm_native=config.vlm_native,
+            num_action_readout=config.num_action_readout,
+            max_state_dim=config.max_state_dim,
+            max_action_dim=config.max_action_dim,
+        )
+
+        # fast-path conditioning projections
+        self.spatial_proj = nn.Linear(self.semantic.visual_hidden_size, d)
+        self.camera_emb = nn.Embedding(config.max_cameras, d)
+        self.state_proj = nn.Linear(config.max_state_dim, d)
+        self.embodiment_emb = EmbodimentEmbedding(config.num_embodiments, d)
+        if config.conditioning == "morph":
+            from .modules.embodiment import MorphologyEncoder
+
+            self.morph_encoder = MorphologyEncoder(d, n_tokens=config.morph_tokens)
+        if config.use_demo_conditioning:
+            from .modules.embodiment import DemoConditioner
+
+            self.demo_encoder = DemoConditioner(
+                d,
+                visual_dim=self.semantic.visual_hidden_size,
+                action_dim=config.max_action_dim,
+                n_tokens_per_example=config.demo_tokens_per_example,
+                n_visual_groups=config.demo_visual_groups,
+                n_action_keys=config.demo_action_keys,
+                hidden_mult=config.demo_hidden_mult,
+            )
+
+        self.spatial_head = (
+            nn.Linear(d, config.spatial_vocab) if config.spatial_loss_weight > 0 else None
+        )
+
+        self.expert = FlowMatchingExpert(
+            action_dim=config.max_action_dim,
+            chunk_size=config.chunk_size,
+            d_model=d,
+            n_layers=config.expert_layers,
+            n_heads=config.expert_heads,
+            min_period=config.min_period,
+            max_period=config.max_period,
+        )
+
+        self._action_queue: deque = deque()
+
+    def reset(self):
+        self._action_queue.clear()
+
+    def get_optim_params(self):
+        return [p for p in self.parameters() if p.requires_grad]
+
+    # ---------------------------------------------------------------- features
+
+    @staticmethod
+    def _last_obs(x: torch.Tensor, ndim: int) -> torch.Tensor:
+        """Drop the n_obs_steps dim if present: (B, S, ...) -> (B, ...)."""
+        return x[:, -1] if x.dim() == ndim + 1 else x
+
+    def _encode_spatial(self, batch: dict, feats: dict | None = None) -> torch.Tensor:
+        """Per-camera spatial tokens for the fast path: (B, num_cams*64, d).
+
+        `feats` lets the caller pass vision-tower outputs computed once per step
+        (see `_vision_features`) instead of re-running the tower here.
+        """
+        cond = []
+        for i, key in enumerate(self.image_keys):
+            tokens = (feats or {}).get(key)
+            if tokens is None:
+                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 _vision_features(self, batch: dict, keys) -> dict:
+        """One vision-tower pass per camera, shared by the fast and slow paths."""
+        return {k: self.semantic.encode_image(self._last_obs(batch[k], 4)) for k in keys}
+
+    def _semantic_latent(self, batch: dict, image=None, image_embeds=None) -> torch.Tensor:
+        """Slow-path latent: (B, K, d), cacheable, conditions the fast path.
+
+        vlm_native=True makes this latent RICH — the LM sequence also carries the
+        robot description and the demo (image, action) exemplars, so the readout
+        tokens can attend over them. The fast path is unchanged: it still gets its
+        own spatial tokens + proprio + expert, just a better-informed latent.
+        """
+        img = image if image is not None else self._last_obs(batch[self.image_keys[0]], 4)
+        if self.config.vlm_native:
+            state = None
+            if self.state_dim:
+                state = pad_vector(self._last_obs(batch[OBS_STATE], 2), self.config.max_state_dim)
+            return self.semantic.forward_unified(
+                image=img,
+                lang_tokens=batch[OBS_LANGUAGE_TOKENS],
+                lang_mask=batch[OBS_LANGUAGE_ATTENTION_MASK],
+                state=state,
+                # demo_only: no robot description — demos are the sole identity cue
+                robot_ids=None if self.config.demo_only else batch.get("morph_text_ids"),
+                robot_mask=None if self.config.demo_only else batch.get("morph_text_mask"),
+                demo_images=batch.get("support_images"),
+                demo_actions=batch.get("support_actions"),
+                n_demo_action_keys=self.config.demo_action_keys_lm,
+            )
+        morph = batch.get("morphology") if self.config.morph_to_slow else None
+        return self.semantic(
+            img,
+            batch[OBS_LANGUAGE_TOKENS],
+            batch[OBS_LANGUAGE_ATTENTION_MASK],
+            morphology=morph,
+            image_embeds=image_embeds,
+        )
+
+    def _conditioning(self, batch: dict) -> torch.Tensor:
+        b = batch[self.image_keys[0]].shape[0]
+        device = batch[self.image_keys[0]].device
+
+        # Run the shared Qwen vision tower ONCE per camera per step. Previously
+        # cam0 went through it twice: here for the spatial tokens, and again
+        # inside the LM call in `_semantic_latent`. The tower is ~45 GFLOPs per
+        # camera, i.e. a quarter of the forward pass thrown away every step.
+        feats = None
+        if self.config.use_spatial_tokens:
+            feats = self._vision_features(batch, self.image_keys)
+        elif self.config.use_semantic_latent and not self.config.vlm_native:
+            feats = self._vision_features(batch, self.image_keys[:1])
+
+        parts = []
+        if self.config.use_semantic_latent:
+            latent = batch.get("semantic_latent")
+            if latent is None:
+                if "latent_image" in batch:  # staleness aug: latent from an older frame,
+                    latent = self._semantic_latent(batch, image=batch["latent_image"])
+                else:                        # ...so the cached cam0 features do NOT apply
+                    cam0 = None if (feats is None or self.config.vlm_native) else feats[self.image_keys[0]]
+                    latent = self._semantic_latent(batch, image_embeds=cam0)
+            parts.append(latent)
+        if self.config.use_spatial_tokens:
+            parts.append(self._encode_spatial(batch, feats))
+        if self.state_dim:
+            state = pad_vector(self._last_obs(batch[OBS_STATE], 2), self.config.max_state_dim)
+            parts.append(self.state_proj(state)[:, None, :])
+        if self.config.demo_only:
+            pass  # no identity token at all — demos in the LM latent are the only cue
+        elif self.config.conditioning == "morph":
+            desc = batch.get("morphology")
+            if desc is None:
+                from .modules.embodiment import MORPH_DIM
+
+                desc = torch.zeros(b, MORPH_DIM, device=device)
+            parts.append(self.morph_encoder(desc.float()))
+        elif self.config.conditioning == "morph_qwen":
+            ids = batch.get("morph_text_ids")
+            mask = batch.get("morph_text_mask")
+            if ids is None:
+                ids = torch.zeros(b, 1, dtype=torch.long, device=device)
+                mask = torch.ones(b, 1, dtype=torch.bool, device=device)
+            parts.append(self.semantic.encode_robot_text(ids, mask))
+        else:
+            emb_id = batch.get("embodiment_id")
+            if emb_id is None:
+                emb_id = torch.zeros(b, dtype=torch.long, device=device)
+            parts.append(self.embodiment_emb(emb_id))
+
+        if self.config.use_demo_conditioning:
+            sup_imgs = batch.get("support_images")  # (B, K, 3, H, W)
+            sup_acts = batch.get("support_actions")  # (B, K, T, A)
+            if sup_imgs is not None and sup_acts is not None:
+                bb, k = sup_imgs.shape[:2]
+                flat_imgs = sup_imgs.reshape(bb * k, *sup_imgs.shape[2:])
+                spatial = self.semantic.encode_image(flat_imgs)  # (B*K, 64, visual_dim)
+                # keep the spatial structure — the encoder does grouped pooling
+                spatial = spatial.view(bb, k, spatial.shape[1], spatial.shape[2])
+                parts.append(self.demo_encoder(spatial, sup_acts))
+        return torch.cat(parts, dim=1)
+
+    # ---------------------------------------------------------------- training
+
+    def forward(self, batch: dict) -> tuple[torch.Tensor, dict]:
+        cond = self._conditioning(batch)
+
+        actions = pad_vector(batch[ACTION], self.config.max_action_dim)  # (B, T, A)
+        b, t, a = actions.shape
+        device = actions.device
+
+        beta = torch.distributions.Beta(self.config.flow_beta_alpha, self.config.flow_beta_beta)
+        time = beta.sample((b,)).to(device) * 0.999 + 0.001
+        noise = torch.randn_like(actions)
+        x_t = time[:, None, None] * noise + (1 - time[:, None, None]) * actions
+        target = noise - actions
+
+        pred = self.expert(x_t, time, cond)
+
+        loss_mask = torch.ones(b, t, a, device=device)
+        if "action_dim_mask" in batch:  # per-sample real dims (mixture training)
+            loss_mask = loss_mask * batch["action_dim_mask"][:, None, :].float()
+        else:
+            loss_mask[:, :, self.action_dim :] = 0  # padded action dims
+        if "action_is_pad" in batch:  # episode-end padding from delta_timestamps
+            loss_mask = loss_mask * (~batch["action_is_pad"])[:, :, None].float()
+
+        loss = (F.mse_loss(pred, target, reduction="none") * loss_mask).sum() / loss_mask.sum()
+        info = {"flow_loss": loss.item()}
+
+        if (
+            self.spatial_head is not None
+            and "spatial_cell" in batch
+            and self.config.use_spatial_tokens
+        ):
+            cells = batch["spatial_cell"]
+            labeled = cells >= 0
+            if labeled.any():
+                logits = self.spatial_head(self._cam0_proj.mean(dim=1))  # (B, vocab)
+                ce = F.cross_entropy(logits[labeled], cells[labeled])
+                loss = loss + self.config.spatial_loss_weight * ce
+                info["spatial_ce"] = ce.item()
+
+        info["loss"] = loss.item()
+        return loss, info
+
+    # --------------------------------------------------------------- inference
+
+    @torch.no_grad()
+    def predict_action_chunk(self, batch: dict) -> torch.Tensor:
+        cond = self._conditioning(batch)
+        b = cond.shape[0]
+        device = cond.device
+
+        x = torch.randn(b, self.config.chunk_size, self.config.max_action_dim, device=device)
+        n = self.config.num_denoise_steps
+        dt = 1.0 / n
+        for i in range(n):
+            t = torch.full((b,), 1.0 - i * dt, device=device)
+            v = self.expert(x, t, cond)
+            x = x - dt * v
+        return x[:, :, : self.action_dim]
+
+    @torch.no_grad()
+    def select_action(self, batch: dict) -> torch.Tensor:
+        if len(self._action_queue) == 0:
+            chunk = self.predict_action_chunk(batch)  # (B, T, A)
+            self._action_queue.extend(chunk.transpose(0, 1)[: self.config.n_action_steps])
+        return self._action_queue.popleft()

tinyvla_b200/src/tinyvla/modules/expert.py

@@ -0,0 +1,170 @@
+"""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
+import torch.nn.functional as F
+
+
+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 _MHA(nn.Module):
+    """Drop-in replacement for torch's batch-first multi-head attention, routed
+    through F.scaled_dot_product_attention (flash / mem-efficient kernels, and
+    far friendlier to torch.compile than the eager fallback).
+
+    Parameter names are kept identical (in_proj_weight / in_proj_bias /
+    out_proj.weight / out_proj.bias) so existing checkpoints load unchanged.
+    `key_padding_mask` keeps torch semantics: True = ignore that key.
+    """
+
+    def __init__(self, dim: int, n_heads: int):
+        super().__init__()
+        assert dim % n_heads == 0
+        self.dim, self.n_heads, self.head_dim = dim, n_heads, dim // n_heads
+        self.in_proj_weight = nn.Parameter(torch.empty(3 * dim, dim))
+        self.in_proj_bias = nn.Parameter(torch.zeros(3 * dim))
+        self.out_proj = nn.Linear(dim, dim)
+        nn.init.xavier_uniform_(self.in_proj_weight)
+        nn.init.zeros_(self.out_proj.bias)
+
+    def forward(self, query, key, value, key_padding_mask=None, need_weights=False):
+        b, lq, d = query.shape
+        w_q, w_k, w_v = self.in_proj_weight.split(d, dim=0)
+        b_q, b_k, b_v = self.in_proj_bias.split(d, dim=0)
+        q = F.linear(query, w_q, b_q)
+        k = F.linear(key, w_k, b_k)
+        v = F.linear(value, w_v, b_v)
+
+        def split(t):
+            return t.view(b, -1, self.n_heads, self.head_dim).transpose(1, 2)
+
+        q, k, v = split(q), split(k), split(v)
+        attn_mask = None
+        if key_padding_mask is not None:
+            # torch convention: True = ignore. SDPA bool mask: True = attend.
+            attn_mask = (~key_padding_mask)[:, None, None, :]
+        o = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask)
+        o = o.transpose(1, 2).reshape(b, lq, d)
+        return self.out_proj(o), None
+
+
+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 = _MHA(dim, n_heads)
+        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 = _MHA(dim, n_heads)
+        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))

tinyvla_b200/src/tinyvla/modules/semantic.py

@@ -0,0 +1,420 @@
+"""Slow semantic path: Qwen3.5-0.8B VLM over [image | instruction | readout slots].
+
+Qwen3.5-0.8B has a native vision tower (12-layer ViT, patch 16, spatial merge 2:
+a 256x256 image becomes 64 visual tokens), so the slow path consumes the raw
+image directly — no hand-rolled projector. The language stack is truncated to
+the first `num_layers` blocks (L/2 layer-skip, SmolVLA-style). The semantic
+latent is the hidden state of K trailing readout slots (pad tokens after the
+instruction; causal attention lets them aggregate the full prefix), projected
+to the expert dimension. At deployment this runs at 0.5-2 Hz and the latent is
+cached between refreshes.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class SemanticPath(nn.Module):
+    def __init__(
+        self,
+        model_name: str = "Qwen/Qwen3.5-0.8B",
+        num_layers: int = 12,
+        num_readout: int = 8,
+        out_dim: int = 512,
+        image_size: int = 256,
+        freeze_lm: bool = True,
+        freeze_vision: bool = True,
+        morph_to_slow: bool = False,
+        morph_dim: int = 16,
+        morph_qwen: bool = False,
+        num_morph_readout: int = 4,
+        vlm_native: bool = False,
+        num_action_readout: int = 16,
+        max_state_dim: int = 16,
+        max_action_dim: int = 8,
+    ):
+        super().__init__()
+        from transformers import AutoModelForImageTextToText, AutoProcessor
+
+        vlm = AutoModelForImageTextToText.from_pretrained(model_name, dtype=torch.float32)
+        model = vlm.model  # Qwen3_5Model: visual tower + language stack
+        # L/2 layer-skip: keep the first num_layers language blocks; lm_head unused
+        model.language_model.layers = model.language_model.layers[:num_layers]
+        model.language_model.config.num_hidden_layers = num_layers
+        self.vlm = model
+        cfg = vlm.config
+        self.hidden_size = cfg.text_config.hidden_size
+        self.image_token_id = cfg.image_token_id
+        self.vision_start_id = cfg.vision_start_token_id
+        self.vision_end_id = cfg.vision_end_token_id
+
+        proc = AutoProcessor.from_pretrained(model_name)
+        self.pad_token_id = proc.tokenizer.pad_token_id or proc.tokenizer.eos_token_id
+        ip = proc.image_processor
+        self.patch_size = ip.patch_size  # 16
+        self.merge_size = ip.merge_size  # 2
+        self.temporal_patch_size = ip.temporal_patch_size  # 2
+        self.register_buffer(
+            "pixel_mean", torch.tensor(ip.image_mean).view(1, 3, 1, 1), persistent=False
+        )
+        self.register_buffer(
+            "pixel_std", torch.tensor(ip.image_std).view(1, 3, 1, 1), persistent=False
+        )
+        self.image_size = image_size
+        grid = image_size // self.patch_size
+        self.num_image_tokens = (grid // self.merge_size) ** 2  # 64 @ 256px
+
+        self.out_proj = nn.Linear(self.hidden_size, out_dim)
+        self.num_readout = num_readout
+        self._fixed_fn = None  # lazily compiled fixed-grid tower path
+        self.morph_to_slow = morph_to_slow
+        if morph_to_slow:
+            # inject robot morphology into the slow LM prefix so the semantic
+            # latent is morphology-aware ("grasp X" differs per embodiment)
+            self.morph_to_lm = nn.Linear(morph_dim, self.hidden_size)
+
+        # LLaVA-style unified mode: every modality becomes tokens inside the LM
+        # sequence, and the action policy reads latent tokens off the LM output.
+        self.vlm_native = vlm_native
+        self.num_action_readout = num_action_readout
+        if vlm_native:
+            self.state_to_lm = nn.Linear(max_state_dim, self.hidden_size)
+            self.action_to_lm = nn.Linear(max_action_dim, self.hidden_size)
+            self.action_readout_proj = nn.Linear(self.hidden_size, out_dim)
+            self.readout_emb = nn.Parameter(torch.zeros(num_action_readout, self.hidden_size))
+            nn.init.normal_(self.readout_emb, std=0.02)
+
+        self.morph_qwen = morph_qwen
+        self.num_morph_readout = num_morph_readout
+        if morph_qwen:
+            # encode the robot TEXT description through the SAME Qwen language
+            # stack (shared weights, text-only, no image) instead of an MLP over
+            # numeric fields — lets the (unfrozen) LM's language understanding
+            # do the morphology encoding instead of a from-scratch projector.
+            self.morph_out_proj = nn.Linear(self.hidden_size, out_dim)
+
+        if freeze_lm:
+            for p in self.vlm.language_model.parameters():
+                p.requires_grad = False
+        # vocab embeddings stay frozen either way (254M tied params, negligible gain)
+        self.vlm.get_input_embeddings().weight.requires_grad = False
+        if freeze_vision:
+            for p in self.vlm.visual.parameters():
+                p.requires_grad = False
+
+    @property
+    def visual_hidden_size(self) -> int:
+        return self.vlm.visual.config.out_hidden_size
+
+    def encode_image(self, images: torch.Tensor) -> torch.Tensor:
+        """Qwen vision tower spatial tokens: (B, 3, H, W) -> (B, 64, out_hidden).
+
+        Shared encoder for both paths; the fast path cross-attends to these
+        directly, bypassing the language stack. The SAME tensor can be handed to
+        `forward(..., image_embeds=...)` so cam0 goes through the tower once per
+        step instead of twice (once here, once inside the LM call).
+
+        Uses a dense fixed-grid fast path: every image here is the same
+        (1, g, g) grid, so the HF packed-varlen tower (which runs B separate
+        SDPA calls PER LAYER and rebuilds pos/rotary tables in Python on every
+        call) is replaced by one batched SDPA per layer over (B, L, D), with the
+        per-image positional constants computed once and cached. Numerically
+        equivalent — per-image attention with equal lengths IS batched attention.
+        """
+        b = images.shape[0]
+        pixel_values, grid_thw = self._patchify(images)
+        fn = self._fixed_fn
+        if fn is None:
+            import os
+
+            if os.environ.get("TINYVLA_NO_COMPILE_VIT"):
+                fn = self._encode_image_fixed
+            else:
+                # static shapes -> compile is safe and worth ~2x on top of the
+                # dense rewrite; recompiles once per new batch size
+                fn = torch.compile(self._encode_image_fixed, dynamic=False)
+            self._fixed_fn = fn
+        try:
+            return fn(b, pixel_values)
+        except Exception:
+            if fn is not self._encode_image_fixed:  # compile blew up -> eager forever
+                self._fixed_fn = self._encode_image_fixed
+                return self._encode_image_fixed(b, pixel_values)
+            raise
+
+    def _encode_image_hf(self, images: torch.Tensor) -> torch.Tensor:
+        """Reference path through the stock HF tower (kept for verification)."""
+        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 _fixed_grid_ctx(self, device):
+        """Per-image positional constants for the fixed (1, g, g) grid, computed
+        ONCE via the tower's own functions (so the layout can't drift) and cached."""
+        g = self.image_size // self.patch_size
+        ctx = getattr(self, "_fixed_ctx", None)
+        if ctx is not None and ctx[0] == (g, str(device)):
+            return ctx[1]
+        visual = self.vlm.visual
+        grid1 = torch.tensor([[1, g, g]], device=device)
+        with torch.no_grad():
+            pos = visual.fast_pos_embed_interpolate(grid1)  # (g*g, hidden)
+            rot = visual.rot_pos_emb(grid1)                 # (g*g, head_dim/2)
+        emb = torch.cat((rot, rot), dim=-1)
+        val = (pos.detach(), emb.cos().detach(), emb.sin().detach())
+        self._fixed_ctx = ((g, str(device)), val)
+        return val
+
+    def _encode_image_fixed(self, b: int, pixel_values: torch.Tensor) -> torch.Tensor:
+        from transformers.models.qwen3_5.modeling_qwen3_5 import rotate_half
+
+        visual = self.vlm.visual
+        g = self.image_size // self.patch_size
+        L = g * g
+        x = visual.patch_embed(pixel_values.to(self.pixel_mean.dtype))  # (B*L, D)
+        pos, cos, sin = self._fixed_grid_ctx(x.device)
+        d = x.shape[-1]
+        x = x.view(b, L, d) + pos[None].to(x.dtype)
+        # (1, 1, L, hd) fp32 — same upcast-to-fp32 rotary as the HF path
+        cosf = cos[None, None].float()
+        sinf = sin[None, None].float()
+        for blk in visual.blocks:
+            attn = blk.attn
+            h = blk.norm1(x)
+            qkv = attn.qkv(h).view(b, L, 3, attn.num_heads, attn.head_dim)
+            q, k, v = qkv.permute(2, 0, 3, 1, 4).unbind(0)  # (B, nh, L, hd)
+            qf, kf = q.float(), k.float()
+            q = ((qf * cosf) + (rotate_half(qf) * sinf)).to(q.dtype)
+            k = ((kf * cosf) + (rotate_half(kf) * sinf)).to(k.dtype)
+            o = F.scaled_dot_product_attention(q, k, v, scale=attn.scaling)
+            o = o.transpose(1, 2).reshape(b, L, d)
+            x = x + attn.proj(o)
+            x = x + blk.mlp(blk.norm2(x))
+        merged = visual.merger(x.reshape(b * L, d))  # (B*L/m^2, out_hidden)
+        return merged.reshape(b, self.num_image_tokens, -1)
+
+    def _grid_thw(self, b: int, device) -> torch.Tensor:
+        g = self.image_size // self.patch_size
+        return torch.tensor([[1, g, g]], device=device).expand(b, -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 | None,  # (B, 3, H, W) in [0, 1]; may be None with image_embeds
+        lang_tokens: torch.Tensor,  # (B, L)
+        lang_mask: torch.Tensor,  # (B, L) bool, True = valid
+        morphology: torch.Tensor | None = None,  # (B, morph_dim) for morph_to_slow
+        image_embeds: torch.Tensor | None = None,  # (B, n_img, C) already through the tower
+    ) -> torch.Tensor:
+        """Returns semantic latent (B, num_readout, out_dim).
+
+        `image_embeds` lets the caller reuse the vision-tower output the fast path
+        already computed for the same frame — the tower is ~45 GFLOPs/camera and
+        cam0 was otherwise encoded twice per training step.
+        """
+        if image_embeds is not None:
+            b, device = image_embeds.shape[0], image_embeds.device
+            pixel_values, grid_thw = None, self._grid_thw(b, device)
+        else:
+            b, device = images.shape[0], images.device
+            pixel_values, grid_thw = self._patchify(images)
+        n_img = self.num_image_tokens
+
+        # <vision_start> <image>*n <vision_end> instruction <pad>*K
+        ids = torch.cat(
+            [
+                torch.full((b, 1), self.vision_start_id, dtype=torch.long, device=device),
+                torch.full((b, n_img), self.image_token_id, dtype=torch.long, device=device),
+                torch.full((b, 1), self.vision_end_id, dtype=torch.long, device=device),
+                lang_tokens,
+                torch.full((b, self.num_readout), self.pad_token_id, dtype=torch.long, device=device),
+            ],
+            dim=1,
+        )
+        attn = torch.cat(
+            [
+                torch.ones(b, n_img + 2, dtype=torch.bool, device=device),
+                lang_mask.bool(),
+                torch.ones(b, self.num_readout, dtype=torch.bool, device=device),
+            ],
+            dim=1,
+        )
+        mm_type = (ids == self.image_token_id).int()  # text 0, image 1
+
+        use_embeds = image_embeds is not None or (self.morph_to_slow and morphology is not None)
+        if use_embeds:
+            # embed tokens, scatter image features, optionally add morphology to the
+            # readout slots, then run the LM on inputs_embeds (keeps M-RoPE via
+            # mm_token_type_ids). This is also the path that consumes precomputed
+            # `image_embeds` — identical maths, one fewer vision-tower pass.
+            inputs_embeds = self.vlm.get_input_embeddings()(ids)
+            if image_embeds is not None:
+                img = image_embeds.reshape(-1, image_embeds.shape[-1])
+            else:
+                img = self.vlm.get_image_features(pixel_values.to(self.pixel_mean.dtype), grid_thw)
+                img = torch.cat(img.pooler_output, dim=0)  # (B*n_img, hidden)
+            img = img.to(inputs_embeds.dtype)
+            imask = (ids == self.image_token_id).unsqueeze(-1).expand_as(inputs_embeds)
+            inputs_embeds = inputs_embeds.masked_scatter(imask, img)
+            if self.morph_to_slow and morphology is not None:
+                morph_tok = self.morph_to_lm(morphology.to(inputs_embeds.dtype))  # (B, hidden)
+                inputs_embeds[:, -self.num_readout:] = inputs_embeds[:, -self.num_readout:] + morph_tok[:, None, :]
+            # With inputs_embeds and no input_ids the model cannot recompute
+            # M-RoPE and falls back to a CACHED rope_deltas from a previous call
+            # (crashes on a different batch size; silently wrong on the same one).
+            # Compute the 3D position ids explicitly from the token layout.
+            position_ids, _ = self.vlm.get_rope_index(
+                ids, image_grid_thw=grid_thw, attention_mask=attn, mm_token_type_ids=mm_type,
+            )
+            out = self.vlm(
+                inputs_embeds=inputs_embeds, image_grid_thw=grid_thw, attention_mask=attn,
+                mm_token_type_ids=mm_type, position_ids=position_ids, use_cache=False,
+            )
+        else:
+            out = self.vlm(
+                input_ids=ids, pixel_values=pixel_values.to(self.pixel_mean.dtype),
+                image_grid_thw=grid_thw, attention_mask=attn, mm_token_type_ids=mm_type,
+                use_cache=False,
+            )
+        readout_h = out.last_hidden_state[:, -self.num_readout :]
+        return self.out_proj(readout_h)
+
+    def encode_robot_text(self, desc_ids: torch.Tensor, desc_mask: torch.Tensor) -> torch.Tensor:
+        """Encode a robot-description text through the SAME (shared-weight) Qwen
+        language stack, text-only (no image, no M-RoPE needed). Trailing pad-token
+        readout slots aggregate the description via causal attention; their hidden
+        states are projected to `num_morph_readout` conditioning tokens.
+
+        desc_ids/desc_mask: (B, L) tokenized "Robot: X, a N-DOF arm on ..." text.
+        Returns (B, num_morph_readout, out_dim).
+        """
+        b = desc_ids.shape[0]
+        device = desc_ids.device
+        ids = torch.cat(
+            [desc_ids, torch.full((b, self.num_morph_readout), self.pad_token_id,
+                                   dtype=torch.long, device=device)],
+            dim=1,
+        )
+        attn = torch.cat(
+            [desc_mask.bool(), torch.ones(b, self.num_morph_readout, dtype=torch.bool, device=device)],
+            dim=1,
+        )
+        out = self.vlm.language_model(
+            input_ids=ids, attention_mask=attn, use_cache=False,
+        )
+        readout_h = out.last_hidden_state[:, -self.num_morph_readout:]
+        return self.morph_out_proj(readout_h)
+
+    def forward_unified(
+        self,
+        image: torch.Tensor,                 # (B, 3, H, W) current observation
+        lang_tokens: torch.Tensor,           # (B, L) task instruction
+        lang_mask: torch.Tensor,             # (B, L)
+        state: torch.Tensor | None = None,   # (B, max_state_dim) proprio
+        robot_ids: torch.Tensor | None = None,    # (B, Lr) robot description text
+        robot_mask: torch.Tensor | None = None,
+        demo_images: torch.Tensor | None = None,  # (B, K, 3, H, W)
+        demo_actions: torch.Tensor | None = None, # (B, K, T, A)
+        n_demo_action_keys: int = 4,
+    ) -> torch.Tensor:
+        """LLaVA-style unified pass: ALL modalities become tokens in one LM
+        sequence, and the policy reads `num_action_readout` latent tokens off the
+        LM output. Demo exemplars live INSIDE the attention stream, so the readout
+        tokens can attend to them selectively (the thing that "conditioning-set"
+        placement could not do).
+
+        Sequence layout:
+          [<vis> main-image <-vis>] [instruction] [robot description]
+          [ <vis> demo_k image <-vis> demo_k action-keyframe tokens ]*K
+          [state token] [action readout slots]
+        Returns (B, num_action_readout, out_dim).
+        """
+        b = image.shape[0]
+        device = image.device
+        emb = self.vlm.get_input_embeddings()
+        n_img = self.num_image_tokens
+
+        # ---- encode all images (main + demos) in one vision pass -------------
+        imgs = [image]
+        k = 0
+        if demo_images is not None:
+            k = demo_images.shape[1]
+            imgs.append(demo_images.reshape(b * k, *demo_images.shape[2:]))
+        all_imgs = torch.cat(imgs, dim=0)
+        pixel_values, grid_thw = self._patchify(all_imgs)
+        feats = self.vlm.get_image_features(pixel_values.to(self.pixel_mean.dtype), grid_thw)
+        feats = torch.cat(feats.pooler_output, dim=0)          # (B*(1+K)*n_img, hidden)
+        feats = feats.view(b * (1 + k), n_img, -1)
+        main_feat = feats[:b]                                   # (B, n_img, hidden)
+        demo_feat = feats[b:].view(b, k, n_img, -1) if k else None
+
+        seq, mm, msk = [], [], []
+
+        def push(x, is_img=False, mask=None):
+            seq.append(x)
+            n = x.shape[1]
+            mm.append(torch.full((b, n), 1 if is_img else 0, dtype=torch.int, device=device))
+            msk.append(mask if mask is not None else torch.ones(b, n, dtype=torch.bool, device=device))
+
+        vs = emb(torch.full((b, 1), self.vision_start_id, dtype=torch.long, device=device))
+        ve = emb(torch.full((b, 1), self.vision_end_id, dtype=torch.long, device=device))
+
+        push(vs); push(main_feat, is_img=True); push(ve)
+        push(emb(lang_tokens), mask=lang_mask.bool())
+        if robot_ids is not None:
+            push(emb(robot_ids), mask=robot_mask.bool())
+        if demo_feat is not None:
+            keys = torch.linspace(0, demo_actions.shape[2] - 1, n_demo_action_keys,
+                                  device=device).long()
+            for i in range(k):
+                push(vs); push(demo_feat[:, i], is_img=True); push(ve)
+                act_tok = self.action_to_lm(demo_actions[:, i, keys].to(main_feat.dtype))
+                push(act_tok)                                    # (B, n_keys, hidden)
+        if state is not None:
+            push(self.state_to_lm(state.to(main_feat.dtype))[:, None, :])
+        push(self.readout_emb[None].expand(b, -1, -1).to(main_feat.dtype))
+
+        inputs_embeds = torch.cat(seq, dim=1)
+        mm_type = torch.cat(mm, dim=1)
+        attn = torch.cat(msk, dim=1)
+        # image_grid_thw must cover only the images actually present in-sequence
+        grid_for_seq = grid_thw[: b * (1 + k)]
+
+        out = self.vlm(
+            inputs_embeds=inputs_embeds,
+            image_grid_thw=grid_for_seq,
+            attention_mask=attn,
+            mm_token_type_ids=mm_type,
+            use_cache=False,
+        )
+        readout = out.last_hidden_state[:, -self.num_action_readout:]
+        return self.action_readout_proj(readout)