feat(lewm): enable gpu parallel rollout validation

docs/lewm-imf-experiment-guide.md

@@ -0,0 +1,471 @@
+# feat-lewm-imf-fusion 实验操作指南
+
+适用 worktree：`/home/droid/project/roboimi/.worktrees/feat-lewm-imf-fusion`
+
+## 0. 先记住当前常用 recipe
+
+当前这条分支最常用的训练/验证配方，直接参考：
+`experiment_suites/2026-04-21-lewm-fromscratch-old9-epoch50-roll5-val-20260421-153037/`
+
+核心约定：
+- agent：`lewm_resnet_query_imf_attnres`
+- from scratch：`train.pretrained_ckpt=null`，`agent.lewm_pretrained_ckpt=null`
+- 训练：`batch_size=32`，`lr=1e-4`，`max_steps=109350`，`save_freq=10000`
+- 数值验证：`train.val_split=0.0` + `train.val_episode_indices=[100]`
+- held-out numeric validation：`train.action_mse_val_freq_epochs=1`
+- rollout validation：`train.rollout_val_freq_epochs=5`，`train.rollout_num_episodes=10`
+- SwanLab：`train.use_swanlab=true`，project=`roboimi-vla`
+
+---
+
+## 1. 分支结构与关键文件
+
+| 路径 | 作用 |
+| --- | --- |
+| `roboimi/demos/vla_scripts/train_vla.py` | 主训练入口；负责数据集、checkpoint、数值验证、训练期 rollout 验证、SwanLab |
+| `roboimi/demos/vla_scripts/eval_vla.py` | 单次 rollout / 离线验证入口；支持 headless、summary、trajectory image/video artifact |
+| `roboimi/vla/conf/config.yaml` | 全局 Hydra 配置；训练默认值都在这里 |
+| `roboimi/vla/conf/eval/eval.yaml` | eval 默认配置；`eval.ckpt_path`、`eval.num_episodes`、artifact 开关都在这里 |
+| `roboimi/vla/conf/agent/lewm_resnet_query_imf_attnres.yaml` | 本分支最常用 agent；LeWM query fusion + IMF AttnRes head |
+| `roboimi/vla/conf/backbone/lewm_resnet_query_fusion.yaml` | LeWM 多视角 ResNet query fusion backbone 配置 |
+| `roboimi/vla/agent_imf.py` | `IMFVLAAgent` 实现；one-step IMF 推理、LeWM loss、LeWM 预训练组件加载 |
+| `roboimi/vla/data/simpe_robot_dataset.py` | HDF5 懒加载数据集；也负责 `episode_indices` 过滤 |
+| `roboimi/vla/scripts/calculate_stats.py` | 重算 `dataset_stats.pkl` |
+| `experiment_suites/2026-04-21-lewm-fromscratch-old9-epoch50-roll5-val-20260421-153037/` | 当前最常用 suite；manifest、notes、launch log、local 启动脚本都在这里 |
+
+补充：
+- 本分支常用 run name 形如 `lewmimf-q08-ph08-ex08-emb384-l12-fromscratch-epoch50-step109350-5090g0-20260421-153037`
+- `q08/ph16/ex08` 这类后缀分别对应 `agent.lewm_query_offsets`、`agent.pred_horizon`、`agent.num_action_steps`
+
+---
+
+## 2. 三台机器与环境
+
+| 机器 | GPU | repo / worktree | Python | 常用数据集路径 |
+| --- | --- | --- | --- | --- |
+| 本地 `droid-z790eagleax` | 1× RTX 5090 32GB | `/home/droid/project/roboimi/.worktrees/feat-lewm-imf-fusion` | `/home/droid/.conda/envs/roboimi/bin/python` | `/home/droid/project/diana_sim/sim_transfer` |
+| 5880 节点 `100.73.14.65` | 2× RTX 5880 Ada 48GB | `/home/droid/roboimi_suite_20260416_lewm_imf_fusion` | `/home/droid/miniforge3/envs/roboimi/bin/python` | `/home/droid/sim_dataset/sim_transfer` |
+| L20 节点 `100.119.99.14` | 8× NVIDIA L20 46GB | `/data/roboimi_suite_20260416_lewm_imf_fusion` | `/home/droid/miniforge3/envs/roboimi/bin/python` | `/data/simtransfer/current` |
+
+连接：
+- 5880：`ssh droid@100.73.14.65`
+- L20：`ssh droid@100.119.99.14`
+
+经验规则：
+- 本地 5090：适合单条 smoke / 小规模主跑 / 本地调参
+- 5880：适合 2 条并行主跑
+- L20：适合大 grid；数据和 run 建议都放 `/data`
+
+---
+
+## 3. 训练流怎么走
+
+`train_vla.py` 的实际流程：
+
+1. 读取 Hydra 配置并打印完整 cfg
+2. 通过 `build_train_val_datasets()` 构建 train/val dataset
+3. 用 `DataLoader` 建 train/val loader
+4. 从 `dataset_dir/dataset_stats.pkl` 读取归一化统计
+5. instantiate `IMFVLAAgent`
+6. 可选加载：
+   - `train.pretrained_ckpt`
+   - `train.resume_ckpt`
+   - `agent.lewm_pretrained_ckpt`
+7. 训练循环里按 `log_freq` 打 train loss / lr
+8. 按 `save_freq` 保存 `checkpoints/vla_model_step_*.pt`
+9. 每个 epoch 结束时，按配置跑：
+   - held-out action MSE
+   - rollout validation
+10. 最后写：
+    - `checkpoints/vla_model_best.pt`
+    - `checkpoints/vla_model_final.pt`
+
+当前 best model 选择逻辑：
+- **第一次拿到 rollout reward 之前**：先用 `val_loss`（或 train loss 回退）挑 best
+- **第一次 rollout 之后**：优先用 `rollout_avg_reward` 挑 best
+
+输出目录一般通过 `hydra.run.dir=...` 固定；否则 Hydra 自己生成。
+
+---
+
+## 4. 验证流怎么走
+
+### 4.1 held-out 数值验证
+
+当前常用做法不是随机切 `val_split`，而是：
+- `train.val_split=0.0`
+- `train.val_episode_indices=[100]`
+- `train.action_mse_val_freq_epochs=1`
+
+这样每个 epoch 结束都会在 `episode_100.hdf5` 上跑一次 `compute_action_mse_validation()`，日志 key 是：
+- 控制台 / `train_vla.log`：`held-out action MSE`
+- SwanLab：`val/action_mse`
+
+### 4.2 rollout 验证
+
+当前训练内 rollout 验证由 `train_vla.py -> run_rollout_validation() -> eval_vla._run_eval()` 触发。
+
+当前这条分支的常用训练内 rollout 约束是：
+- `train.rollout_val_freq_epochs=5`
+- `train.rollout_num_episodes=10`
+- `train.rollout_validate_on_checkpoint=false`
+- 强制 headless
+- 强制 `verbose_action=false`
+- 强制 `record_video=false`
+- 强制 `save_trajectory_image=true`
+- 强制 `trajectory_image_camera_name=front`
+- 强制 `save_summary_json=true`
+
+当前已经修正为**配置驱动的 rollout device / worker 路径**：
+- `train.rollout_device`：默认跟随 `train.device`
+- `train.rollout_num_workers`：默认 `null`
+  - 当 rollout 设备是 CPU 时，自动退化为 `1`
+  - 当 rollout 设备是 CUDA 时，自动推断为 `min(train.rollout_num_episodes, 8)`
+- `train.rollout_cuda_devices`：默认 `null`，等价于当前可见逻辑 GPU `[0]`
+- `train.rollout_response_timeout_s`
+- `train.rollout_server_startup_timeout_s`
+
+所以现在：
+- 训练在 `cuda` 上时，**训练期 rollout 默认会走 GPU**
+- 如果 `rollout_num_workers > 1`，就会自动走并行 rollout
+- 可以是 **单 GPU 多 worker 共用一个 inference server**
+- 也可以是 **多 GPU 多 server 分摊 worker**
+
+训练内 rollout artifact 默认落到：
+`<hydra.run.dir>/rollout_artifacts/<checkpoint_stem>/`
+
+常见文件：
+- `rollout_summary.json`
+- `rollout_front_ep01_trajectory.png` ... `rollout_front_ep10_trajectory.png`
+
+日志重点看：
+- `Epoch X rollout 平均奖励`
+- `最佳模型已更新`
+
+---
+
+## 5. 数据集加载与 `val_episode_indices` 机制
+
+### 5.1 数据集格式
+
+`SimpleRobotDataset` 读取 `dataset_dir` 下的 `episode_*.hdf5`，每个 episode 文件里至少要有：
+- `action`
+- `observations/qpos`
+- `observations/images/{cam_name}`
+
+当前常用相机：
+- `r_vis`
+- `top`
+- `front`
+
+### 5.2 懒加载行为
+
+`roboimi/vla/data/simpe_robot_dataset.py` 是按帧懒加载，不会一次性把整套 HDF5 全读进内存。
+
+它会：
+- 扫描目录下的 HDF5 文件
+- 用文件名里的 episode 编号（如 `episode_100.hdf5` -> `100`）建立 `available_episode_indices`
+- 在 worker 内做 HDF5 文件句柄 LRU 缓存
+
+### 5.3 `val_episode_indices` 怎么切
+
+`build_train_val_datasets()` 的逻辑是：
+
+1. 先 instantiate 一次完整 dataset
+2. 读取 `dataset.available_episode_indices`
+3. 检查 `train.val_episode_indices` 是否都存在
+4. 用 `episode_indices=` 再各 instantiate 一次：
+   - train dataset = 全部 episode - held-out episode
+   - val dataset = 只包含 held-out episode
+
+因此：
+- `train.val_episode_indices=[100]` 的意思是“把 `episode_100.hdf5` 整个拿去做 held-out val”
+- 如果 episode 不存在，会直接报错
+- 如果你把所有 episode 都塞进 `val_episode_indices`，也会直接报错，因为训练集会变空
+
+### 5.4 图像 resize 与 LeWM 附加字段
+
+dataset 侧 resize 默认来自：
+- `data.image_resize_shape`
+- 如果 backbone 额外覆盖，则优先 `agent.vision_backbone.dataset_image_resize_shape`
+
+返回 batch 除了常规：
+- `observation.state`
+- `observation.<cam>`
+- `action`
+
+还会在 LeWM 打开时返回：
+- `lewm.observation.state`
+- `lewm.observation.<cam>`
+- `lewm.future.state`
+- `lewm.future.<cam>`
+
+### 5.5 统计文件
+
+训练和推理都默认依赖 `dataset_stats.pkl`。数据集更新后重算：
+
+```bash
+/home/droid/.conda/envs/roboimi/bin/python roboimi/vla/scripts/calculate_stats.py \
+  --dataset_dir /home/droid/project/diana_sim/sim_transfer
+```
+
+远端只要把 `--dataset_dir` 换成对应主机路径即可。
+
+---
+
+## 6. SwanLab 行为
+
+当前配置默认值里 `train.use_swanlab=false`，但本分支常用 recipe 基本都显式开：
+- `train.use_swanlab=true`
+- `train.swanlab_project=roboimi-vla`
+- `train.swanlab_run_name=<run_name>`
+
+`train_vla.py` 的 SwanLab 行为：
+- 初始化时上传 `train` / `data` / `agent` 三段 config
+- 训练中记录：
+  - `train/loss`
+  - `train/lr`
+  - `train/best_loss`
+  - `train/step`
+- checkpoint 验证时记录：
+  - `val/loss`
+- held-out 数值验证时记录：
+  - `val/action_mse`
+- rollout 验证时记录：
+  - `rollout/avg_reward`
+  - `rollout/epoch`
+- 训练结束时记录：
+  - `final/checkpoint_path`
+  - `final/best_checkpoint_path`
+
+训练期 rollout 生成的前视图轨迹 PNG 会 best-effort 上传到 SwanLab；失败只会 warning，不会让训练中断。
+
+---
+
+## 7. 并行 rollout 说明
+
+### 7.1 这套能力从哪里来
+
+本分支的并行 rollout 方向不是 DataLoader 并行，而是 **`eval_vla.py` 的 multiprocess rollout path**。
+参考来源：
+`/home/droid/project/roboimi/.worktrees/multiprocess-rollout/roboimi/demos/vla_scripts/eval_vla.py`
+
+那条路径的控制参数是：
+- `eval.num_workers`
+- `eval.cuda_devices`
+
+语义是：
+- `eval.num_workers`：环境 worker 数，按 episode 切分
+- `eval.cuda_devices`：推理 server 绑定到哪些逻辑 GPU
+
+### 7.2 两种常见模式
+
+1. **单机单卡，多 worker 共用同一张 GPU**
+   - 典型：本地 5090 只有 1 卡，但想让 4 个 rollout worker 并行跑环境
+   - 形式：`eval.device=cuda eval.num_workers=4 'eval.cuda_devices=[0]'`
+   - 这时是 **1 个 CUDA inference server + 4 个 env worker**
+
+2. **单机多卡，多 server 分摊 worker**
+   - 典型：5880 有 2 卡，L20 有多卡
+   - 形式：`eval.device=cuda eval.num_workers=8 'eval.cuda_devices=[0,1]'`
+   - worker 会按 round-robin 分到多个 server 上
+
+### 7.3 操作上要注意什么
+
+- 并行 rollout 依赖 **多进程 eval 路径**，不是 `train.num_workers`
+- `train.num_workers` 是 DataLoader worker，和 rollout 并行不是一回事
+- `eval.num_workers > 1` 时必须 `eval.headless=true`
+- worker 数会自动 cap 到 `eval.num_episodes`
+- multiprocess rollout 当前已经支持 **per-episode trajectory image PNG**；多 worker 时每个 worker 会在自己的 artifact 子目录下写图，summary 会带回对应路径
+- 但多 worker 时仍然不要同时要求：
+  - `eval.record_video=true`
+  - `eval.save_trajectory=true`
+  - `eval.save_trajectory_npz=true`
+- `eval.save_trajectory_image=true` 现在是可以开的，适合并行 reward + 定性检查一起做
+
+### 7.4 并行 rollout 命令模板
+
+**5090 单卡 4 worker：**
+
+```bash
+/home/droid/.conda/envs/roboimi/bin/python roboimi/demos/vla_scripts/eval_vla.py \
+  agent=lewm_resnet_query_imf_attnres \
+  data.dataset_dir=/home/droid/project/diana_sim/sim_transfer \
+  train.device=cuda eval.device=cuda eval.headless=true eval.verbose_action=false \
+  eval.ckpt_path=/home/droid/project/roboimi/.worktrees/feat-lewm-imf-fusion/runs/<run_name>/checkpoints/vla_model_best.pt \
+  eval.num_episodes=10 eval.num_workers=4 'eval.cuda_devices=[0]' \
+  eval.save_summary_json=true eval.artifact_dir=/tmp/lewm_parallel_eval_5090
+```
+
+**5880 双卡 8 worker：**
+
+```bash
+/home/droid/miniforge3/envs/roboimi/bin/python roboimi/demos/vla_scripts/eval_vla.py \
+  agent=lewm_resnet_query_imf_attnres \
+  data.dataset_dir=/home/droid/sim_dataset/sim_transfer \
+  train.device=cuda eval.device=cuda eval.headless=true eval.verbose_action=false \
+  eval.ckpt_path=/home/droid/roboimi_suite_20260416_lewm_imf_fusion/runs/<run_name>/checkpoints/vla_model_best.pt \
+  eval.num_episodes=10 eval.num_workers=8 'eval.cuda_devices=[0,1]' \
+  eval.save_summary_json=true eval.artifact_dir=/tmp/lewm_parallel_eval_5880
+```
+
+---
+
+## 8. 当前常用命令 / 脚本
+
+### 8.1 本地 5090：直接用 suite 脚本
+
+现成脚本：
+`experiment_suites/2026-04-21-lewm-fromscratch-old9-epoch50-roll5-val-20260421-153037/launch_local_5090.sh`
+
+运行：
+
+```bash
+bash experiment_suites/2026-04-21-lewm-fromscratch-old9-epoch50-roll5-val-20260421-153037/launch_local_5090.sh
+```
+
+### 8.2 本地 5090：手动启动同 recipe
+
+```bash
+/home/droid/.conda/envs/roboimi/bin/python roboimi/demos/vla_scripts/train_vla.py \
+  agent=lewm_resnet_query_imf_attnres \
+  data.dataset_dir=/home/droid/project/diana_sim/sim_transfer \
+  'agent.lewm_query_offsets=[8]' \
+  agent.pred_horizon=8 \
+  agent.num_action_steps=8 \
+  train.device=cuda \
+  train.batch_size=32 \
+  train.lr=0.0001 \
+  train.max_steps=109350 \
+  train.num_workers=4 \
+  train.save_freq=10000 \
+  train.rollout_validate_on_checkpoint=false \
+  train.rollout_val_freq_epochs=5 \
+  train.rollout_num_episodes=10 \
+  train.val_split=0.0 \
+  'train.val_episode_indices=[100]' \
+  train.action_mse_val_freq_epochs=1 \
+  train.use_swanlab=true \
+  train.swanlab_project=roboimi-vla \
+  train.swanlab_run_name=lewmimf-q08-ph08-ex08-emb384-l12-fromscratch-epoch50-step109350-5090g0-20260421-153037 \
+  train.pretrained_ckpt=null \
+  agent.lewm_pretrained_ckpt=null \
+  hydra.run.dir=/home/droid/project/roboimi/.worktrees/feat-lewm-imf-fusion/runs/lewmimf-q08-ph08-ex08-emb384-l12-fromscratch-epoch50-step109350-5090g0-20260421-153037
+```
+
+### 8.3 5880：常用命令模板
+
+```bash
+ssh droid@100.73.14.65
+cd /home/droid/roboimi_suite_20260416_lewm_imf_fusion
+/home/droid/miniforge3/envs/roboimi/bin/python roboimi/demos/vla_scripts/train_vla.py \
+  agent=lewm_resnet_query_imf_attnres \
+  data.dataset_dir=/home/droid/sim_dataset/sim_transfer \
+  'agent.lewm_query_offsets=[8]' \
+  agent.pred_horizon=16 \
+  agent.num_action_steps=8 \
+  train.device=cuda train.batch_size=32 train.lr=0.0001 train.max_steps=109350 \
+  train.num_workers=4 train.save_freq=10000 train.rollout_validate_on_checkpoint=false \
+  train.rollout_val_freq_epochs=5 train.rollout_num_episodes=10 train.val_split=0.0 \
+  'train.val_episode_indices=[100]' train.action_mse_val_freq_epochs=1 \
+  train.use_swanlab=true train.swanlab_project=roboimi-vla \
+  train.swanlab_run_name=lewmimf-q08-ph16-ex08-emb384-l12-fromscratch-epoch50-step109350-5880g0-20260421-153037 \
+  train.pretrained_ckpt=null agent.lewm_pretrained_ckpt=null \
+  hydra.run.dir=/home/droid/roboimi_suite_20260416_lewm_imf_fusion/runs/lewmimf-q08-ph16-ex08-emb384-l12-fromscratch-epoch50-step109350-5880g0-20260421-153037
+```
+
+### 8.4 L20：常用命令模板
+
+```bash
+ssh droid@100.119.99.14
+cd /data/roboimi_suite_20260416_lewm_imf_fusion
+/home/droid/miniforge3/envs/roboimi/bin/python roboimi/demos/vla_scripts/train_vla.py \
+  agent=lewm_resnet_query_imf_attnres \
+  data.dataset_dir=/data/simtransfer/current \
+  'agent.lewm_query_offsets=[16]' \
+  agent.pred_horizon=16 \
+  agent.num_action_steps=16 \
+  train.device=cuda train.batch_size=32 train.lr=0.0001 train.max_steps=109350 \
+  train.num_workers=4 train.save_freq=10000 train.rollout_validate_on_checkpoint=false \
+  train.rollout_val_freq_epochs=5 train.rollout_num_episodes=10 train.val_split=0.0 \
+  'train.val_episode_indices=[100]' train.action_mse_val_freq_epochs=1 \
+  train.use_swanlab=true train.swanlab_project=roboimi-vla \
+  train.swanlab_run_name=lewmimf-q16-ph16-ex16-emb384-l12-fromscratch-epoch50-step109350-l20g0-20260421-153037 \
+  train.pretrained_ckpt=null agent.lewm_pretrained_ckpt=null \
+  hydra.run.dir=/data/roboimi_suite_20260416_lewm_imf_fusion/runs/lewmimf-q16-ph16-ex16-emb384-l12-fromscratch-epoch50-step109350-l20g0-20260421-153037
+```
+
+### 8.5 单次离线验证（当前分支已支持并行）
+
+**单 GPU / 4 worker：**
+
+```bash
+/home/droid/.conda/envs/roboimi/bin/python roboimi/demos/vla_scripts/eval_vla.py \
+  agent=lewm_resnet_query_imf_attnres \
+  data.dataset_dir=/home/droid/project/diana_sim/sim_transfer \
+  train.device=cuda eval.device=cuda \
+  eval.ckpt_path=/home/droid/project/roboimi/.worktrees/feat-lewm-imf-fusion/runs/<run_name>/checkpoints/vla_model_best.pt \
+  eval.num_episodes=10 eval.num_workers=4 'eval.cuda_devices=[0]' \
+  eval.headless=true eval.verbose_action=false \
+  eval.save_summary_json=true eval.save_trajectory_image=true \
+  eval.trajectory_image_camera_name=front \
+  eval.artifact_dir=/tmp/lewm_eval_front
+```
+
+**训练内启用并行 GPU rollout（推荐显式写清楚）**：
+
+```bash
+/home/droid/.conda/envs/roboimi/bin/python roboimi/demos/vla_scripts/train_vla.py \
+  agent=lewm_resnet_query_imf_attnres \
+  data.dataset_dir=/home/droid/project/diana_sim/sim_transfer \
+  'agent.lewm_query_offsets=[8]' \
+  agent.pred_horizon=8 \
+  agent.num_action_steps=8 \
+  train.device=cuda \
+  train.batch_size=32 \
+  train.lr=0.0001 \
+  train.max_steps=109350 \
+  train.num_workers=4 \
+  train.save_freq=10000 \
+  train.rollout_val_freq_epochs=5 \
+  train.rollout_num_episodes=10 \
+  train.rollout_device=cuda \
+  train.rollout_num_workers=4 \
+  'train.rollout_cuda_devices=[0]' \
+  train.rollout_validate_on_checkpoint=false \
+  train.val_split=0.0 \
+  'train.val_episode_indices=[100]' \
+  train.action_mse_val_freq_epochs=1 \
+  train.use_swanlab=true \
+  train.swanlab_project=roboimi-vla \
+  train.swanlab_run_name=<run_name> \
+  hydra.run.dir=/home/droid/project/roboimi/.worktrees/feat-lewm-imf-fusion/runs/<run_name>
+```
+
+### 8.6 监控日志
+
+```bash
+tail -f runs/<run_name>/launch.stdout.log
+tail -f runs/<run_name>/train_vla.log
+```
+
+远端就把 `runs/<run_name>` 换成 manifest 里的绝对路径。
+
+---
+
+## 9. 操作建议
+
+- **优先以 suite 的 `manifest.json` / `notes.md` / `launch_logs/*.launch.log` 为准**，不要手写一套和历史 run 不一致的命令
+- 要做当前常用验证，就显式加上：
+  - `train.val_split=0.0`
+  - `train.val_episode_indices=[100]`
+  - `train.action_mse_val_freq_epochs=1`
+  - `train.rollout_val_freq_epochs=5`
+  - `train.rollout_num_episodes=10`
+- 本分支如果要对比不同 horizon / action-step，尽量只改：
+  - `agent.lewm_query_offsets`
+  - `agent.pred_horizon`
+  - `agent.num_action_steps`
+- 想复现 2026-04-21 那轮 from-scratch 结果时，记得同时设：
+  - `train.pretrained_ckpt=null`
+  - `agent.lewm_pretrained_ckpt=null`

docs/superpowers/plans/2026-04-23-sim-air-insert-ring-bar.md

@@ -1,311 +0,0 @@
-# sim_air_insert_ring_bar Implementation Plan
-
-> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking.
-
-**Goal:** Add an independent dual-Diana MuJoCo task `sim_air_insert_ring_bar` with a square ring block, a square bar block, staged rewards, strict finite-geometry in-air insertion success detection, and a task-specific scripted policy.
-
-**Architecture:** Reuse the current dual-Diana EE-control stack and environment factory, but add a task-specific scene XML, robot asset entrypoint, sampling helpers, and a new task-specific environment module. Keep `sim_transfer` untouched while introducing pure-Python geometry helpers and focused tests so reward/success behavior can be regression tested without requiring a full MuJoCo rollout in every test.
-
-**Tech Stack:** Python, unittest, MuJoCo XML assets, existing dual-Diana environment classes, Hydra-compatible task naming/config patterns.
-
----
-
-## File Structure / Responsibilities
-
-- **Create:** `roboimi/assets/models/manipulators/DianaMed/ring_bar_objects.xml`
-  - Defines the rigid ring body and bar body, each with a free joint and stable box-based geoms.
-- **Create:** `roboimi/assets/models/manipulators/DianaMed/bi_diana_ring_bar_ee.xml`
-  - Scene entrypoint that includes the shared world/table/robot assets plus the new object XML.
-- **Modify:** `roboimi/assets/robots/diana_med.py`
-  - Add a task-specific robot asset class for the new scene XML without changing existing `BiDianaMed` behavior.
-- **Modify:** `roboimi/utils/act_ex_utils.py`
-  - Add deterministic helpers to sample left/right planar placement regions for ring and bar objects.
-- **Modify:** `roboimi/utils/constants.py`
-  - Register the new task name and default metadata.
-- **Create:** `roboimi/envs/double_air_insert_env.py`
-  - New task-specific environment, finite-geometry success helpers, reset logic, reward logic, and task factory branch.
-- **Modify:** `roboimi/envs/double_pos_ctrl_env.py`
-  - Route `make_sim_env()` to the new task-specific environment while keeping current `sim_transfer` logic unchanged.
-- **Create:** `roboimi/demos/diana_air_insert_policy.py`
-  - Task-specific waypoint/open-loop scripted policy for grasp-lift-align-insert.
-- **Modify:** `roboimi/demos/vla_scripts/eval_vla.py`
-  - Reset the new task with the correct sampled task state instead of assuming a single transfer box pose.
-- **Create:** `tests/test_air_insert_env.py`
-  - Focused unit tests for sampling, reset helpers, reward progression, and strict success detection.
-- **Modify:** `tests/test_eval_vla_headless.py`
-  - Add coverage that headless evaluation dispatches the correct reset sampler for the new task.
-- **Modify:** `tests/test_robot_asset_paths.py`
-  - Verify the new robot asset class resolves its XML path correctly independent of cwd.
-
----
-
-### Task 1: Add failing tests for task registration, samplers, and asset wiring
-
-**Files:**
-- Create: `tests/test_air_insert_env.py`
-- Modify: `tests/test_eval_vla_headless.py`
-- Modify: `tests/test_robot_asset_paths.py`
-- Modify: `roboimi/utils/act_ex_utils.py` (later in implementation)
-- Modify: `roboimi/utils/constants.py` (later in implementation)
-- Modify: `roboimi/assets/robots/diana_med.py` (later in implementation)
-- Modify: `roboimi/envs/double_pos_ctrl_env.py` (later in implementation)
-- Create: `roboimi/envs/double_air_insert_env.py` (minimal stub in this task)
-
-- [ ] **Step 1: Write failing tests for task config and sampling helpers**
-
-Add tests in `tests/test_air_insert_env.py` covering:
-- `SIM_TASK_CONFIGS['sim_air_insert_ring_bar']` exists
-- `sample_air_insert_ring_bar_pose()` (or equivalent helper) returns ring/bar positions with fixed z and correct left/right planar ranges
-- output structure is explicit and easy for reset/eval code to consume
-
-- [ ] **Step 2: Write failing tests for environment factory dispatch and robot asset resolution**
-
-Add tests covering:
-- `make_sim_env('sim_air_insert_ring_bar', headless=True)` dispatches to the new environment with rendering disabled
-- a new robot asset class resolves the new XML path independent of cwd, similar to the existing `BiDianaMed` test pattern
-
-- [ ] **Step 3: Write failing tests for eval reset helper dispatch**
-
-Extend `tests/test_eval_vla_headless.py` so headless eval can reset the new task using the new sampler instead of hard-coding `sample_transfer_pose()`.
-
-- [ ] **Step 4: Run the targeted tests to verify they fail for the expected missing-feature reasons**
-
-Run:
-`/home/droid/.conda/envs/roboimi/bin/python -m unittest tests.test_air_insert_env tests.test_eval_vla_headless tests.test_robot_asset_paths -v`
-
-Expected:
-- FAIL because the new task config/helper/class/dispatch branch does not exist yet
-
-- [ ] **Step 5: Implement the minimal production code to satisfy the new task registration and helper tests**
-
-Implement only enough to make the new tests pass:
-- add new task config entry
-- add the new placement sampler
-- add the new robot asset class
-- create a minimal importable `double_air_insert_env.py` stub and class/function surface needed for factory dispatch tests
-- add the factory dispatch branch / headless wiring
-- update eval reset dispatch for the new task
-
-- [ ] **Step 6: Re-run the targeted tests to verify they pass**
-
-Run:
-`/home/droid/.conda/envs/roboimi/bin/python -m unittest tests.test_air_insert_env tests.test_eval_vla_headless tests.test_robot_asset_paths -v`
-
-Expected:
-- PASS for the new registration/sampler/dispatch/asset tests
-
-- [ ] **Step 7: Commit Task 1**
-
-Run:
-`git add tests/test_air_insert_env.py tests/test_eval_vla_headless.py tests/test_robot_asset_paths.py roboimi/utils/act_ex_utils.py roboimi/utils/constants.py roboimi/assets/robots/diana_med.py roboimi/envs/double_pos_ctrl_env.py roboimi/envs/double_air_insert_env.py roboimi/demos/vla_scripts/eval_vla.py && git commit -m "feat(env): register sim air insert ring bar task"`
-
----
-
-### Task 2: Add the MuJoCo ring+bar scene assets and reset helpers
-
-**Files:**
-- Create: `roboimi/assets/models/manipulators/DianaMed/ring_bar_objects.xml`
-- Create: `roboimi/assets/models/manipulators/DianaMed/bi_diana_ring_bar_ee.xml`
-- Create or Modify: `roboimi/envs/double_air_insert_env.py`
-- Modify: `tests/test_air_insert_env.py`
-
-- [ ] **Step 1: Write failing tests for object reset helpers and scene-specific joint naming assumptions**
-
-In `tests/test_air_insert_env.py`, add unit tests for helper functions that:
-- write ring pose to `ring_block_joint` from the named task-state mapping
-- write bar pose to `bar_block_joint` from the named task-state mapping
-- read back `env_state` as a stable 14D vector `[ring_pos, ring_quat, bar_pos, bar_quat]`
-
-Use fake `mj_data` objects so tests stay fast and deterministic.
-
-- [ ] **Step 2: Run the focused test slice and verify it fails**
-
-Run:
-`/home/droid/.conda/envs/roboimi/bin/python -m unittest tests.test_air_insert_env -v`
-
-Expected:
-- FAIL because reset/state helper functions and joint conventions are not implemented yet
-
-- [ ] **Step 3: Implement the scene XML files and reset/state helper code**
-
-Implement:
-- the object XML with one rigid ring body and one rigid bar body
-- the task scene XML entrypoint using the shared world/table/robot includes
-- reset helper(s) in `double_air_insert_env.py` that set qpos for both free joints with fixed quaternions
-- task-state accessor(s) returning both object poses in a stable structure
-
-- [ ] **Step 4: Re-run the focused test slice and verify it passes**
-
-Run:
-`/home/droid/.conda/envs/roboimi/bin/python -m unittest tests.test_air_insert_env -v`
-
-Expected:
-- PASS for reset/state helper tests
-
-- [ ] **Step 5: Commit Task 2**
-
-Run:
-`git add roboimi/assets/models/manipulators/DianaMed/ring_bar_objects.xml roboimi/assets/models/manipulators/DianaMed/bi_diana_ring_bar_ee.xml roboimi/envs/double_air_insert_env.py tests/test_air_insert_env.py && git commit -m "feat(scene): add ring and bar insertion scene assets"`
-
----
-
-### Task 3: Implement strict reward and finite-geometry success detection
-
-**Files:**
-- Modify: `roboimi/envs/double_air_insert_env.py`
-- Modify: `tests/test_air_insert_env.py`
-
-- [ ] **Step 1: Write failing tests for reward stages and strict success detection**
-
-Add tests in `tests/test_air_insert_env.py` for:
-- left contact stage reward
-- right contact stage reward
-- ring lifted off table stage
-- bar lifted off table stage
-- positive success case where a finite bar truly passes through the aperture
-- negative case where the centerline would pass but the finite square body would clip
-- negative case where the bar has not crossed the ring thickness direction enough
-- negative case where one/both objects are still on the table
-
-Structure the tests around pure helper functions and light fake contact/state objects so the geometry logic is directly regression tested.
-
-- [ ] **Step 2: Run the focused tests and verify they fail for missing reward/success logic**
-
-Run:
-`/home/droid/.conda/envs/roboimi/bin/python -m unittest tests.test_air_insert_env -v`
-
-Expected:
-- FAIL because the staged reward and finite-geometry insertion logic are not implemented yet
-
-- [ ] **Step 3: Implement minimal strict success helpers and reward logic**
-
-Implement in `roboimi/envs/double_air_insert_env.py`:
-- pure helper(s) for transforming bar geometry into ring-local coordinates
-- finite-geometry insertion predicate (not centerline-only)
-- table-contact / airborne checks
-- staged reward function returning the highest achieved stage with `max_reward = 5`
-
-- [ ] **Step 4: Re-run the focused tests to verify the logic passes**
-
-Run:
-`/home/droid/.conda/envs/roboimi/bin/python -m unittest tests.test_air_insert_env -v`
-
-Expected:
-- PASS for reward and success-detection regression tests
-
-- [ ] **Step 5: Commit Task 3**
-
-Run:
-`git add roboimi/envs/double_air_insert_env.py tests/test_air_insert_env.py && git commit -m "feat(env): add strict air insertion reward and success logic"`
-
----
-
-### Task 4: Add the scripted policy and integration smoke coverage
-
-**Files:**
-- Create: `roboimi/demos/diana_air_insert_policy.py`
-- Modify: `roboimi/demos/diana_record_sim_episodes.py`
-- Modify: `tests/test_air_insert_env.py`
-- Optionally Modify: `roboimi/demos/vla_scripts/eval_vla.py` (only if integration gaps remain after Task 1)
-
-- [ ] **Step 1: Write failing tests for scripted-policy action shape and basic generation**
-
-Add tests covering:
-- the new policy produces a 16D action
-- trajectory generation accepts sampled named task state without error
-- the first action is a valid open-gripper safe pose command
-- a deterministic nominal smoke path (with canonical sampled state or fake env shim) reaches the intended terminal interface contract without shape/reward mismatches
-
-Keep the tests unit-level; do not require a full MuJoCo rollout for every assertion.
-
-- [ ] **Step 2: Write failing tests for the scripted rollout entrypoint and a real headless smoke path**
-
-Add coverage for both:
-- the standard scripted rollout entrypoint (`roboimi/demos/diana_record_sim_episodes.py`) can select the new task sampler/policy instead of remaining sim_transfer-only
-- a deterministic integration/smoke test that instantiates `make_sim_env('sim_air_insert_ring_bar', headless=True)`, resets with sampled named task state, and steps a few actions or scripted-policy outputs using the real task XML and task-specific wiring
-
-- [ ] **Step 3: Run the scripted-policy tests and verify they fail**
-
-Run:
-`/home/droid/.conda/envs/roboimi/bin/python -m unittest tests.test_air_insert_env -v`
-
-Expected:
-- FAIL because the new scripted policy does not exist yet
-
-- [ ] **Step 4: Implement the waypoint-based scripted policy**
-
-Implement a conservative open-loop policy with phases:
-- safe wait pose
-- above-target approach
-- descend + grasp
-- dual lift
-- airborne meeting alignment
-- bar push-through insertion
-
-Use fixed orientations for version 1 and follow the existing repository style from `diana_policy.py`.
-
-- [ ] **Step 5: Re-run the scripted-policy tests to verify they pass**
-
-Run:
-`/home/droid/.conda/envs/roboimi/bin/python -m unittest tests.test_air_insert_env -v`
-
-Expected:
-- PASS for scripted-policy tests
-
-- [ ] **Step 6: Run the combined verification suite for this feature**
-
-Run:
-`/home/droid/.conda/envs/roboimi/bin/python -m unittest tests.test_air_insert_env tests.test_eval_vla_headless tests.test_eval_vla_rollout_artifacts tests.test_train_vla_rollout_validation tests.test_robot_asset_paths -v`
-
-Expected:
-- PASS with 0 failures
-
-- [ ] **Step 6b: Run the mandatory real headless smoke check**
-
-Run a focused smoke command that instantiates the real task, resets with sampled state, and steps a few actions using the new scripted policy or a deterministic action sequence.
-
-Example command (adjust module/test helper if needed):
-`/home/droid/.conda/envs/roboimi/bin/python -m unittest tests.test_air_insert_env.AirInsertEnvSmokeTest -v`
-
-Expected:
-- PASS, proving the real XML/assets/env wiring instantiate and step correctly in headless mode
-
-- [ ] **Step 7: Commit Task 4**
-
-Run:
-`git add roboimi/demos/diana_air_insert_policy.py tests/test_air_insert_env.py tests/test_eval_vla_headless.py tests/test_robot_asset_paths.py roboimi/demos/vla_scripts/eval_vla.py && git commit -m "feat(policy): add scripted air insertion policy"`
-
----
-
-### Task 5: Final verification and implementation review
-
-**Files:**
-- Review all files touched above
-
-- [ ] **Step 1: Run fresh end-to-end verification before claiming completion**
-
-Run:
-`/home/droid/.conda/envs/roboimi/bin/python -m unittest tests.test_air_insert_env tests.test_eval_vla_headless tests.test_robot_asset_paths -v`
-
-Expected:
-- PASS with 0 failures
-
-- [ ] **Step 2: Inspect git status and recent commits**
-
-Run:
-`git status --short && git log --oneline --decorate -n 8`
-
-Expected:
-- only intended feature files modified / committed
-
-- [ ] **Step 3: Request final code review for the completed feature**
-
-Use the requesting-code-review skill against the full diff from the feature branch starting point to current HEAD.
-
-- [ ] **Step 4: Address any review findings and re-run verification if code changes**
-
-If fixes are made, repeat the unittest command from Step 1.
-
-- [ ] **Step 5: Hand off using finishing-a-development-branch**
-
-After verification and review, use the finishing-a-development-branch skill to decide merge / PR / cleanup.

docs/superpowers/plans/2026-05-25-native-smolvla-model-migration.md

@@ -1,184 +0,0 @@
-# Native SmolVLA Model Migration Implementation Plan
-
-> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking.
-
-**Goal:** Migrate only the SmolVLA model core into RoboIMI and expose it as a native VLA agent without importing the full LeRobot package.
-
-**Architecture:** Add a focused `roboimi.vla.models.smolvla` model package and a `SmolVLANativeAgent` wrapper that speaks the existing RoboIMI train/eval interface. Keep normalization, queues, camera ordering, and task fallback in the agent; keep model math in the migrated model package.
-
-**Tech Stack:** PyTorch, Transformers, Hydra/OmegaConf, unittest/pytest-style tests, existing RoboIMI normalization and VLA scripts.
-
----
-
-## File Structure
-
-- Create `roboimi/vla/models/smolvla/__init__.py`: package exports.
-- Create `roboimi/vla/models/smolvla/configuration.py`: lightweight config dataclass.
-- Create `roboimi/vla/models/smolvla/modeling.py`: model utilities and `VLAFlowMatching`.
-- Create `roboimi/vla/models/smolvla/smolvlm_with_expert.py`: adapted VLM/expert core.
-- Create `roboimi/vla/agent_smolvla_native.py`: RoboIMI agent wrapper.
-- Create `roboimi/vla/conf/agent/smolvla_native.yaml`: Hydra config.
-- Create `tests/test_smolvla_native_agent.py`: TDD tests for wrapper behavior.
-- Create `tests/test_smolvla_native_modeling.py`: TDD tests for utilities/config where useful.
-
-## Task 1: Wrapper behavior tests and minimal agent skeleton
-
-**Files:**
-- Create: `tests/test_smolvla_native_agent.py`
-- Create: `roboimi/vla/agent_smolvla_native.py`
-
-- [ ] **Step 1: Write failing tests for task fallback, camera order, queue, and fake model calls**
-
-Implement tests using fake tokenizer and fake model. Test names:
-
-- `test_compute_loss_orders_cameras_tokenizes_task_and_calls_native_model`
-- `test_unknown_task_uses_configured_task_description`
-- `test_missing_camera_raises_clear_error`
-- `test_predict_action_chunk_denormalizes_fake_model_output`
-- `test_select_action_uses_action_queue_before_recomputing`
-
-- [ ] **Step 2: Run tests and verify import failure**
-
-Run: `python -m unittest tests.test_smolvla_native_agent -v`
-
-Expected: FAIL because `roboimi.vla.agent_smolvla_native` does not exist.
-
-- [ ] **Step 3: Implement minimal `SmolVLANativeAgent` skeleton**
-
-Implement constructor injection for `model` and `tokenizer`, plus `_resolve_task`, `_order_images`, `_tokenize_tasks`, `reset`, `_prepare_observation_batch`, `compute_loss`, `predict_action_chunk`, and `select_action`. Use fake model in tests; real model construction can raise a clear ImportError until Task 3.
-
-- [ ] **Step 4: Run tests and verify pass**
-
-Run: `python -m unittest tests.test_smolvla_native_agent -v`
-
-Expected: PASS.
-
-## Task 2: Config and modeling utility tests
-
-**Files:**
-- Create: `tests/test_smolvla_native_modeling.py`
-- Create: `roboimi/vla/models/smolvla/configuration.py`
-- Create: `roboimi/vla/models/smolvla/modeling.py`
-- Create: `roboimi/vla/models/smolvla/__init__.py`
-
-- [ ] **Step 1: Write failing tests**
-
-Cover:
-
-- `NativeSmolVLAConfig` validates `n_action_steps <= chunk_size`.
-- `pad_vector` pads and rejects truncation.
-- `resize_with_pad` preserves batch/channel shape and output size.
-- `make_att_2d_masks` matches prefix-LM mask semantics.
-
-- [ ] **Step 2: Run tests and verify failure**
-
-Run: `python -m unittest tests.test_smolvla_native_modeling -v`
-
-Expected: FAIL because package/modeling code is missing.
-
-- [ ] **Step 3: Implement config and utility functions**
-
-Port minimal functions from external SmolVLA while removing LeRobot imports.
-
-- [ ] **Step 4: Run tests and verify pass**
-
-Run: `python -m unittest tests.test_smolvla_native_modeling -v`
-
-Expected: PASS.
-
-## Task 3: Migrate model core
-
-**Files:**
-- Modify: `roboimi/vla/models/smolvla/modeling.py`
-- Create: `roboimi/vla/models/smolvla/smolvlm_with_expert.py`
-- Modify: `roboimi/vla/agent_smolvla_native.py`
-
-- [ ] **Step 1: Write fake-backed integration tests for lazy real model construction**
-
-Patch `VLAFlowMatching` and tokenizer loader so `SmolVLANativeAgent(model=None, tokenizer=None)` constructs a fake model without real downloads. Verify config fields are passed.
-
-- [ ] **Step 2: Run test and verify failure**
-
-Run: `python -m unittest tests.test_smolvla_native_agent -v`
-
-Expected: FAIL because real construction path is not implemented.
-
-- [ ] **Step 3: Port `SmolVLMWithExpertModel` and `VLAFlowMatching`**
-
-Copy only model-core logic from external files, replacing LeRobot dependencies with local config and local helpers. Preserve `forward`, `sample_actions`, `denoise_step`, image resize/pad, state/action padding, and language token inputs.
-
-- [ ] **Step 4: Wire agent lazy construction**
-
-If `model` is not supplied, create `NativeSmolVLAConfig`, tokenizer, and `VLAFlowMatching`. Keep `model` injection path for tests.
-
-- [ ] **Step 5: Run unit tests**
-
-Run:
-
-```bash
-python -m unittest tests.test_smolvla_native_agent tests.test_smolvla_native_modeling -v
-```
-
-Expected: PASS.
-
-## Task 4: Hydra config
-
-**Files:**
-- Create: `roboimi/vla/conf/agent/smolvla_native.yaml`
-- Modify: `tests/test_smolvla_native_agent.py`
-
-- [ ] **Step 1: Add failing Hydra compose test**
-
-Test composing `agent=smolvla_native` exposes `_target_`, `action_dim=16`, `obs_dim=16`, `camera_names=${data.camera_names}`, and model config fields.
-
-- [ ] **Step 2: Run and verify failure**
-
-Run: `python -m unittest tests.test_smolvla_native_agent -v`
-
-Expected: FAIL because yaml is missing.
-
-- [ ] **Step 3: Add yaml config**
-
-Add `smolvla_native.yaml` with `_target_: roboimi.vla.agent_smolvla_native.SmolVLANativeAgent` and sane defaults.
-
-- [ ] **Step 4: Run test and verify pass**
-
-Run: `python -m unittest tests.test_smolvla_native_agent -v`
-
-Expected: PASS.
-
-## Task 5: Regression checks
-
-**Files:**
-- No production changes expected unless tests reveal integration gaps.
-
-- [ ] **Step 1: Run focused existing tests**
-
-Run:
-
-```bash
-python -m unittest tests.test_smolvla_prefix_encoder tests.test_smolvla_imf_agent tests.test_eval_vla_execution -v
-```
-
-Expected: PASS.
-
-- [ ] **Step 2: Search for forbidden import**
-
-Run:
-
-```bash
-rg -n "import lerobot|from lerobot" roboimi/vla/models/smolvla roboimi/vla/agent_smolvla_native.py
-```
-
-Expected: no matches.
-
-- [ ] **Step 3: Commit**
-
-Run:
-
-```bash
-git add roboimi/vla/models/smolvla roboimi/vla/agent_smolvla_native.py roboimi/vla/conf/agent/smolvla_native.yaml tests/test_smolvla_native_agent.py tests/test_smolvla_native_modeling.py docs/superpowers/specs/2026-05-25-native-smolvla-model-migration-design.md docs/superpowers/plans/2026-05-25-native-smolvla-model-migration.md
-git commit -m "feat(vla): add native SmolVLA model agent"
-```
-
-Expected: commit succeeds.

docs/superpowers/specs/2026-04-23-sim-air-insert-ring-bar-design.md

@@ -1,316 +0,0 @@
-# sim_air_insert_ring_bar Design
-
-## Summary
-
-Add a new independent MuJoCo simulation task named `sim_air_insert_ring_bar` that keeps the existing dual-Diana tabletop setup but replaces the single transfer box with two randomized objects:
-
-- a square ring block grasped by the left arm
-- a square bar block grasped by the right arm
-
-The task is to pick both objects off the table and complete an in-air insertion where the bar truly passes through the ring aperture. The existing `sim_transfer` task must remain unchanged.
-
-## Goals
-
-- Reuse the current dual-Diana EE-control simulation stack
-- Keep the same table/base robot arrangement as the existing transfer task
-- Add an independent task entrypoint and scene definition
-- Randomize planar placement of both objects within left/right task-specific regions
-- Implement reward staging for contact, lift, and successful in-air insertion
-- Add a scripted policy that performs pick, lift, align, and in-air insertion
-- Preserve compatibility with existing environment creation, evaluation, and rollout patterns
-
-## Non-Goals
-
-- No random yaw in the first version
-- No visual servoing or closed-loop insertion controller
-- No general multi-task environment framework refactor
-- No guarantee that the VLA training stack is immediately tuned for this new task
-- No replacement or behavior change for `sim_transfer`
-
-## Task Name
-
-Use a new task name:
-
-- `sim_air_insert_ring_bar`
-
-This task should be exposed alongside `sim_transfer`, not as a replacement.
-
-## Scene Geometry
-
-### Shared Base Scene
-
-Keep the dual Diana robot, the table, and the existing camera layout conceptually unchanged.
-
-### Ring Block
-
-Represent the square ring as a rigid free body composed from simple MuJoCo box geoms rather than an external mesh.
-
-Dimensions:
-
-- outer side length: 68 mm
-- inner aperture side length: 32 mm
-- thickness: 18 mm
-- ring wall width: 18 mm
-
-The ring should behave as a single object body with a single free joint.
-
-### Bar Block
-
-Represent the bar as a rigid free body with a single box geom.
-
-Dimensions:
-
-- length: 90 mm
-- cross-section: 18 mm x 18 mm
-
-The bar should also be a single free-joint body.
-
-## Initial Placement / Reset
-
-The first version uses position-only randomization with fixed orientation. Reset sampling stays **caller-driven**, matching the existing `sim_transfer` usage pattern in rollout/eval code: a helper samples task state, then callers pass that state into `env.reset(...)`.
-
-Use an explicit sampled task-state structure with named fields:
-
-- `ring_pos`: 3D position
-- `ring_quat`: fixed 4D quaternion for version 1
-- `bar_pos`: 3D position
-- `bar_quat`: fixed 4D quaternion for version 1
-
-Behavior:
-
-- ring block: randomized only in a left-side planar sampling region
-- bar block: randomized only in a right-side planar sampling region
-- both objects start flat on the table
-- both objects use fixed orientation at reset
-- no random yaw, tilt, or flip in this version
-
-The sampling regions should be chosen conservatively so that:
-
-- the left arm can comfortably reach and grasp the ring
-- the right arm can comfortably reach and grasp the bar
-- scripted open-loop pick trajectories remain feasible
-
-## Control / Action Interface
-
-Reuse the current 16D EE-space action convention already used by the dual-Diana position-control environment:
-
-- left arm EE pose: 7D (`xyz + quat`)
-- right arm EE pose: 7D (`xyz + quat`)
-- left gripper command: 1D
-- right gripper command: 1D
-
-The new task should continue using EE targets transformed through the existing IK-based control path.
-
-## Environment Structure
-
-Implement this as a new task-specific environment path while reusing the existing dual-Diana simulation base where possible.
-
-Expected responsibilities:
-
-- scene instantiation for the ring+bar setup
-- task reset for randomized object placement
-- environment-state accessors for both objects
-- reward computation
-- in-air insertion success detection
-
-The environment factory must dispatch by task name and leave the `sim_transfer` branch unchanged.
-
-## Observation / Environment State
-
-The task should retain the current observation structure style used by the dual-Diana environment:
-
-- `qpos`
-- multi-camera images
-
-For task state access, the environment should expose a stable `env_state` vector with this exact order:
-
-- `ring_pos[0:3]`
-- `ring_quat[3:7]`
-- `bar_pos[7:10]`
-- `bar_quat[10:14]`
-
-This 14D state should be sufficient for scripted-policy debugging and future rollout analysis, while reset itself remains caller-driven via the named task-state helper structure above.
-
-## Reward Design
-
-Use staged rewards in the same spirit as the current task, returning the highest achieved stage rather than accumulating one-time sparse bonuses per event.
-
-Maximum reward:
-
-- `max_reward = 5`
-
-Reward stages:
-
-1. left gripper touches the ring block
-2. right gripper touches the bar block
-3. ring block is lifted off the table
-4. bar block is lifted off the table
-5. while both objects are off the table, the bar truly passes through the ring aperture
-
-Notes:
-
-- contact rewards are intended as grasp-progress stages
-- lift rewards require the object to be off the table, not merely touched
-- final success reward only applies when both objects are airborne
-
-## Success Detection
-
-Success must **not** be based on a centerline-only check.
-
-A centerline-only test is insufficient because:
-
-- the bar has thickness, so a centerline can pass through while the body cannot
-- a square bar with imperfect orientation can have its centerline inside the aperture while its corners still collide with the ring
-
-### Required Success Semantics
-
-A successful insertion requires all of the following:
-
-1. the ring is off the table
-2. the bar is off the table
-3. the bar has actually crossed through the ring thickness direction
-4. the bar’s finite square cross-section fits through the square aperture during that crossing
-
-### Recommended Detection Approach
-
-Use a task-level geometric check in Python rather than relying on contact alone.
-
-Implementation intent:
-
-- transform the bar geometry into the ring’s local frame
-- reason about the bar as a finite oriented box (not a line)
-- verify that the bar has crossed the ring thickness direction
-- verify that the portion of the bar passing the aperture fits within the inner square opening, accounting for the bar’s cross-section and orientation
-
-This geometric check is the primary success test.
-
-### Role of Contacts
-
-Contacts may still be used for:
-
-- grasp-stage rewards
-- debugging / diagnostics
-
-But contact alone should **not** be the sole criterion for insertion success, since:
-
-- a true clean insertion may have limited aperture-wall contact
-- persistent contact can also happen while the bar is jammed and not actually inserted
-
-## Scripted Policy
-
-Add a new task-specific scripted policy for `sim_air_insert_ring_bar`.
-
-### Policy Intent
-
-The first version prioritizes a conservative, reliable open-loop demonstration rather than an optimized trajectory.
-
-### Action Phases
-
-The scripted policy should follow these phases:
-
-1. move both arms to safe initial / waiting poses with grippers open
-2. move left arm above the ring and right arm above the bar
-3. descend and grasp the assigned objects
-4. lift both objects clear of the table
-5. move both objects to an airborne meeting region above the table
-6. hold the ring stably while aligning the bar with the aperture
-7. push the bar along the intended insertion direction until the geometric success condition is met
-
-### Grasp Assignment
-
-- left arm: ring only
-- right arm: bar only
-
-### Motion Style
-
-Keep the current repository style:
-
-- waypoint-based trajectory definition
-- open-loop interpolation between waypoints
-- fixed grasp orientation in the first version
-
-No adaptive replanning is required for the first version.
-
-## Files / Integration Scope
-
-The implementation is expected to add task-specific files rather than broadly refactoring the codebase.
-
-Likely additions / changes:
-
-- a new MuJoCo scene XML for the ring+bar task
-- one or more XML fragments defining the two new objects
-- a new task-specific dual-Diana environment file
-- robot asset wiring for the new scene XML
-- reset sampling helpers for the new task
-- task registration in constants / environment factory paths
-- a new scripted policy file
-- focused tests for task creation, reset, rewards, success detection, and scripted policy shape/smoke behavior
-
-## Testing Requirements
-
-At minimum, add regression coverage for:
-
-### Environment Creation
-
-- the new task can be created via the task factory
-- the existing `sim_transfer` task remains unchanged
-
-### Reset / Sampling
-
-- ring reset positions are inside the left sampling region
-- bar reset positions are inside the right sampling region
-- reset orientation is fixed as intended
-
-### Environment State
-
-- environment-state access returns both object poses in the expected structure
-
-### Success Detection
-
-Must include both positive and negative cases.
-
-Positive case:
-
-- a configuration where the finite bar truly passes through the ring aperture is detected as success
-
-Negative cases:
-
-- centerline-inside but finite body would clip the aperture
-- not enough depth / not actually crossing the ring thickness direction
-- one or both objects still on the table
-
-### Reward Logic
-
-- left contact stage
-- right contact stage
-- ring lift stage
-- bar lift stage
-- final success stage with `max_reward = 5`
-
-### Scripted Policy
-
-At minimum:
-
-- policy emits valid 16D actions
-- trajectory generation does not error
-- rollout smoke path can step through the new environment
-
-## Risks / Constraints
-
-- MuJoCo contact naming must remain stable enough for stage rewards
-- geometric insertion checks must be strict enough to avoid false positives but not so brittle that numerically valid insertions are missed
-- scripted open-loop insertion may require conservative alignment and lift heights to keep the first version reliable
-
-## Acceptance Criteria
-
-The feature is complete when all of the following are true:
-
-- `sim_air_insert_ring_bar` is creatable as an independent task
-- the scene contains the dual Diana, table, ring block, and bar block
-- reset randomizes ring and bar positions in left/right planar regions with fixed orientation
-- the environment exposes task state for both objects
-- staged rewards progress to `max_reward = 5`
-- final success is based on finite-geometry insertion semantics, not a centerline-only shortcut
-- a new scripted policy can execute the intended pick-lift-align-insert behavior in the new environment
-- a canonical nominal smoke path (unit-level or deterministic integration-level) exists for the new scripted-policy interface so success is not judged purely by interpretation
-- existing `sim_transfer` behavior is preserved

docs/superpowers/specs/2026-05-25-native-smolvla-model-migration-design.md

@@ -1,86 +0,0 @@
-# Native SmolVLA Model Migration Design
-
-## Goal
-
-Migrate only the SmolVLA model core from `/data/lerobot-imf-attnres-exp/lerobot-imf-attnres` into this `roboimi` project so Diana simulation can train/evaluate through the existing `roboimi.vla` agent interface without importing or installing the full LeRobot tree.
-
-## Non-goals
-
-- Do not vendor the full `lerobot` package.
-- Do not change the current Python environment to LeRobot 0.5.x requirements.
-- Do not alter Diana environment action semantics.
-- Do not change `train_vla.py` or `eval_vla.py` main control flow unless a minimal compatibility hook is unavoidable.
-- Do not implement RTC in the first pass.
-
-## Architecture
-
-Add a native RoboIMI SmolVLA package under `roboimi/vla/models/smolvla/`. It will contain a lightweight config dataclass, a minimally adapted `SmolVLMWithExpertModel`, and a `VLAFlowMatching` implementation. A new `SmolVLANativeAgent` wraps the model with the existing RoboIMI agent contract: `compute_loss`, `predict_action_chunk`, `select_action`, `reset`, and `get_normalization_stats`.
-
-The new code will preserve the original model math where practical, but replace LeRobot framework dependencies with local constants, tokenizer handling, queue management, and `roboimi.vla.models.normalization.NormalizationModule`.
-
-## Data flow
-
-Training batch input remains the current RoboIMI format:
-
-```python
-{
-  "images": {cam: Tensor[B, T, C, H, W]},
-  "qpos": Tensor[B, T, 16],
-  "action": Tensor[B, H, 16],
-  "action_is_pad": optional BoolTensor[B, H],
-  "task": optional list[str],
-}
-```
-
-`SmolVLANativeAgent` normalizes `qpos` and `action` using RoboIMI dataset stats, tokenizes task strings, and passes images/state/language/action into the native SmolVLA model. In inference, the model returns normalized action chunks; the agent denormalizes them to 16-dim Diana EE actions.
-
-## Components
-
-### `roboimi/vla/models/smolvla/configuration.py`
-
-Defines `NativeSmolVLAConfig`, a small dataclass with fields needed by the model: state/action dimensions, padding dimensions, image resize target, tokenizer settings, VLM model name, VLM loading flags, expert layer configuration, sampling steps, dtype/device behavior, and compile flags.
-
-### `roboimi/vla/models/smolvla/smolvlm_with_expert.py`
-
-Migrates the model-core helper from the old repo. It should depend only on PyTorch and Transformers. It will expose `SmolVLMWithExpertModel` and attention helpers.
-
-### `roboimi/vla/models/smolvla/modeling.py`
-
-Defines model-core utilities (`resize_with_pad`, `pad_vector`, `make_att_2d_masks`, sinusoidal time embedding) plus `VLAFlowMatching`, with `forward` and `sample_actions`.
-
-### `roboimi/vla/agent_smolvla_native.py`
-
-RoboIMI-native agent wrapper. It owns tokenizer, normalization, task fallback, camera ordering, observation/action queues, loss masking, and denormalized rollout actions.
-
-### `roboimi/vla/conf/agent/smolvla_native.yaml`
-
-Hydra config for the native model, defaulting to Diana 16-dim state/action and configurable camera names.
-
-## Error handling
-
-- Missing configured camera raises `ValueError` with expected/missing names.
-- Task list length not matching batch size raises `ValueError`.
-- State/action dimensions exceeding configured `max_state_dim` / `max_action_dim` raises `ValueError`.
-- Missing Transformers SmolVLM classes raises `ImportError` explaining the required package.
-- Invalid action chunk shape raises `RuntimeError` explaining expected `(B,H,A)`.
-
-## Testing
-
-Use TDD with fake VLM/tokenizer/model components first, so tests do not download or instantiate the real SmolVLM. Cover:
-
-1. Config and Hydra instantiation with fake injected components.
-2. Camera ordering and missing-camera errors.
-3. Task fallback and newline/tokenization behavior.
-4. Loss path shape/mask behavior using a fake native model.
-5. `predict_action_chunk` normalization/denormalization and shape.
-6. `select_action` queue behavior.
-
-A later smoke test may instantiate the real model in an environment that already has the required Transformers/weights, but the first implementation must pass without external downloads.
-
-## Acceptance criteria
-
-- `agent=smolvla_native` can be composed by Hydra.
-- Unit tests pass without importing external `/data/.../src/lerobot`.
-- The new production code has no `import lerobot`.
-- The agent accepts the same batch/observation structure used by current train/eval scripts.
-- The agent emits 16-dim denormalized Diana EE actions for rollout.

render.py

@@ -1,200 +0,0 @@
-import mujoco
-from mujoco import viewer
-import sys
-import numpy as np
-import time
-import threading
-
-
-class MjBasicRenderer:
-    def __new__(cls, *args, **kwargs):
-        return super().__new__(cls)
-
-    def __init__(self, mj_model=None, mj_data=None):
-        # keyboard flag
-        self.render_paused = True
-        self.exit_flag = False
-        # init param
-        self.mj_model = mj_model
-        self.mj_data = mj_data
-        self.renderer = "viewer"  # default
-        self.viewer = None
-        self._image = None
-
-        # Set up mujoco viewer
-        self.image_renderer = mujoco.Renderer(self.mj_model)
-
-    def __del__(self):
-        pass
-
-    def _init_renderer(self):
-        """Initialize renderer, choose official renderer with "viewer"(joined from version 2.3.3),
-        another renderer with "mujoco_viewer"
-        """
-
-        def key_callback(keycode):
-            if keycode == 32:  # space
-                self.render_paused = not self.render_paused
-            elif keycode == 256:  # escape
-                self.exit_flag = not self.exit_flag
-
-        if self.renderer == "viewer":
-            # This function does not block, allowing user code to continue execution.
-            self.viewer = viewer.launch_passive(
-                self.mj_model,
-                self.mj_data,
-                key_callback=key_callback,
-                show_left_ui=False,
-                show_right_ui=False,
-            )
-            self.set_renderer_config()
-        else:
-            raise ValueError("Invalid renderer for some reason.")
-
-    def render(self):
-        """mujoco render"""
-        if self.viewer is not None and self.render_paused is True:
-            if self.viewer.is_running() and self.exit_flag is False:
-                self.viewer: viewer.Handle
-                self.viewer.sync()
-            else:
-                self.viewer.close()
-
-    def set_renderer_config(self):
-        """Setup mujoco global config while using viewer as renderer.
-        It should be noted that the render thread need locked.
-        """
-        self.viewer.cam.lookat = np.array([0.4, 0, 0.5])
-        self.viewer.cam.azimuth -= 0.005
-        with self.viewer.lock():
-            self.viewer.opt.flags[mujoco.mjtVisFlag.mjVIS_CONTACTPOINT] = int(
-                self.mj_data.time % 2
-            )
-
-
-try:
-    import cv2
-except ImportError:
-    print("Could not import cv2, please install it to enable camera viewer.")
-
-
-class MjMultiRenderer(MjBasicRenderer):
-
-    # __slots__=('mj_model','mj_data','renderer','enable_camera_viewer')
-    def __new__(cls, *args, **kwargs):
-        return super().__new__(cls)
-
-    def __init__(
-        self,
-        mj_model=None,
-        mj_data=None,
-        renderer=None,
-        enable_camera_viewer=False,
-        enable_depth=False,
-    ):
-        super().__init__(mj_model, mj_data)
-        self._depth = None
-        self.renderer = renderer
-        self._init_renderer()
-        self.enable_camera_viewer = enable_camera_viewer
-        if self.enable_camera_viewer:
-            self.enable_depth = enable_depth
-            self._init_window()
-        else:
-            self.enable_depth = False
-            print("No Camera View")
-
-    def __del__(self):
-        self.close()
-
-    def _init_renderer(self):
-        """
-        Initialize renderer, choose official renderer with "viewer"(joined from version 2.3.3)
-        """
-        if self.renderer == "unity":
-            # TODO: Support unity renderer.
-            raise ValueError("Unity renderer init failed for no supporting reason")
-        elif self.renderer == "viewer":
-            super()._init_renderer()
-            print("mujoco viewer init !")
-        else:
-            raise ValueError("renderer init failed for some reason.")
-
-    def _init_window(self, name="Camera view"):
-        if not self.enable_depth:
-            cv2.namedWindow(name, cv2.WINDOW_NORMAL)
-        else:
-            cv2.namedWindow(name, cv2.WINDOW_NORMAL)
-            cv2.namedWindow("Camera depth view", cv2.WINDOW_NORMAL)
-
-    def render(self):
-        """render mujoco"""
-        if self.renderer == "viewer":
-            super().render()
-        elif self.renderer == "unity":
-            # TODO: Support unity renderer.
-            raise ValueError("Unity renderer not supported now.")
-        else:
-            raise ValueError("Invalid renderer for some reason.")
-
-    def render(self):
-        """mujoco render"""
-        if self.viewer is not None and self.render_paused is True:
-            if self.viewer.is_running() and self.exit_flag is False:
-                self.viewer: viewer.Handle
-                self.viewer.sync()
-            else:
-                self.viewer.close()
-
-    def camera_render(self, cam=None):
-        if self.enable_camera_viewer:
-            if not self.enable_depth:
-                rgb, depth = self.render_from_camera(cam)
-                rgb = cv2.resize(rgb, (1920, 1600))
-                cv2.imshow("Camera view", rgb)
-                cv2.waitKey(1)
-
-            else:
-                rgb, depth = self.render_from_camera(cam)
-                cv2.imshow("Camera view", rgb)
-                cv2.imshow("Camera depth view", depth)
-                cv2.waitKey(1)
-
-        else:
-            print("camera info disable")
-            return
-
-    def render_from_camera(self, cam=None):
-        self.image_renderer.update_scene(self.mj_data, camera=cam)
-        if self.enable_depth is True:
-            self.image_renderer.enable_depth_rendering()
-            org = self.image_renderer.render()
-            depth = org[:, :]
-            self.image_renderer.disable_depth_rendering()
-            org = self.image_renderer.render()
-            image = org[:, :, ::-1]
-        else:
-            org = self.image_renderer.render()
-            image = org[:, :, ::-1]
-            depth = np.zeros([240, 320])
-        return image, depth
-
-    def close(self):
-        """close the environment."""
-        if self.enable_camera_viewer and self.viewer.is_running() == False:
-            cv2.destroyAllWindows()
-        self.viewer.close()
-        # sys.exit(0)
-
-    # def get_cam_intrinsic(self, fovy=45.0, width=320, height=240):
-    #     aspect = width * 1.0 / height
-    #     fovx = np.degrees(2 * np.arctan(aspect * np.tan(np.radians(fovy / 2))))
-
-    #     cx = 0.5 * width
-    #     cy = 0.5 * height
-    #     fx = cx / np.tan(fovx * np.pi / 180 * 0.5)
-    #     fy = cy / np.tan(fovy * np.pi / 180 * 0.5)
-
-    #     K = np.array([[fx, 0, cx],
-    #                   [0, fy, cy],
-    #                   [0, 0, 1]], dtype=np.float32)

roboimi/assets/models/manipulators/DianaMed/BiDianaMed_rethink.xml

@@ -76,7 +76,7 @@
                                                 <body name="ee_cam_left" pos="0.00 0.046 -0.075" euler="0.0 0.0 -0.0">
                                                         <inertial pos="0 0 0" quat="1 0 0 0" mass="0" diaginertia="0 0 0" />
                                                         <geom type="mesh" contype="1" conaffinity="1" group="1" rgba="0.69804 0.69804 0.69804 1" mesh="realsense_cam" />
-                                                        <camera name="rs_cam_left" mode="fixed" pos="0.0 0.0 0.01" euler="0.0 9.4 0.0 " fovy="50" resolution="1920 1200"/>
+                                                        <camera name="rs_cam_left" mode="fixed" pos="0.0 0.0 -0.25" euler="0.0 9.4 0.0 " fovy="15" resolution="1920 1200"/>
                                                 </body>
                                             </body>
                                             <body name="l_finger_left" pos="0 0.01 0.0444">

roboimi/assets/models/manipulators/DianaMed/bi_diana_socket_peg_ee.xml

@@ -1,6 +0,0 @@
-<mujoco model="bi_diana_socket_peg">
-    <include file="./empty_world.xml" />
-    <include file="./table_square.xml" />
-    <include file="./socket_peg_objects.xml" />
-    <include file="./BiDianaMed_rethink.xml" />
-</mujoco>

roboimi/assets/models/manipulators/DianaMed/socket_peg_objects.xml

@@ -1,19 +0,0 @@
-<mujoco model="socket_peg_objects">
-  <worldbody>
-     <body name="peg" pos="0.12 0.90 0.46">
-          <joint name="red_peg_joint" type="free" frictionloss="0.01" />
-          <inertial pos="0 0 0" mass="0.05" diaginertia="0.002 0.002 0.002" />
-          <geom condim="4" solimp="2 1 0.01" solref="0.01 1" friction="1 0.005 0.0001" pos="0 0 0" size="0.06 0.01 0.01" type="box" name="red_peg" rgba="1 0 0 1" />
-     </body>
-
-     <body name="socket" pos="-0.12 0.90 0.472">
-          <joint name="blue_socket_joint" type="free" frictionloss="0.01" />
-          <inertial pos="0 0 0" mass="0.05" diaginertia="0.002 0.002 0.002" />
-          <geom condim="4" solimp="2 1 0.01" solref="0.01 1" friction="1 0.05 0.001" pos="0 0 -0.02" size="0.06 0.018 0.002" type="box" name="socket-1" rgba="0 0 1 1" />
-          <geom condim="4" solimp="2 1 0.01" solref="0.01 1" friction="1 0.05 0.001" pos="0 0 0.02" size="0.06 0.018 0.002" type="box" name="socket-2" rgba="0 0 1 1" />
-          <geom condim="4" solimp="2 1 0.01" solref="0.01 1" friction="1 0.05 0.001" pos="0 0.02 0" size="0.06 0.002 0.018" type="box" name="socket-3" rgba="0 0 1 1" />
-          <geom condim="4" solimp="2 1 0.01" solref="0.01 1" friction="1 0.05 0.001" pos="0 -0.02 0" size="0.06 0.002 0.018" type="box" name="socket-4" rgba="0 0 1 1" />
-          <geom condim="4" solimp="2 1 0.01" solref="0.01 1" friction="1 0.005 0.0001" pos="0 0 0" size="0.04 0.01 0.01" type="box" name="pin" rgba="1 0 0 1" />
-     </body>
-  </worldbody>
-</mujoco>

roboimi/assets/models/manipulators/DianaMed/table_square.xml

@@ -7,6 +7,7 @@
           <geom name="table" condim="4" contype="1" conaffinity="1"  type="box" rgba="0.4 0.4 0.4 1" size="0.62 0.62 0.01" density="1500" friction="0.9 0.9 0.9"/>
       </body>
       <camera name="top" pos="0.0 1.0 2.0" fovy="44" mode="targetbody" target="table"/>
+      <camera name="angle" pos="0.0 0.0 2.0" fovy="37" mode="targetbody" target="table"/>
       <camera name="front" pos="0 0 0.8" fovy="65" mode="fixed" quat="0.7071 0.7071 0 0"/>
   </worldbody>
 </mujoco>

roboimi/assets/robots/diana_med.py

@@ -91,39 +91,3 @@ class BiDianaMed(ArmBase):
         """ Robot's init joint position. """
         return np.array([0.0, 0.0, 0.0, 1.57, 0.0, 0.0, 0.0])
     
-
-class BiDianaMedSocketPeg(ArmBase):
-    def __init__(self):
-        super().__init__(
-            name="Bidiana_socket_peg",
-            urdf_path="roboimi/assets/models/manipulators/DianaMed/DualDianaMed.urdf",
-            xml_path="roboimi/assets/models/manipulators/DianaMed/bi_diana_socket_peg_ee.xml",
-            gripper=None
-        )
-        self.left_arm = self.Arm(self, 'single', self.urdf_path)
-        self.left_arm.set_Arm_base_link('left_base_link')
-        self.left_arm.set_Arm_ee_link('left_link7')
-        self.left_arm.InitKDL
-        self.left_arm.joint_index = ['l_j1','l_j2','l_j3','l_j4','l_j5','l_j6','l_j7']
-        self.left_arm.gripper_index = ['l_finger_joint_left','r_finger_joint_left']
-        self.left_arm.actuator_index = ['a1_l','a2_l','a3_l','a4_l','a5_l','a6_l','a7_l','gripper_left']
-        self.left_arm.setArmInitPose(self.init_qpos)
-        self.arms.append(self.left_arm)
-        self.right_arm = self.Arm(self,'single', self.urdf_path)
-        self.right_arm.set_Arm_base_link('right_base_link')
-        self.right_arm.set_Arm_ee_link('right_link7')
-        self.right_arm.InitKDL
-        self.right_arm.joint_index = ['r_j1','r_j2','r_j3','r_j4','r_j5','r_j6','r_j7']
-        self.right_arm.gripper_index = ['l_finger_joint_right','r_finger_joint_right']
-        self.right_arm.actuator_index = ['a1_r','a2_r','a3_r','a4_r','a5_r','a6_r','a7_r','gripper_right']
-        self.right_arm.setArmInitPose(self.init_qpos)
-        self.arms.append(self.right_arm)
-        self.jnt_num = self.left_arm.jnt_num + self.right_arm.jnt_num
-        self.kp = 500 * np.ones(self.jnt_num)
-        self.kd = 44.57 * np.ones(self.jnt_num)
-
-    @property
-    def init_qpos(self):
-        """ Robot's init joint position. """
-        return np.array([0.0, 0.0, 0.0, 1.57, 0.0, 0.0, 0.0])
-    

roboimi/demos/diana_air_insert_policy.py

@@ -1,184 +0,0 @@
-import numpy as np
-from pyquaternion import Quaternion
-
-from roboimi.demos.diana_policy import PolicyBase
-
-
-class TestAirInsertPolicy(PolicyBase):
-    ACTION_OBJECT_Z_OFFSET = 0.078
-    SOCKET_GRASP_OFFSET = np.array([0.0, 0.0, 0.0], dtype=np.float64)
-    PEG_GRASP_OFFSET = np.array([0.0, 0.0, 0.0], dtype=np.float64)
-    SOCKET_OUTER_GRASP_STRATEGY = "socket_outer"
-    LEGACY_GRASP_STRATEGY = "legacy"
-    SOCKET_HOLD_Z = 0.85
-    PEG_INSERT_START_OFFSET = np.array([0.105, 0.0, 0.0], dtype=np.float64)
-    INSERT_END_T = 580
-    LEFT_SOCKET_GRIPPER_CLOSED = -100
-    RIGHT_PEG_GRIPPER_CLOSED = -100
-    SOCKET_APPROACH_Z = 1.05
-    EPISODE_END_T = 600
-
-    def __init__(self, inject_noise=False, grasp_strategy=SOCKET_OUTER_GRASP_STRATEGY):
-        super().__init__(inject_noise=inject_noise)
-        valid_strategies = {
-            self.SOCKET_OUTER_GRASP_STRATEGY,
-            self.LEGACY_GRASP_STRATEGY,
-        }
-        if grasp_strategy not in valid_strategies:
-            raise ValueError(
-                f"Unsupported air insert grasp_strategy={grasp_strategy!r}; "
-                f"expected one of {sorted(valid_strategies)}"
-            )
-        self.grasp_strategy = grasp_strategy
-
-    def generate_trajectory(self, task_state):
-        return self._generate_socket_peg_trajectory(task_state)
-
-    def _generate_socket_peg_trajectory(self, task_state):
-        socket_xyz = np.asarray(task_state["socket_pos"], dtype=np.float64)
-        peg_xyz = np.asarray(task_state["peg_pos"], dtype=np.float64)
-
-        init_mocap_pose_left = np.array(
-            [
-                -0.17297014,
-                1.00485877,
-                1.32773627,
-                7.06825181e-01,
-                8.20281078e-06,
-                -7.07388269e-01,
-                -5.20399313e-06,
-            ],
-            dtype=np.float64,
-        )
-        init_mocap_pose_right = np.array(
-            [
-                0.17297014,
-                0.9951369,
-                1.32773623,
-                2.59463975e-06,
-                7.07388269e-01,
-                5.59551158e-06,
-                7.06825181e-01,
-            ],
-            dtype=np.float64,
-        )
-
-        left_init_quat = Quaternion(init_mocap_pose_left[3:])
-        right_init_quat = Quaternion(init_mocap_pose_right[3:])
-
-        left_pick_quat = (
-            left_init_quat * Quaternion(axis=[0.0, 1.0, 0.0], degrees=45)
-        ).elements
-        right_pick_quat = (
-            right_init_quat * Quaternion(axis=[0.0, 1.0, 0.0], degrees=45)
-        ).elements
-
-        socket_hold_action = np.array(
-            [socket_xyz[0] - 0.078, socket_xyz[1], self.SOCKET_HOLD_Z], dtype=np.float64
-        )
-
-        peg_init_xyz = peg_xyz + np.array(
-            [0.078, 0.0, self.ACTION_OBJECT_Z_OFFSET + 0.01]
-        )
-        peg_lift_center = np.array(
-            [peg_xyz[0] + 0.078, socket_hold_action[1], self.SOCKET_HOLD_Z - 0.01],
-            dtype=np.float64,
-        )
-        # The front camera looks along +Y, so visual right-to-left insertion is
-        # world +X -> -X.  With the socket XML in identity orientation, its
-        # tunnel axis is local/world X, so the peg approaches from +X and stops
-        # when its leading face reaches the socket's internal pin.
-        peg_insert_end_center = np.array(
-            [
-                socket_hold_action[0] + 0.078 * 2 + 0.04 + 0.06 - 0.01,
-                socket_hold_action[1],
-                self.SOCKET_HOLD_Z - 0.01,
-            ],
-            dtype=np.float64,
-        )
-
-        self.left_trajectory = [
-            {
-                "t": 1,
-                "xyz": init_mocap_pose_left[:3],
-                "quat": init_mocap_pose_left[3:],
-                "gripper": 100,
-            },
-            {
-                "t": 130,
-                "xyz": socket_xyz
-                + np.array([-0.078, 0.0, self.ACTION_OBJECT_Z_OFFSET]),
-                "quat": left_pick_quat,
-                "gripper": 100,
-            },
-            {
-                "t": 180,
-                "xyz": socket_xyz
-                + np.array([-0.078, 0.0, self.ACTION_OBJECT_Z_OFFSET]),
-                "quat": left_pick_quat,
-                "gripper": self.LEFT_SOCKET_GRIPPER_CLOSED,
-            },
-            {
-                "t": 350,
-                "xyz": socket_hold_action,
-                "quat": left_pick_quat,
-                "gripper": self.LEFT_SOCKET_GRIPPER_CLOSED,
-            },
-            {
-                "t": self.EPISODE_END_T,
-                "xyz": socket_hold_action,
-                "quat": left_pick_quat,
-                "gripper": self.LEFT_SOCKET_GRIPPER_CLOSED,
-            },
-        ]
-
-        self.right_trajectory = [
-            {
-                "t": 1,
-                "xyz": init_mocap_pose_right[:3],
-                "quat": init_mocap_pose_right[3:],
-                "gripper": 100,
-            },
-            {
-                "t": 80,
-                "xyz": peg_init_xyz,
-                "quat": right_pick_quat,
-                "gripper": 100,
-            },
-            {
-                "t": 150,
-                "xyz": peg_init_xyz,
-                "quat": right_pick_quat,
-                "gripper": 100,
-            },
-            {
-                "t": 180,
-                "xyz": peg_init_xyz,
-                "quat": right_pick_quat,
-                "gripper": self.RIGHT_PEG_GRIPPER_CLOSED,
-            },
-            {
-                "t": 350,
-                "xyz": peg_init_xyz,
-                "quat": right_pick_quat,
-                "gripper": self.RIGHT_PEG_GRIPPER_CLOSED,
-            },
-            {
-                "t": 450,
-                "xyz": peg_lift_center,
-                "quat": right_pick_quat,
-                "gripper": self.RIGHT_PEG_GRIPPER_CLOSED,
-            },
-            {
-                "t": self.INSERT_END_T,
-                "xyz": peg_insert_end_center,
-                "quat": right_pick_quat,
-                "gripper": self.RIGHT_PEG_GRIPPER_CLOSED,
-            },
-            {
-                "t": self.EPISODE_END_T,
-                "xyz": peg_insert_end_center,
-                "quat": right_pick_quat,
-                "gripper": self.RIGHT_PEG_GRIPPER_CLOSED,
-            },
-        ]

roboimi/demos/diana_record_sim_episodes.py

@@ -2,11 +2,9 @@ import time
 import os
 import numpy as np
 from roboimi.envs.double_pos_ctrl_env import make_sim_env
-from roboimi.demos.diana_air_insert_policy import TestAirInsertPolicy
+from diana_policy import TestPickAndTransferPolicy
-from roboimi.demos.diana_policy import TestPickAndTransferPolicy
 import cv2
-from roboimi.utils.act_ex_utils import sample_air_insert_socket_peg_state, sample_transfer_pose
+from roboimi.utils.act_ex_utils import sample_transfer_pose
-from roboimi.utils.constants import SIM_TASK_CONFIGS
 from roboimi.utils.streaming_episode_writer import StreamingEpisodeWriter
 
 import pathlib
@@ -14,34 +12,16 @@ HOME_PATH = str(pathlib.Path(__file__).parent.resolve())
 DATASET_DIR = HOME_PATH + '/dataset'
 
 
-def sample_task_state(task_name):
+def main():
-    if task_name == 'sim_transfer':
+    task_name = 'sim_transfer'
-        return sample_transfer_pose()
+    dataset_dir = DATASET_DIR + '/sim_transfer'   #SIM_TASK_CONFIGS[task_name]['dataset_dir']
-    if task_name == 'sim_air_insert_socket_peg':
+    num_episodes = 100       #SIM_TASK_CONFIGS[task_name]['num_episodes']
-        return sample_air_insert_socket_peg_state()
-    raise NotImplementedError(f'Unsupported scripted rollout task: {task_name}')
-
-
-def make_policy(task_name, inject_noise=False, grasp_strategy=None):
-    if task_name == 'sim_transfer':
-        return TestPickAndTransferPolicy(inject_noise)
-    if task_name == 'sim_air_insert_socket_peg':
-        if grasp_strategy is None:
-            return TestAirInsertPolicy(inject_noise)
-        return TestAirInsertPolicy(inject_noise, grasp_strategy=grasp_strategy)
-    raise NotImplementedError(f'Unsupported scripted rollout task: {task_name}')
-
-
-def main(task_name='sim_transfer'):
-    task_cfg = SIM_TASK_CONFIGS[task_name]
-    dataset_dir = task_cfg['dataset_dir']
-    num_episodes = 100
     inject_noise = False
 
-    episode_len = task_cfg['episode_len']
+    episode_len  = 700      #SIM_TASK_CONFIGS[task_name]['episode_len']
-    camera_names = task_cfg['camera_names']
+    camera_names = ['angle','r_vis', 'top', 'front']   #SIM_TASK_CONFIGS[task_name]['camera_names']
     image_size = (256, 256)
-    if task_name in {'sim_transfer', 'sim_air_insert_socket_peg'}:
+    if task_name == 'sim_transfer':
         print(task_name)
     else:
         raise NotImplementedError
@@ -49,7 +29,7 @@ def main(task_name='sim_transfer'):
     success = []
     
     env = make_sim_env(task_name)
-    policy = make_policy(task_name, inject_noise=inject_noise)
+    policy = TestPickAndTransferPolicy(inject_noise)
 
     # 等待osmesa完全启动后再开始收集数据
     print("等待osmesa线程启动...")
@@ -61,8 +41,8 @@ def main(task_name='sim_transfer'):
         max_reward = float('-inf')
         print(f'\n{episode_idx=}')
         print('Rollout out EE space scripted policy')
-        task_state = sample_task_state(task_name)
+        box_pos = sample_transfer_pose()
-        env.reset(task_state)
+        env.reset(box_pos)
         episode_writer = StreamingEpisodeWriter(
             dataset_path=os.path.join(dataset_dir, f'episode_{episode_idx}.hdf5'),
             max_timesteps=episode_len,
@@ -70,7 +50,7 @@ def main(task_name='sim_transfer'):
             image_size=image_size,
         )
         for step in range(episode_len):
-            raw_action = policy.predict(task_state, step)
+            raw_action = policy.predict(box_pos,step)
             env.step(raw_action)
             env.render()
             sum_reward += env.rew

roboimi/demos/vla_scripts/eval_vla.py

@@ -30,10 +30,7 @@ from hydra.utils import instantiate
 from einops import rearrange
 
 from roboimi.envs.double_pos_ctrl_env import make_sim_env
-from roboimi.utils.act_ex_utils import (
+from roboimi.utils.act_ex_utils import sample_transfer_pose
-    sample_air_insert_socket_peg_state,
-    sample_transfer_pose,
-)
 from roboimi.vla.eval_utils import execute_policy_action
 
 sys.path.append(os.getcwd())
@@ -113,7 +110,6 @@ def prepare_observation(
     obs: Dict,
     camera_names: list,
     image_resize_shape: Optional[tuple[int, int]] = (224, 224),
-    task_description: Optional[str] = None,
 ) -> Dict:
     """
     将环境观测转换为 agent 格式。
@@ -140,34 +136,7 @@ def prepare_observation(
     # 转换 qpos: numpy -> tensor
     qpos = torch.from_numpy(obs['qpos']).float()
 
-    observation = {'qpos': qpos, 'images': images}
+    return {'qpos': qpos, 'images': images}
-    if 'task' in obs:
-        observation['task'] = obs['task']
-    elif task_description is not None:
-        observation['task'] = task_description
-    return observation
-
-
-def _normalize_resize_shape(shape) -> Optional[tuple[int, int]]:
-    if shape is None:
-        return None
-    normalized = tuple(int(v) for v in shape)
-    if len(normalized) != 2:
-        raise ValueError(f'image resize shape must contain exactly two values, got {normalized}')
-    return normalized
-
-
-def _resolve_eval_image_resize_shape(cfg: DictConfig) -> Optional[tuple[int, int]]:
-    image_resize_shape = cfg.get('data', {}).get('image_resize_shape', (224, 224))
-    agent_cfg = cfg.agent
-    if 'eval_image_resize_shape' in agent_cfg:
-        return _normalize_resize_shape(agent_cfg.get('eval_image_resize_shape'))
-    for backbone_key in ('vision_backbone', 'condition_encoder'):
-        backbone_cfg = agent_cfg.get(backbone_key, None)
-        if backbone_cfg is not None and 'eval_image_resize_shape' in backbone_cfg:
-            image_resize_shape = backbone_cfg.get('eval_image_resize_shape')
-            break
-    return _normalize_resize_shape(image_resize_shape)
 
 
 def _resolve_policy_camera_names(cfg: DictConfig) -> list[str]:
@@ -187,7 +156,6 @@ def _new_local_policy_queues(obs_horizon: int) -> dict[str, deque]:
     return {
         'qpos': deque(maxlen=int(obs_horizon)),
         'images': deque(maxlen=int(obs_horizon)),
-        'task': deque(maxlen=int(obs_horizon)),
         'action': deque(),
     }
 
@@ -203,8 +171,6 @@ def _populate_local_policy_queues(
             camera_name: image.detach().clone()
             for camera_name, image in observation['images'].items()
         })
-    if 'task' in observation:
-        queues['task'].append(observation['task'])
 
 
 def _prepare_local_policy_batch(
@@ -233,10 +199,7 @@ def _prepare_local_policy_batch(
         ).unsqueeze(0)
         for camera_name in ordered_camera_names
     }
-    batch = {'qpos': batch_qpos, 'images': batch_images}
+    return {'qpos': batch_qpos, 'images': batch_images}
-    if queues.get('task'):
-        batch['task'] = [list(queues['task'])[-1]]
-    return batch
 
 
 def _enqueue_predicted_actions(
@@ -244,14 +207,13 @@ def _enqueue_predicted_actions(
     predicted_actions: Any,
     obs_horizon: int,
     num_action_steps: int,
-    action_chunk_start: Optional[int] = None,
 ) -> None:
     if isinstance(predicted_actions, np.ndarray):
         predicted_actions = torch.from_numpy(predicted_actions)
     if predicted_actions.ndim == 2:
         predicted_actions = predicted_actions.unsqueeze(0)
 
-    start = int(obs_horizon) - 1 if action_chunk_start is None else int(action_chunk_start)
+    start = int(obs_horizon) - 1
     end = start + int(num_action_steps)
     executable_actions = predicted_actions[:, start:end]
     for action_index in range(executable_actions.shape[1]):
@@ -261,29 +223,23 @@ def _enqueue_predicted_actions(
 
 
 def _serialize_policy_batch(batch: Dict[str, torch.Tensor]) -> dict[str, Any]:
-    serialized = {
+    return {
         'qpos': batch['qpos'].detach().cpu().numpy().astype(np.float32, copy=True),
         'images': {
             camera_name: image.detach().cpu().numpy().astype(np.float32, copy=True)
             for camera_name, image in batch['images'].items()
         },
     }
-    if 'task' in batch:
-        serialized['task'] = batch['task']
-    return serialized
 
 
 def _deserialize_policy_batch(batch: dict[str, Any], device: str) -> Dict[str, torch.Tensor]:
-    deserialized = {
+    return {
         'qpos': torch.as_tensor(batch['qpos'], dtype=torch.float32, device=device),
         'images': {
             camera_name: torch.as_tensor(image, dtype=torch.float32, device=device)
             for camera_name, image in batch['images'].items()
         },
     }
-    if 'task' in batch:
-        deserialized['task'] = batch['task']
-    return deserialized
 
 
 class _LocalPolicyRunner:
@@ -319,7 +275,6 @@ class _RemotePolicyRunner:
         camera_names: list[str],
         obs_horizon: int,
         num_action_steps: int,
-        action_chunk_start: Optional[int] = None,
         response_timeout_s: float = 30.0,
     ):
         self.worker_index = int(worker_index)
@@ -329,7 +284,6 @@ class _RemotePolicyRunner:
         self.camera_names = list(camera_names)
         self.obs_horizon = int(obs_horizon)
         self.num_action_steps = int(num_action_steps)
-        self.action_chunk_start = None if action_chunk_start is None else int(action_chunk_start)
         self.response_timeout_s = float(response_timeout_s)
         self.local_queues = _new_local_policy_queues(self.obs_horizon)
         self.uses_local_model = False
@@ -376,7 +330,6 @@ class _RemotePolicyRunner:
                 predicted_actions=response['actions'],
                 obs_horizon=self.obs_horizon,
                 num_action_steps=self.num_action_steps,
-                action_chunk_start=self.action_chunk_start,
             )
 
         if not self.local_queues['action']:
@@ -1011,14 +964,6 @@ def _close_env(env):
         viewer.close()
 
 
-def _sample_task_reset_state(task_name: str):
-    if task_name == 'sim_air_insert_socket_peg':
-        return sample_air_insert_socket_peg_state()
-    if 'sim_transfer' in task_name:
-        return sample_transfer_pose()
-    raise NotImplementedError(f'Unsupported eval task reset sampling: {task_name}')
-
-
 def _print_eval_config(cfg: DictConfig):
     # 打印配置
     print("=" * 80)
@@ -1059,8 +1004,6 @@ def _run_eval_episode_plans(
     eval_cfg = cfg.eval
     device = str(eval_cfg.device)
     camera_names = list(eval_cfg.camera_names)
-    image_resize_shape = _resolve_eval_image_resize_shape(cfg)
-    task_description = eval_cfg.get('task_description', None)
     artifact_paths = artifact_paths or _resolve_artifact_paths(eval_cfg)
     video_recorder = _RolloutVideoRecorder(
         output_path=artifact_paths['video_mp4'],
@@ -1087,19 +1030,17 @@ def _run_eval_episode_plans(
         episode_max_rewards = []
         for plan in episode_plans:
             episode_idx = int(plan['episode_index'])
-            task_state = plan.get('box_pos')
+            box_pos = plan.get('box_pos')
-            if task_state is None:
+            if box_pos is None:
-                task_state = _sample_task_reset_state(str(eval_cfg.task_name))
+                box_pos = sample_transfer_pose()
-            elif isinstance(task_state, np.ndarray):
+            box_pos = np.asarray(box_pos, dtype=np.float32)
-                task_state = np.asarray(task_state, dtype=np.float32)
-
 
             if show_progress:
                 print(f"\n{'='*60}")
                 print(f"回合 {episode_idx + 1}/{eval_cfg.num_episodes}")
                 print(f"{'='*60}\n")
 
-            env.reset(task_state)
+            env.reset(box_pos)
 
             # 为新回合重置 rollout policy 状态
             policy_runner.reset()
@@ -1135,12 +1076,7 @@ def _run_eval_episode_plans(
                     video_recorder.write(video_frame)
 
                     # 准备给 agent 的观测
-                    observation = prepare_observation(
+                    observation = prepare_observation(obs, camera_names)
-                        obs,
-                        camera_names,
-                        image_resize_shape=image_resize_shape,
-                        task_description=task_description,
-                    )
                     end_preprocess = time.perf_counter()
 
                     # 选择动作（本地 agent 或远端 inference server）
@@ -1413,7 +1349,6 @@ def _run_remote_eval_worker(
     eval_cfg = cfg.eval
     agent_cfg = cfg.agent
     num_action_steps = int(agent_cfg.get('num_action_steps', eval_cfg.get('num_queries', 1)))
-    action_chunk_start = agent_cfg.get('action_chunk_start', None)
     policy_runner = _RemotePolicyRunner(
         worker_index=worker_index,
         server_index=server_index,
@@ -1422,7 +1357,6 @@ def _run_remote_eval_worker(
         camera_names=_resolve_policy_camera_names(cfg),
         obs_horizon=int(agent_cfg.get('obs_horizon', eval_cfg.obs_horizon)),
         num_action_steps=num_action_steps,
-        action_chunk_start=action_chunk_start,
         response_timeout_s=float(eval_cfg.get('response_timeout_s', 300.0)),
     )
     return _run_eval_episode_plans(

roboimi/demos/vla_scripts/train_vla.py

@@ -71,19 +71,6 @@ from hydra.utils import instantiate
 
 log = logging.getLogger(__name__)
 
-
-SMOLVLA_NATIVE_TRAINING_PRESET = {
-    'lr': 1e-4,
-    'betas': (0.9, 0.95),
-    'eps': 1e-8,
-    'weight_decay': 1e-10,
-    'grad_clip': 10.0,
-    'warmup_steps': 1000,
-    'scheduler_type': 'cosine',
-    'scheduler_decay_steps': 30000,
-    'scheduler_decay_lr': 2.5e-6,
-}
-
 # 注册列表长度解析器（用于配置中如 ${len:${data.camera_names}}）
 if not OmegaConf.has_resolver("len"):
     OmegaConf.register_new_resolver("len", lambda x: len(x))
@@ -131,6 +118,127 @@ def recursive_to_device(data, device):
     return data
 
 
+def build_agent_input(batch_data):
+    agent_input = {
+        'images': {
+            cam_name.replace('observation.', ''): value
+            for cam_name, value in batch_data.items()
+            if cam_name.startswith('observation.') and cam_name != 'observation.state'
+        },
+        'qpos': batch_data['observation.state'],
+        'action': batch_data['action'],
+    }
+
+    if 'action_is_pad' in batch_data:
+        agent_input['action_is_pad'] = batch_data['action_is_pad']
+
+    lewm_images = {
+        cam_name.replace('lewm.observation.', ''): value
+        for cam_name, value in batch_data.items()
+        if cam_name.startswith('lewm.observation.') and cam_name != 'lewm.observation.state'
+    }
+    if lewm_images:
+        agent_input['lewm_images'] = lewm_images
+    if 'lewm.observation.state' in batch_data:
+        agent_input['lewm_qpos'] = batch_data['lewm.observation.state']
+
+    lewm_future_images = {
+        cam_name.replace('lewm.future.', ''): value
+        for cam_name, value in batch_data.items()
+        if cam_name.startswith('lewm.future.') and cam_name != 'lewm.future.state'
+    }
+    if lewm_future_images:
+        agent_input['lewm_future_images'] = lewm_future_images
+    if 'lewm.future.state' in batch_data:
+        agent_input['lewm_future_qpos'] = batch_data['lewm.future.state']
+
+    return agent_input
+
+
+def _instantiate_dataset(cfg, dataset_image_resize_shape, episode_indices=None):
+    kwargs = {'image_resize_shape': dataset_image_resize_shape}
+    if episode_indices is not None:
+        kwargs['episode_indices'] = episode_indices
+    return instantiate(cfg.data, **kwargs)
+
+
+def build_train_val_datasets(cfg, dataset_image_resize_shape):
+    val_episode_indices = cfg.train.get('val_episode_indices', None)
+    if val_episode_indices:
+        dataset = _instantiate_dataset(cfg, dataset_image_resize_shape)
+        available_episode_indices = list(getattr(dataset, 'available_episode_indices', []))
+        if not available_episode_indices:
+            raise ValueError('显式 val_episode_indices 需要数据集暴露 available_episode_indices')
+        requested_val_episode_indices = sorted(int(idx) for idx in val_episode_indices)
+        available_set = set(available_episode_indices)
+        missing = sorted(set(requested_val_episode_indices) - available_set)
+        if missing:
+            raise ValueError(
+                f'val_episode_indices {missing} 不存在于数据集可用 episodes {available_episode_indices}'
+            )
+        train_episode_indices = [
+            idx for idx in available_episode_indices
+            if idx not in set(requested_val_episode_indices)
+        ]
+        if not train_episode_indices:
+            raise ValueError('显式 val_episode_indices 不能覆盖全部 episodes，训练集将为空')
+
+        train_dataset = _instantiate_dataset(
+            cfg,
+            dataset_image_resize_shape,
+            episode_indices=train_episode_indices,
+        )
+        val_dataset = _instantiate_dataset(
+            cfg,
+            dataset_image_resize_shape,
+            episode_indices=requested_val_episode_indices,
+        )
+        return dataset, train_dataset, val_dataset, requested_val_episode_indices
+
+    dataset = _instantiate_dataset(cfg, dataset_image_resize_shape)
+    val_split = float(cfg.train.get('val_split', 0.1))
+    seed = int(cfg.train.get('seed', 42))
+    val_size = int(len(dataset) * val_split)
+    train_size = len(dataset) - val_size
+    if val_size > 0:
+        train_dataset, val_dataset = random_split(
+            dataset,
+            [train_size, val_size],
+            generator=torch.Generator().manual_seed(seed)
+        )
+    else:
+        train_dataset, val_dataset = dataset, None
+    return dataset, train_dataset, val_dataset, None
+
+
+def compute_action_mse_validation(agent, val_loader, device):
+    if val_loader is None:
+        return None
+
+    was_training = agent.training
+    agent.eval()
+    total_squared_error = 0.0
+    total_count = 0.0
+    with torch.no_grad():
+        for val_batch in val_loader:
+            val_batch = recursive_to_device(val_batch, device)
+            val_input = build_agent_input(val_batch)
+            pred_actions = agent.predict_action_chunk(val_input)
+            target_actions = val_input['action']
+            squared_error = (pred_actions - target_actions).pow(2)
+            action_is_pad = val_input.get('action_is_pad', None)
+            if action_is_pad is not None:
+                mask = (~action_is_pad).unsqueeze(-1).to(squared_error.dtype)
+                total_squared_error += (squared_error * mask).sum().item()
+                total_count += mask.sum().item() * squared_error.shape[-1]
+            else:
+                total_squared_error += squared_error.sum().item()
+                total_count += target_actions.numel()
+    if was_training:
+        agent.train()
+    return total_squared_error / max(total_count, 1.0)
+
+
 def resolve_resume_checkpoint(resume_ckpt, checkpoint_dir):
     """
     解析恢复训练用的 checkpoint 路径。
@@ -167,7 +275,7 @@ def get_lr_schedule_with_warmup(optimizer, warmup_steps, max_steps, scheduler_ty
     Args:
         optimizer: PyTorch 优化器
         warmup_steps: 预热步数
-        max_steps: 余弦衰减步数
+        max_steps: 总训练步数
         scheduler_type: 预热后的调度器类型 ('cosine' 或 'constant')
         min_lr: 最小学习率（用于余弦衰减）
 
@@ -180,24 +288,16 @@ def get_lr_schedule_with_warmup(optimizer, warmup_steps, max_steps, scheduler_ty
     min_lr_ratio = min_lr / base_lr if base_lr > 0 else 0.0
 
     def lr_lambda(step):
-        # LeRobot CosineDecayWithWarmupSchedulerConfig 的线性预热：
+        # 预热阶段：从 0 线性增加到 1
-        # 从一个很小的非零 LR 开始，避免首步完全为 0。
         if step < warmup_steps:
-            if step <= 0:
+            return float(step) / float(max(1, warmup_steps))
-                return 1.0 / float(max(1, warmup_steps + 1))
-            frac = 1.0 - float(step) / float(max(1, warmup_steps))
-            return (1.0 / float(max(1, warmup_steps + 1)) - 1.0) * frac + 1.0
 
         # 预热后阶段
         if scheduler_type == 'cosine':
-            # 与 LeRobot SmolVLA/PI0 的 CosineDecayWithWarmupSchedulerConfig 对齐：
+            # 从 1 到 min_lr_ratio 的余弦退火
-            # 1) 余弦衰减步数是固定的 num_decay_steps（SmolVLA 默认 30k），
+            progress = float(step - warmup_steps) / float(max(1, max_steps - warmup_steps))
-            # 2) 超过 decay_steps 后 clamp 到 decay_lr，不能继续 cos() 进入下一周期；
+            cosine_decay = 0.5 * (1.0 + math.cos(math.pi * progress))
-            #    否则 150k 训练会显示成“正弦波”式反复升降。
+            return max(min_lr_ratio, cosine_decay)
-            decay_steps = max(1, int(max_steps))
-            clamped_step = min(max(int(step), 0), decay_steps)
-            cosine_decay = 0.5 * (1.0 + math.cos(math.pi * clamped_step / decay_steps))
-            return (1.0 - min_lr_ratio) * cosine_decay + min_lr_ratio
         else:
             # 恒定学习率
             return 1.0
@@ -205,119 +305,7 @@ def get_lr_schedule_with_warmup(optimizer, warmup_steps, max_steps, scheduler_ty
     return LambdaLR(optimizer, lr_lambda)
 
 
-def _is_smolvla_native_agent_config(agent_cfg) -> bool:
+def build_training_optimizer(agent, lr, weight_decay):
-    target = str(agent_cfg.get('_target_', '')) if hasattr(agent_cfg, 'get') else ''
-    return target == 'roboimi.vla.agent_smolvla_native.SmolVLANativeAgent'
-
-
-def resolve_training_recipe(cfg):
-    """Return optimizer/scheduler knobs, using LeRobot SmolVLA defaults for smolvla_native."""
-    if _is_smolvla_native_agent_config(cfg.agent):
-        preset = SMOLVLA_NATIVE_TRAINING_PRESET
-        scheduler_steps = int(cfg.train.get(
-            'scheduler_decay_steps',
-            cfg.train.get('lr_scheduler_steps', cfg.train.max_steps),
-        ))
-        return {
-            'lr': float(preset['lr']),
-            'betas': tuple(preset['betas']),
-            'eps': float(preset['eps']),
-            'weight_decay': float(preset['weight_decay']),
-            'grad_clip': float(preset['grad_clip']),
-            'warmup_steps': int(preset['warmup_steps']),
-            'scheduler_type': str(preset['scheduler_type']),
-            'scheduler_steps': scheduler_steps,
-            'min_lr': float(preset['scheduler_decay_lr']),
-            'preset_name': 'lerobot_smolvla',
-        }
-
-    return {
-        'lr': float(cfg.train.lr),
-        'betas': (0.9, 0.999),
-        'eps': 1e-8,
-        'weight_decay': float(cfg.train.get('weight_decay', 1e-5)),
-        'grad_clip': float(cfg.train.get('grad_clip', 1.0)),
-        'warmup_steps': int(cfg.train.get('warmup_steps', 500)),
-        'scheduler_type': str(cfg.train.get('scheduler_type', 'cosine')),
-        'scheduler_steps': int(cfg.train.max_steps),
-        'min_lr': float(cfg.train.get('min_lr', 1e-6)),
-        'preset_name': 'config',
-    }
-
-
-def _instantiate_dataset(cfg, dataset_image_resize_shape, episode_indices=None):
-    kwargs = {'image_resize_shape': dataset_image_resize_shape}
-    if episode_indices is not None:
-        kwargs['episode_indices'] = episode_indices
-    return instantiate(cfg.data, **kwargs)
-
-
-def _resolve_dataset_image_resize_shape(cfg):
-    dataset_image_resize_shape = cfg.data.get('image_resize_shape', (224, 224))
-    agent_cfg = cfg.agent
-    if 'dataset_image_resize_shape' in agent_cfg:
-        return agent_cfg.get('dataset_image_resize_shape')
-    for backbone_key in ('vision_backbone', 'condition_encoder'):
-        backbone_cfg = agent_cfg.get(backbone_key, None)
-        if backbone_cfg is not None and 'dataset_image_resize_shape' in backbone_cfg:
-            dataset_image_resize_shape = backbone_cfg.get('dataset_image_resize_shape')
-            break
-    return dataset_image_resize_shape
-
-
-def build_train_val_datasets(cfg, dataset_image_resize_shape):
-    val_episode_indices = cfg.train.get('val_episode_indices', None)
-    if val_episode_indices:
-        dataset = _instantiate_dataset(cfg, dataset_image_resize_shape)
-        available_episode_indices = list(getattr(dataset, 'available_episode_indices', []))
-        if not available_episode_indices:
-            raise ValueError('显式 val_episode_indices 需要数据集暴露 available_episode_indices')
-
-        requested_val_episode_indices = sorted(int(idx) for idx in val_episode_indices)
-        available_set = set(available_episode_indices)
-        missing = sorted(set(requested_val_episode_indices) - available_set)
-        if missing:
-            raise ValueError(
-                f'val_episode_indices {missing} 不存在于数据集可用 episodes {available_episode_indices}'
-            )
-
-        val_set = set(requested_val_episode_indices)
-        train_episode_indices = [
-            idx for idx in available_episode_indices
-            if idx not in val_set
-        ]
-        if not train_episode_indices:
-            raise ValueError('显式 val_episode_indices 不能覆盖全部 episodes，训练集将为空')
-
-        train_dataset = _instantiate_dataset(
-            cfg,
-            dataset_image_resize_shape,
-            episode_indices=train_episode_indices,
-        )
-        val_dataset = _instantiate_dataset(
-            cfg,
-            dataset_image_resize_shape,
-            episode_indices=requested_val_episode_indices,
-        )
-        return dataset, train_dataset, val_dataset, requested_val_episode_indices
-
-    dataset = _instantiate_dataset(cfg, dataset_image_resize_shape)
-    val_split = float(cfg.train.get('val_split', 0.1))
-    seed = int(cfg.train.get('seed', 42))
-    val_size = int(len(dataset) * val_split)
-    train_size = len(dataset) - val_size
-    if val_size > 0:
-        train_dataset, val_dataset = random_split(
-            dataset,
-            [train_size, val_size],
-            generator=torch.Generator().manual_seed(seed)
-        )
-    else:
-        train_dataset, val_dataset = dataset, None
-    return dataset, train_dataset, val_dataset, None
-
-
-def build_training_optimizer(agent, lr, weight_decay, betas=(0.9, 0.999), eps=1e-8):
     """为训练脚本构建优化器，优先复用任意 head 自带的参数分组。"""
     trainable_params = [param for param in agent.parameters() if param.requires_grad]
     noise_pred_net = getattr(agent, 'noise_pred_net', None)
@@ -325,7 +313,7 @@ def build_training_optimizer(agent, lr, weight_decay, betas=(0.9, 0.999), eps=1e
     use_head_groups = callable(get_optim_groups)
 
     if not use_head_groups:
-        return AdamW(trainable_params, lr=lr, weight_decay=weight_decay, betas=betas, eps=eps)
+        return AdamW(trainable_params, lr=lr, weight_decay=weight_decay)
 
     head_groups = []
     grouped_param_ids = set()
@@ -367,7 +355,33 @@ def build_training_optimizer(agent, lr, weight_decay, betas=(0.9, 0.999), eps=1e
     if grouped_param_ids != all_trainable_param_ids:
         raise ValueError('Optimizer parameter groups must include each trainable parameter exactly once')
 
-    return AdamW(optim_groups, lr=lr, weight_decay=weight_decay, betas=betas, eps=eps)
+    return AdamW(optim_groups, lr=lr, weight_decay=weight_decay)
+
+
+def load_state_dict_ignoring_shape_mismatches(module, incoming_state_dict):
+    """Load only checkpoint tensors whose keys exist locally and whose shapes match."""
+    current_state_dict = module.state_dict()
+    compatible_state_dict = {}
+    mismatched_keys = []
+    missing_keys = []
+
+    for key, value in incoming_state_dict.items():
+        if key not in current_state_dict:
+            missing_keys.append(key)
+            continue
+        if current_state_dict[key].shape != value.shape:
+            mismatched_keys.append(key)
+            continue
+        compatible_state_dict[key] = value
+
+    merged_state_dict = dict(current_state_dict)
+    merged_state_dict.update(compatible_state_dict)
+    module.load_state_dict(merged_state_dict, strict=True)
+    return {
+        'loaded_keys': sorted(compatible_state_dict.keys()),
+        'missing_keys': sorted(missing_keys),
+        'mismatched_keys': sorted(mismatched_keys),
+    }
 
 
 def _init_swanlab(cfg):
@@ -414,12 +428,7 @@ def _init_swanlab(cfg):
         init_kwargs['experiment_name'] = run_name
 
     try:
-        run = swanlab.init(**init_kwargs)
+        swanlab.init(**init_kwargs)
-        if run is not None:
-            try:
-                setattr(swanlab, "_roboimi_swanlab_run", run)
-            except Exception:
-                pass
     except Exception as exc:
         raise RuntimeError(
             f"SwanLab logging is enabled, but SwanLab init/login failed: {exc}"
@@ -428,24 +437,6 @@ def _init_swanlab(cfg):
     return swanlab
 
 
-def _log_swanlab_init_details(swanlab_module):
-    if swanlab_module is None:
-        return
-    run = getattr(swanlab_module, "_roboimi_swanlab_run", None)
-    if run is None:
-        run = getattr(swanlab_module, "run", None)
-    if run is None:
-        return
-
-    url = getattr(run, "url", None)
-    swanlog_dir = getattr(run, "swanlog_dir", None)
-    log.info(
-        "🦢 SwanLab initialized%s%s",
-        f" | url={url}" if url else "",
-        f" | swanlog_dir={swanlog_dir}" if swanlog_dir else "",
-    )
-
-
 def _log_to_swanlab(swanlab_module, payload, step=None):
     if swanlab_module is None:
         return
@@ -524,13 +515,7 @@ def _run_training(cfg: DictConfig):
     log.info(f"🚀 开始 VLA 训练 (设备: {cfg.train.device})")
     _configure_cuda_runtime(cfg)
     swanlab_module = _init_swanlab(cfg)
-    _log_swanlab_init_details(swanlab_module)
     try:
-        action_mse_val_freq_epochs = int(cfg.train.get('action_mse_val_freq_epochs', 0) or 0)
-        explicit_val_episode_indices = cfg.train.get('val_episode_indices', None)
-        if action_mse_val_freq_epochs > 0 and not explicit_val_episode_indices:
-            raise ValueError('action_mse_val_freq_epochs > 0 requires train.val_episode_indices')
-
         # 创建检查点目录
         run_output_dir = _resolve_run_output_dir()
         checkpoint_dir = run_output_dir / "checkpoints"
@@ -542,7 +527,10 @@ def _run_training(cfg: DictConfig):
         # =========================================================================
         log.info("📦 加载数据集...")
         try:
-            dataset_image_resize_shape = _resolve_dataset_image_resize_shape(cfg)
+            dataset_image_resize_shape = cfg.data.get('image_resize_shape', (224, 224))
+            vision_backbone_cfg = cfg.agent.get('vision_backbone', None)
+            if vision_backbone_cfg is not None and 'dataset_image_resize_shape' in vision_backbone_cfg:
+                dataset_image_resize_shape = vision_backbone_cfg.get('dataset_image_resize_shape')
             dataset, train_dataset, val_dataset, explicit_val_episode_indices = (
                 build_train_val_datasets(cfg, dataset_image_resize_shape)
             )
@@ -668,18 +656,23 @@ def _run_training(cfg: DictConfig):
                 try:
                     checkpoint = torch.load(ckpt_path, map_location=cfg.train.device)
 
-                    # 只加载模型权重（不加载 optimizer、scheduler）
+                    load_info = load_state_dict_ignoring_shape_mismatches(
-                    missing_keys, unexpected_keys = agent.load_state_dict(
+                        agent,
                         checkpoint['model_state_dict'],
-                        strict=False  # 允许部分加载（结构不完全匹配时）
                     )
 
                     log.info(f"✅ [Finetune] 模型权重加载成功")
 
-                    if missing_keys:
+                    if load_info['missing_keys']:
-                        log.warning(f"⚠️  [Finetune] 缺少的键 ({len(missing_keys)} 个): {missing_keys[:5]}...")
+                        log.warning(
-                    if unexpected_keys:
+                            f"⚠️  [Finetune] checkpoint 中存在本地模型没有的键 ({len(load_info['missing_keys'])} 个): "
-                        log.warning(f"⚠️  [Finetune] 多余的键 ({len(unexpected_keys)} 个): {unexpected_keys[:5]}...")
+                            f"{load_info['missing_keys'][:5]}..."
+                        )
+                    if load_info['mismatched_keys']:
+                        log.warning(
+                            f"⚠️  [Finetune] 因形状不匹配而跳过的键 ({len(load_info['mismatched_keys'])} 个): "
+                            f"{load_info['mismatched_keys'][:5]}..."
+                        )
 
                     log.info(f"📊 [Finetune] 预训练信息: 步骤={checkpoint.get('step', 'N/A')}, 损失={checkpoint.get('loss', 'N/A')}")
                     log.info(f"📈 [Finetune] 使用新的训练配置（lr={cfg.train.lr}, max_steps={cfg.train.max_steps}）")
@@ -694,35 +687,21 @@ def _run_training(cfg: DictConfig):
         # =========================================================================
         # 4. 设置优化器与学习率调度器
         # =========================================================================
-        training_recipe = resolve_training_recipe(cfg)
+        weight_decay = float(cfg.train.get('weight_decay', 1e-5))
-        weight_decay = training_recipe['weight_decay']
+        grad_clip = float(cfg.train.get('grad_clip', 1.0))
-        grad_clip = training_recipe['grad_clip']
+
-        train_lr = training_recipe['lr']
+        optimizer = build_training_optimizer(agent, lr=cfg.train.lr, weight_decay=weight_decay)
-        optimizer = build_training_optimizer(
+        log.info(f"🔧 优化器: AdamW (学习率={cfg.train.lr}, weight_decay={weight_decay})")
-            agent,
-            lr=train_lr,
-            weight_decay=weight_decay,
-            betas=training_recipe['betas'],
-            eps=training_recipe['eps'],
-        )
-        log.info(
-            "🔧 优化器: AdamW (preset=%s, 学习率=%s, betas=%s, eps=%s, weight_decay=%s)",
-            training_recipe['preset_name'],
-            train_lr,
-            training_recipe['betas'],
-            training_recipe['eps'],
-            weight_decay,
-        )
 
         # 设置带预热的学習率调度器
-        warmup_steps = training_recipe['warmup_steps']
+        warmup_steps = int(cfg.train.get('warmup_steps', 500))
-        scheduler_type = training_recipe['scheduler_type']
+        scheduler_type = cfg.train.get('scheduler_type', 'cosine')
-        min_lr = training_recipe['min_lr']
+        min_lr = float(cfg.train.get('min_lr', 1e-6))
 
         scheduler = get_lr_schedule_with_warmup(
             optimizer,
             warmup_steps=warmup_steps,
-            max_steps=training_recipe['scheduler_steps'],
+            max_steps=cfg.train.max_steps,
             scheduler_type=scheduler_type,
             min_lr=min_lr
         )
@@ -816,26 +795,6 @@ def _run_training(cfg: DictConfig):
         # =========================================================================
         log.info("🏋️ 开始训练循环...")
 
-        def build_agent_input(batch_data):
-            """构建 agent 输入格式"""
-            images = {}
-            # SimpleRobotDataset 返回 observation.{cam_name} 格式
-            for cam_name in cfg.data.camera_names:
-                key = f"observation.{cam_name}"
-                if key in batch_data:
-                    images[cam_name] = batch_data[key]
-
-            agent_input = {
-                'images': images,
-                'qpos': batch_data['observation.state'],  # SimpleRobotDataset 使用 observation.state
-                'action': batch_data['action'],
-                'action_is_pad': batch_data.get('action_is_pad', None)  # 传递padding mask
-            }
-            if 'task' in batch_data:
-                agent_input['task'] = batch_data['task']
-
-            return agent_input
-
         def save_checkpoint(checkpoint_path: Path, step: int, loss_value, val_loss=None, rollout_avg_reward=None):
             agent_stats = agent.get_normalization_stats()
             torch.save({
@@ -928,6 +887,7 @@ def _run_training(cfg: DictConfig):
         pbar = tqdm(range(start_step, cfg.train.max_steps), desc="训练中", ncols=100)
 
         steps_per_epoch = len(train_loader)
+        action_mse_val_freq_epochs = int(cfg.train.get('action_mse_val_freq_epochs', 0) or 0)
         rollout_val_freq_epochs = int(cfg.train.get('rollout_val_freq_epochs', 0) or 0)
         rollout_validation_enabled = rollout_val_freq_epochs > 0
         best_loss = resume_best_loss
@@ -1006,6 +966,15 @@ def _run_training(cfg: DictConfig):
                     },
                     step=step,
                 )
+                if hasattr(agent, 'get_last_loss_breakdown'):
+                    loss_breakdown = agent.get_last_loss_breakdown()
+                    extra_train_metrics = {
+                        f"train/{key}": value
+                        for key, value in loss_breakdown.items()
+                        if value is not None and key != 'loss'
+                    }
+                    if extra_train_metrics:
+                        _log_to_swanlab(swanlab_module, extra_train_metrics, step=step)
 
             # =====================================================================
             # 检查点保存与验证
@@ -1081,31 +1050,6 @@ def _run_training(cfg: DictConfig):
                 completed_steps // steps_per_epoch
                 if steps_per_epoch > 0 else 0
             )
-            should_run_action_mse_val = (
-                val_loader is not None
-                and explicit_val_episode_indices is not None
-                and action_mse_val_freq_epochs > 0
-                and steps_per_epoch > 0
-                and completed_steps % steps_per_epoch == 0
-                and completed_epoch > 0
-                and completed_epoch % action_mse_val_freq_epochs == 0
-            )
-            if should_run_action_mse_val:
-                val_loss = run_validation()
-                if val_loss is not None:
-                    log.info(
-                        f"步骤 {step}/{cfg.train.max_steps} | Epoch {completed_epoch} "
-                        f"held-out action MSE: {val_loss:.6f}"
-                    )
-                    _log_to_swanlab(
-                        swanlab_module,
-                        {
-                            'val/action_mse': val_loss,
-                            'val/epoch': completed_epoch,
-                        },
-                        step=step,
-                    )
-
             should_run_epoch_rollout = (
                 rollout_validation_enabled
                 and steps_per_epoch > 0
@@ -1113,6 +1057,33 @@ def _run_training(cfg: DictConfig):
                 and completed_epoch > 0
                 and completed_epoch % rollout_val_freq_epochs == 0
             )
+            should_run_action_mse_validation = (
+                action_mse_val_freq_epochs > 0
+                and val_loader is not None
+                and steps_per_epoch > 0
+                and completed_steps % steps_per_epoch == 0
+                and completed_epoch > 0
+                and completed_epoch % action_mse_val_freq_epochs == 0
+            )
+            if should_run_action_mse_validation:
+                action_mse = compute_action_mse_validation(
+                    agent,
+                    val_loader,
+                    cfg.train.device,
+                )
+                if action_mse is not None:
+                    log.info(
+                        f"步骤 {step}/{cfg.train.max_steps} | Epoch {completed_epoch} "
+                        f"held-out action MSE: {action_mse:.6f}"
+                    )
+                    _log_to_swanlab(
+                        swanlab_module,
+                        {
+                            'val/action_mse': action_mse,
+                            'val/action_mse_epoch': completed_epoch,
+                        },
+                        step=step,
+                    )
             if should_run_epoch_rollout:
                 if checkpoint_path is None:
                     checkpoint_path = checkpoint_dir / f"vla_model_step_{step}.pt"

roboimi/envs/double_air_insert_env.py

@@ -1,157 +0,0 @@
-import copy as cp
-import time
-
-import numpy as np
-
-from roboimi.envs.double_base import DualDianaMed
-from roboimi.envs.double_pos_ctrl_env import DualDianaMed_Pos_Ctrl
-
-
-SOCKET_JOINT_NAME = "blue_socket_joint"
-PEG_JOINT_NAME = "red_peg_joint"
-REQUIRED_TASK_STATE_KEYS = ("socket_pos", "socket_quat", "peg_pos", "peg_quat")
-SOCKET_GEOM_NAMES = ("socket-1", "socket-2", "socket-3", "socket-4")
-SOCKET_SUCCESS_GEOM_NAMES = ("pin",)
-SOCKET_BODY_GEOM_NAMES = SOCKET_GEOM_NAMES + SOCKET_SUCCESS_GEOM_NAMES
-PEG_GEOM_NAMES = ("red_peg",)
-LEFT_GRIPPER_GEOM_NAMES = (
-    "l_finger_left",
-    "r_finger_left",
-    "l_fingertip_g0_left",
-    "r_fingertip_g0_left",
-    "l_fingerpad_g0_left",
-    "r_fingerpad_g0_left",
-    "l_fingertip_g0_vis_left",
-    "r_fingertip_g0_vis_left",
-)
-RIGHT_GRIPPER_GEOM_NAMES = (
-    "l_finger_right",
-    "r_finger_right",
-    "l_fingertip_g0_right",
-    "r_fingertip_g0_right",
-    "l_fingerpad_g0_right",
-    "r_fingerpad_g0_right",
-    "l_fingertip_g0_vis_right",
-    "r_fingertip_g0_vis_right",
-)
-TABLE_GEOM_NAME = "table"
-
-
-def _set_free_joint_pose(joint, position, quat):
-    joint.qpos[:3] = np.asarray(position, dtype=np.float64)
-    joint.qpos[3:7] = np.asarray(quat, dtype=np.float64)
-
-
-def set_socket_peg_task_state(mj_data, task_state):
-    if not isinstance(task_state, dict) or tuple(task_state.keys()) != REQUIRED_TASK_STATE_KEYS:
-        raise ValueError(
-            "task_state must be an ordered dict-like mapping with keys "
-            "socket_pos, socket_quat, peg_pos, peg_quat"
-        )
-
-    _set_free_joint_pose(
-        mj_data.joint(SOCKET_JOINT_NAME),
-        task_state["socket_pos"],
-        task_state["socket_quat"],
-    )
-    _set_free_joint_pose(
-        mj_data.joint(PEG_JOINT_NAME),
-        task_state["peg_pos"],
-        task_state["peg_quat"],
-    )
-
-
-def get_socket_peg_env_state(mj_data):
-    socket_qpos = cp.deepcopy(np.asarray(mj_data.joint(SOCKET_JOINT_NAME).qpos[:7], dtype=np.float64))
-    peg_qpos = cp.deepcopy(np.asarray(mj_data.joint(PEG_JOINT_NAME).qpos[:7], dtype=np.float64))
-    return np.concatenate([socket_qpos, peg_qpos], dtype=np.float64)
-
-
-def _normalize_contact_pairs(contact_pairs):
-    return {frozenset(pair) for pair in contact_pairs}
-
-
-def _has_any_object_contact(contact_set, object_geom_names, other_geom_names):
-    return any(
-        frozenset((object_geom_name, other_geom_name)) in contact_set
-        for object_geom_name in object_geom_names
-        for other_geom_name in other_geom_names
-    )
-
-
-def _object_is_airborne(contact_set, object_geom_names):
-    return not _has_any_object_contact(contact_set, object_geom_names, (TABLE_GEOM_NAME,))
-
-
-def peg_inserted_into_socket(contact_pairs):
-    contact_set = _normalize_contact_pairs(contact_pairs)
-    return frozenset((PEG_GEOM_NAMES[0], SOCKET_SUCCESS_GEOM_NAMES[0])) in contact_set
-
-
-def compute_air_insert_reward(contact_pairs, env_state=None):
-    del env_state  # kept for API compatibility with rollout/eval code paths
-    contact_set = _normalize_contact_pairs(contact_pairs)
-    reward = 0
-
-    if _has_any_object_contact(contact_set, SOCKET_GEOM_NAMES, LEFT_GRIPPER_GEOM_NAMES):
-        reward += 1
-    if _has_any_object_contact(contact_set, PEG_GEOM_NAMES, RIGHT_GRIPPER_GEOM_NAMES):
-        reward += 1
-
-    socket_airborne = _object_is_airborne(contact_set, SOCKET_BODY_GEOM_NAMES)
-    peg_airborne = _object_is_airborne(contact_set, PEG_GEOM_NAMES)
-    if socket_airborne:
-        reward += 1
-    if peg_airborne:
-        reward += 1
-
-    if socket_airborne and peg_airborne and peg_inserted_into_socket(contact_pairs):
-        reward += 1
-
-    return reward
-
-
-class DualDianaMed_Air_Insert(DualDianaMed_Pos_Ctrl):
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self.max_reward = 5
-
-    def reset(self, task_state):
-        set_socket_peg_task_state(self.mj_data, task_state)
-        DualDianaMed.reset(self)
-        self.top = None
-        self.r_vis = None
-        self.l_vis = None
-        self.front = None
-        if not self.is_render:
-            self._update_camera_images_sync()
-            return
-        self.cam_flage = True
-        while self.cam_flage:
-            if (
-                type(self.top) == type(None)
-                or type(self.r_vis) == type(None)
-                or type(self.l_vis) == type(None)
-                or type(self.front) == type(None)
-            ):
-                time.sleep(0.001)
-            else:
-                self.cam_flage = False
-
-    def step(self, action=np.zeros(16)):
-        super().step(action)
-        self.rew = self._get_reward()
-        self.obs = self._get_obs()
-
-    def get_env_state(self):
-        return get_socket_peg_env_state(self.mj_data)
-
-    def _get_reward(self):
-        contact_pairs = []
-        for collision_num in range(self.mj_data.ncon):
-            geom1 = self.mj_data.contact[collision_num].geom1
-            geom2 = self.mj_data.contact[collision_num].geom2
-            contact_pairs.append(
-                (self.getID2Name("geom", geom1), self.getID2Name("geom", geom2))
-            )
-        return compute_air_insert_reward(contact_pairs, self.get_env_state())

roboimi/envs/double_base.py

@@ -52,6 +52,7 @@ class DualDianaMed(MujocoEnv):
         self.r_vis = None
         self.l_vis = None
         self.top = None
+        self.angle = None
         self.front = None
         self.obs = None
         
@@ -91,6 +92,7 @@ class DualDianaMed(MujocoEnv):
 
     def step(self,action):
         self.compute_qpos = action  #for observation !
+        self.obs = self._get_obs()
         if self.interpolator_left is not None and self.interpolator_right is not None:
             self.interpolator_left.updateInput(action[:7], control_cycle=self.base_time)
             self.interpolator_right.updateInput(action[7:-2], control_cycle=self.base_time)
@@ -103,7 +105,6 @@ class DualDianaMed(MujocoEnv):
 
             super().step(action)
         self.base_time = time.time() - ctrl_cur_time
-        self.obs = self._get_obs()
             
 
     def preStep(self, action):
@@ -167,9 +168,10 @@ class DualDianaMed(MujocoEnv):
         obs['qpos'] = self.get_obs_qpos
         obs['action'] = self.compute_qpos
         obs['images'] = dict()
+        obs['images']['top']   = self.top
+        obs['images']['angle'] = self.angle
         obs['images']['r_vis'] = self.r_vis
         obs['images']['l_vis'] = self.l_vis
-        obs['images']['top'] = self.top
         obs['images']['front'] = self.front
         return obs
 
@@ -178,9 +180,10 @@ class DualDianaMed(MujocoEnv):
             self._update_camera_images_sync()
         obs = collections.OrderedDict()
         obs['images'] = dict()
+        obs['images']['top']   = self.top
+        obs['images']['angle'] = self.angle
         obs['images']['r_vis'] = self.r_vis
         obs['images']['l_vis'] = self.l_vis
-        obs['images']['top'] = self.top
         obs['images']['front'] = self.front
         return obs
     
@@ -199,12 +202,14 @@ class DualDianaMed(MujocoEnv):
 
     @property
     def cam_view(self):
-        if self.cam == 'r_vis':
+        if self.cam == 'top':
+            return self.top
+        elif self.cam == 'angle':
+            return self.angle
+        elif self.cam == 'r_vis':
             return self.r_vis
         elif self.cam == 'l_vis':
             return self.l_vis
-        elif self.cam == 'top':
-            return self.top
         elif self.cam == 'front':
             return self.front
         else:
@@ -225,6 +230,8 @@ class DualDianaMed(MujocoEnv):
         self.l_vis = img_renderer.render()[:, :, ::-1]
         img_renderer.update_scene(self.mj_data, camera="top")
         self.top = img_renderer.render()[:, :, ::-1]
+        img_renderer.update_scene(self.mj_data, camera="angle")
+        self.angle = img_renderer.render()[:, :, ::-1]
         img_renderer.update_scene(self.mj_data, camera="front")
         self.front = img_renderer.render()[:, :, ::-1]
 

roboimi/envs/double_pos_ctrl_env.py

@@ -73,8 +73,8 @@ class DualDianaMed_Pos_Ctrl(DualDianaMed):
         self.mj_data.joint('red_box_joint').qpos[6] = 0.0
         super().reset()
         self.top = None
+        self.angle = None
         self.r_vis = None
-        self.l_vis = None
         self.front = None
         if not self.is_render:
             self._update_camera_images_sync()
@@ -83,8 +83,8 @@ class DualDianaMed_Pos_Ctrl(DualDianaMed):
         t=0
         while self.cam_flage:
             if(type(self.top)==type(None) 
+               or type(self.angle)==type(None) 
                or type(self.r_vis)==type(None)
-               or type(self.l_vis)==type(None)
                or type(self.front)==type(None)):
                 time.sleep(0.001)
                 t+=1
@@ -137,18 +137,6 @@ class DualDianaMed_Pos_Ctrl(DualDianaMed):
         
 
 def make_sim_env(task_name, headless=False):
-    if task_name == 'sim_air_insert_socket_peg':
-        from roboimi.assets.robots.diana_med import BiDianaMedSocketPeg
-        from roboimi.envs.double_air_insert_env import DualDianaMed_Air_Insert
-
-        env = DualDianaMed_Air_Insert(
-            robot=BiDianaMedSocketPeg(),
-            is_render=not headless,
-            control_freq=30,
-            is_interpolate=True,
-            cam_view='front'
-        )
-        return env
     if 'sim_transfer' in task_name:
         from roboimi.assets.robots.diana_med import BiDianaMed
         env = DualDianaMed_Pos_Ctrl(
@@ -156,7 +144,7 @@ def make_sim_env(task_name, headless=False):
             is_render=not headless,
             control_freq=30,
             is_interpolate=True,
-            cam_view='top'
+            cam_view='angle'
         )
         return env
     else:
@@ -182,3 +170,4 @@ if __name__ == "__main__":
         env.step(action)
         if env.is_render:
             env.render()
+        

roboimi/scripts/refresh_experiment_suite_status.py

@@ -0,0 +1,267 @@
+#!/usr/bin/env python3
+
+from __future__ import annotations
+
+import argparse
+import datetime as dt
+import json
+import pathlib
+import re
+import shlex
+import subprocess
+from collections import defaultdict
+from typing import Any
+
+
+STEP_PAT = re.compile(r"步骤\s+(\d+)/(\d+)")
+BAR_PAT = re.compile(r"\|\s*(\d+)/(\d+)")
+
+
+def normalize_chunks(text: str):
+    for part in re.split(r"[\r\n]+", text):
+        part = part.strip()
+        if part:
+            yield part
+
+
+def parse_latest_line(text: str) -> tuple[str, int | None]:
+    latest_line = ""
+    latest_step = None
+    for line in normalize_chunks(text):
+        if "步骤" not in line and "训练中:" not in line:
+            continue
+        latest_line = line
+        match = STEP_PAT.search(line) or BAR_PAT.search(line)
+        if match:
+            latest_step = int(match.group(1))
+    return latest_line, latest_step
+
+
+def now_iso() -> str:
+    return dt.datetime.now(
+        dt.timezone(dt.timedelta(hours=8)),
+    ).isoformat(timespec="seconds")
+
+
+def run_cmd(cmd: list[str], check: bool = True) -> subprocess.CompletedProcess[str]:
+    return subprocess.run(cmd, capture_output=True, text=True, check=check)
+
+
+def probe_local(run: dict[str, Any]) -> dict[str, Any]:
+    pid = str(run["pid"])
+    ps = run_cmd(["ps", "-p", pid, "-o", "pid=,stat=,etime=,args="], check=False)
+    log_path = pathlib.Path(run["log_path"])
+    latest_line = ""
+    latest_step = None
+    if log_path.exists():
+        latest_line, latest_step = parse_latest_line(log_path.read_text(errors="replace"))
+    return {
+        "alive": bool(ps.stdout.strip()),
+        "ps": ps.stdout.strip(),
+        "log_exists": log_path.exists(),
+        "latest_line": latest_line,
+        "latest_step": latest_step,
+    }
+
+
+def remote_probe(host: str, remote_user: str, runs: list[dict[str, Any]]) -> dict[str, dict[str, Any]]:
+    payload = [
+        {
+            "run_id": run["run_id"],
+            "pid": str(run["pid"]),
+            "log_path": run["log_path"],
+        }
+        for run in runs
+    ]
+    remote_py = r"""
+import json
+import pathlib
+import re
+import subprocess
+import sys
+
+payload = json.loads(sys.argv[1])
+step_pat = re.compile(r"步骤\s+(\d+)/(\d+)")
+bar_pat = re.compile(r"\|\s*(\d+)/(\d+)")
+
+def normalize_chunks(text):
+    for part in re.split(r"[\r\n]+", text):
+        part = part.strip()
+        if part:
+            yield part
+
+def parse_latest_line(text):
+    latest_line = ""
+    latest_step = None
+    for line in normalize_chunks(text):
+        if "步骤" not in line and "训练中:" not in line:
+            continue
+        latest_line = line
+        match = step_pat.search(line) or bar_pat.search(line)
+        if match:
+            latest_step = int(match.group(1))
+    return latest_line, latest_step
+
+out = {}
+for item in payload:
+    try:
+        ps = subprocess.run(
+            ["ps", "-p", item["pid"], "-o", "pid=,stat=,etime=,args="],
+            capture_output=True,
+            text=True,
+            check=False,
+        )
+        log_path = pathlib.Path(item["log_path"])
+        latest_line = ""
+        latest_step = None
+        if log_path.exists():
+            latest_line, latest_step = parse_latest_line(log_path.read_text(errors="replace"))
+        out[item["run_id"]] = {
+            "alive": bool(ps.stdout.strip()),
+            "ps": ps.stdout.strip(),
+            "log_exists": log_path.exists(),
+            "latest_line": latest_line,
+            "latest_step": latest_step,
+        }
+    except Exception as exc:
+        out[item["run_id"]] = {
+            "alive": False,
+            "ps": "",
+            "log_exists": False,
+            "latest_line": "",
+            "latest_step": None,
+            "error": str(exc),
+        }
+print(json.dumps(out, ensure_ascii=False))
+"""
+    remote_target = host if "@" in host else f"{remote_user}@{host}"
+    remote_cmd = (
+        f"python3 -c {shlex.quote(remote_py)} "
+        f"{shlex.quote(json.dumps(payload, ensure_ascii=False))}"
+    )
+    try:
+        res = run_cmd(
+            [
+                "ssh",
+                "-F",
+                "/dev/null",
+                "-o",
+                "BatchMode=yes",
+                "-o",
+                "StrictHostKeyChecking=accept-new",
+                remote_target,
+                remote_cmd,
+            ]
+        )
+        return json.loads(res.stdout)
+    except subprocess.CalledProcessError as exc:
+        error = (exc.stderr or exc.stdout or str(exc)).strip()
+        return {
+            run["run_id"]: {
+                "alive": False,
+                "ps": "",
+                "log_exists": False,
+                "latest_line": "",
+                "latest_step": None,
+                "error": f"ssh_failed: {error}",
+            }
+            for run in runs
+        }
+
+
+def append_notes(notes_path: pathlib.Path, snapshot_at: str, runs: list[dict[str, Any]]) -> None:
+    lines = [f"\n## Status snapshot {snapshot_at}"]
+    for run in runs:
+        lines.append(
+            (
+                f"- {run['run_id']}: host={run['host']} gpu={run['gpu']} "
+                f"alive={run.get('alive', False)} step={run.get('latest_step')} "
+                f"pid={run['pid']}"
+            )
+        )
+        if run.get("latest_line"):
+            lines.append(f"  - latest_line: `{run['latest_line']}`")
+        if run.get("error"):
+            lines.append(f"  - error: `{run['error']}`")
+    with notes_path.open("a", encoding="utf-8") as f:
+        f.write("\n".join(lines) + "\n")
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("suite_dir", type=pathlib.Path)
+    parser.add_argument("--remote-user", default="droid")
+    parser.add_argument("--append-notes", action="store_true")
+    args = parser.parse_args()
+
+    suite_dir = args.suite_dir.resolve()
+    status_path = suite_dir / "status.json"
+    notes_path = suite_dir / "notes.md"
+    monitor_dir = suite_dir / "monitor_logs"
+    monitor_dir.mkdir(parents=True, exist_ok=True)
+
+    status = json.loads(status_path.read_text(encoding="utf-8"))
+    runs: list[dict[str, Any]] = status["runs"]
+    snapshot_at = now_iso()
+
+    by_host: dict[str, list[dict[str, Any]]] = defaultdict(list)
+    for run in runs:
+        by_host[run["host"]].append(run)
+
+    results: dict[str, dict[str, Any]] = {}
+    for host, host_runs in by_host.items():
+        if host == "local":
+            for run in host_runs:
+                results[run["run_id"]] = probe_local(run)
+        else:
+            results.update(remote_probe(host, args.remote_user, host_runs))
+
+    alive_count = 0
+    for run in runs:
+        result = results[run["run_id"]]
+        run["alive"] = result["alive"]
+        run["ps"] = result["ps"]
+        run["log_exists"] = result["log_exists"]
+        run["latest_line"] = result["latest_line"]
+        run["latest_step"] = result["latest_step"]
+        run["last_verified_at"] = snapshot_at
+        if "error" in result:
+            run["error"] = result["error"]
+        else:
+            run.pop("error", None)
+        run["status"] = "running" if result["alive"] else "stopped"
+        alive_count += int(result["alive"])
+
+    status["last_verified_at"] = snapshot_at
+    status["alive_count"] = alive_count
+    status["total_runs"] = len(runs)
+
+    status_path.write_text(json.dumps(status, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
+
+    snapshot_payload = {
+        "suite_name": status.get("suite_name"),
+        "snapshot_at": snapshot_at,
+        "alive_count": alive_count,
+        "total_runs": len(runs),
+        "runs": {run["run_id"]: results[run["run_id"]] for run in runs},
+    }
+    timestamp_slug = snapshot_at.replace(":", "").replace("+", "_").replace("-", "")
+    snapshot_path = monitor_dir / f"status-{timestamp_slug}.json"
+    snapshot_path.write_text(
+        json.dumps(snapshot_payload, ensure_ascii=False, indent=2) + "\n",
+        encoding="utf-8",
+    )
+
+    if args.append_notes:
+        append_notes(notes_path, snapshot_at, runs)
+
+    print(json.dumps(snapshot_payload, ensure_ascii=False, indent=2))
+    print(f"\nstatus_json={status_path}")
+    print(f"snapshot_json={snapshot_path}")
+    if args.append_notes:
+        print(f"notes_md={notes_path}")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

roboimi/utils/act_ex_utils.py

@@ -1,6 +1,5 @@
 import numpy as np
 
-
 def sample_insertion_pose():
         # Peg
         x_range = [0.1, 0.2]
@@ -37,22 +36,3 @@ def sample_transfer_pose():
 
 
         return box_position
-
-
-def sample_air_insert_socket_peg_state():
-        socket_position = np.random.uniform(
-            low=np.array([-0.20, 0.80, 0.472], dtype=np.float32),
-            high=np.array([-0.10, 1.00, 0.472], dtype=np.float32),
-        )
-        peg_position = np.random.uniform(
-            low=np.array([0.10, 0.80, 0.46], dtype=np.float32),
-            high=np.array([0.20, 1.00, 0.46], dtype=np.float32),
-        )
-        socket_quat = np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32)
-        peg_quat = np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32)
-        return {
-            "socket_pos": socket_position.astype(np.float32, copy=False),
-            "socket_quat": socket_quat,
-            "peg_pos": peg_position.astype(np.float32, copy=False),
-            "peg_quat": peg_quat,
-        }

roboimi/utils/constants.py

@@ -20,14 +20,7 @@ SIM_TASK_CONFIGS = {
         'dataset_dir': DATASET_DIR + '/sim_transfer',
         'num_episodes': 20,
         'episode_len': 700,
-        'camera_names': ['r_vis', 'top', 'front'],
+        'camera_names': ['top','r_vis','front'],
-        'xml_dir': HOME_PATH + '/assets'
-    },
-    'sim_air_insert_socket_peg': {
-        'dataset_dir': DATASET_DIR + '/sim_air_insert_socket_peg',
-        'num_episodes': 20,
-        'episode_len': 750,
-        'camera_names': ['l_vis', 'r_vis', 'front'],
         'xml_dir': HOME_PATH + '/assets'
     },
 
@@ -59,3 +52,13 @@ PUPPET_GRIPPER_JOINT_NORMALIZE_FN = lambda x: (x - PUPPET_GRIPPER_JOINT_CLOSE) /
 MASTER_GRIPPER_JOINT_UNNORMALIZE_FN = lambda x: x * (MASTER_GRIPPER_JOINT_OPEN - MASTER_GRIPPER_JOINT_CLOSE) + MASTER_GRIPPER_JOINT_CLOSE
 PUPPET_GRIPPER_JOINT_UNNORMALIZE_FN = lambda x: x * (PUPPET_GRIPPER_JOINT_OPEN - PUPPET_GRIPPER_JOINT_CLOSE) + PUPPET_GRIPPER_JOINT_CLOSE
 MASTER2PUPPET_JOINT_FN = lambda x: PUPPET_GRIPPER_JOINT_UNNORMALIZE_FN(MASTER_GRIPPER_JOINT_NORMALIZE_FN(x))
+
+MASTER_GRIPPER_VELOCITY_NORMALIZE_FN = lambda x: x / (MASTER_GRIPPER_POSITION_OPEN - MASTER_GRIPPER_POSITION_CLOSE)
+PUPPET_GRIPPER_VELOCITY_NORMALIZE_FN = lambda x: x / (PUPPET_GRIPPER_POSITION_OPEN - PUPPET_GRIPPER_POSITION_CLOSE)
+
+MASTER_POS2JOINT = lambda x: MASTER_GRIPPER_POSITION_NORMALIZE_FN(x) * (MASTER_GRIPPER_JOINT_OPEN - MASTER_GRIPPER_JOINT_CLOSE) + MASTER_GRIPPER_JOINT_CLOSE
+MASTER_JOINT2POS = lambda x: MASTER_GRIPPER_POSITION_UNNORMALIZE_FN((x - MASTER_GRIPPER_JOINT_CLOSE) / (MASTER_GRIPPER_JOINT_OPEN - MASTER_GRIPPER_JOINT_CLOSE))
+PUPPET_POS2JOINT = lambda x: PUPPET_GRIPPER_POSITION_NORMALIZE_FN(x) * (PUPPET_GRIPPER_JOINT_OPEN - PUPPET_GRIPPER_JOINT_CLOSE) + PUPPET_GRIPPER_JOINT_CLOSE
+PUPPET_JOINT2POS = lambda x: PUPPET_GRIPPER_POSITION_UNNORMALIZE_FN((x - PUPPET_GRIPPER_JOINT_CLOSE) / (PUPPET_GRIPPER_JOINT_OPEN - PUPPET_GRIPPER_JOINT_CLOSE))
+
+MASTER_GRIPPER_JOINT_MID = (MASTER_GRIPPER_JOINT_OPEN + MASTER_GRIPPER_JOINT_CLOSE)/2

roboimi/vla/agent.py

@@ -28,6 +28,7 @@ class VLAAgent(nn.Module):
         num_action_steps=8,   # 每次推理实际执行多少步动作
         head_type='unet',     # Policy head类型: 'unet' 或 'transformer'
         cond_projector=None,  # 可选：将视觉+状态条件投影到head期望维度
+        extra_condition_tokens: int = 0,  # 可选：额外条件token数量（例如未来预测embedding）
     ):
         super().__init__()
         # 保存参数
@@ -39,6 +40,9 @@ class VLAAgent(nn.Module):
         self.num_action_steps = num_action_steps
         self.inference_steps = inference_steps
         self.head_type = head_type  # 'unet' 或 'transformer'
+        self.extra_condition_tokens = int(extra_condition_tokens)
+        if self.extra_condition_tokens < 0:
+            raise ValueError(f"extra_condition_tokens must be >= 0, got {self.extra_condition_tokens}")
         agent_camera_names = tuple(camera_names) if camera_names is not None else None
         backbone_camera_names = getattr(vision_backbone, 'camera_names', None)
         backbone_camera_names = tuple(backbone_camera_names) if backbone_camera_names is not None else None
@@ -71,11 +75,14 @@ class VLAAgent(nn.Module):
             stats=dataset_stats,
             normalization_type=normalization_type
         )
+        self.dataset_stats = dataset_stats
 
         self.vision_encoder = vision_backbone
+        self.state_encoder = state_encoder
         if self.camera_names is not None:
             self.vision_encoder.camera_names = self.camera_names
         self.condition_tokens_per_step = int(getattr(self.vision_encoder, 'tokens_per_step', 1))
+        self.state_feature_dim = int(getattr(self.state_encoder, 'output_dim', obs_dim))
         joint_vision_dim = getattr(self.vision_encoder, 'joint_output_dim', None)
         if joint_vision_dim is not None:
             per_token_vision_dim = int(joint_vision_dim)
@@ -87,8 +94,11 @@ class VLAAgent(nn.Module):
             else:
                 per_token_vision_dim = int(single_cam_feat_dim) * int(num_cams)
 
-        self.condition_sequence_length = self.obs_horizon * self.condition_tokens_per_step
+        self.history_condition_sequence_length = self.obs_horizon * self.condition_tokens_per_step
-        self.raw_per_step_cond_dim = per_token_vision_dim + obs_dim
+        self.condition_sequence_length = (
+            self.history_condition_sequence_length + self.extra_condition_tokens
+        )
+        self.raw_per_step_cond_dim = per_token_vision_dim + self.state_feature_dim
         if cond_projector is None:
             self.cond_projector = None
             self.per_step_cond_dim = self.raw_per_step_cond_dim
@@ -139,7 +149,6 @@ class VLAAgent(nn.Module):
                 global_cond_dim=self.global_cond_dim
             )
 
-        self.state_encoder = state_encoder
         self.action_encoder = action_encoder
 
         # 初始化队列（用于在线推理）
@@ -220,7 +229,7 @@ class VLAAgent(nn.Module):
                     f"条件维度不匹配: got {cond.shape[-1]}, expected {self.per_step_cond_dim}"
                 )
             cond = cond.reshape(batch_size, obs_steps * token_count, self.per_step_cond_dim)
-            expected_length = self.condition_sequence_length
+            expected_length = self.history_condition_sequence_length
             if cond.shape[1] != expected_length:
                 raise RuntimeError(
                     f"条件序列长度不匹配: got {cond.shape[1]}, expected {expected_length}"

roboimi/vla/agent_imf.py

@@ -1,9 +1,12 @@
 from __future__ import annotations
 
 from contextlib import nullcontext
-from typing import Dict, Optional
+from collections import deque
+from pathlib import Path
+from typing import Any, Dict, Mapping, Optional, Sequence
 
 import torch
+import torch.nn as nn
 import torch.nn.functional as F
 
 from roboimi.vla.agent import VLAAgent
@@ -15,14 +18,87 @@ except ImportError:  # pragma: no cover
 
 
 class IMFVLAAgent(VLAAgent):
-    def __init__(self, *args, inference_steps: int = 1, **kwargs):
+    def __init__(
+        self,
+        *args,
+        inference_steps: int = 1,
+        lewm_history_horizon: Optional[int] = None,
+        lewm_query_offsets: Optional[Sequence[int]] = None,
+        lewm_predictor: Optional[nn.Module] = None,
+        lewm_pred_projector: Optional[nn.Module] = None,
+        future_decoder: Optional[nn.Module] = None,
+        future_query_init_std: float = 0.02,
+        lewm_sigreg: Optional[nn.Module] = None,
+        lewm_sigreg_weight: float = 0.09,
+        lewm_loss_weight: float = 0.0,
+        lewm_pretrained_ckpt: Optional[str | Path | Mapping[str, Any]] = None,
+        **kwargs,
+    ):
         if inference_steps != 1:
             raise ValueError(
                 'IMFVLAAgent only supports one-step inference; '
                 f'inference_steps must be 1, got {inference_steps}.'
             )
+        lewm_query_offsets = tuple(int(offset) for offset in (lewm_query_offsets or ()))
+        inferred_extra_condition_tokens = len(lewm_query_offsets) if lewm_query_offsets else 0
+        default_extra_condition_tokens = (
+            0 if future_decoder is not None else inferred_extra_condition_tokens
+        )
+        kwargs.setdefault('extra_condition_tokens', default_extra_condition_tokens)
+        self.__dict__['lewm_history_horizon'] = int(lewm_history_horizon or kwargs.get('obs_horizon', 1))
+        self.__dict__['lewm_query_offsets'] = lewm_query_offsets
+        self.__dict__['lewm_predictor'] = lewm_predictor
+        self.__dict__['lewm_pred_projector'] = lewm_pred_projector or nn.Identity()
+        self.__dict__['future_decoder'] = future_decoder
+        self.__dict__['future_query_tokens'] = None
+        self.__dict__['future_query_init_std'] = float(future_query_init_std)
+        self.__dict__['lewm_sigreg'] = lewm_sigreg
+        self.__dict__['lewm_sigreg_weight'] = float(lewm_sigreg_weight)
+        self.__dict__['lewm_loss_weight'] = float(lewm_loss_weight)
+        self.__dict__['_last_loss_breakdown'] = {
+            'action_loss': 0.0,
+            'lewm_pred_loss': 0.0,
+            'lewm_sigreg_loss': 0.0,
+            'lewm_loss': 0.0,
+            'loss': 0.0,
+        }
         super().__init__(*args, inference_steps=inference_steps, **kwargs)
         self.inference_steps = 1
+        self.lewm_history_horizon = int(lewm_history_horizon or self.obs_horizon)
+        self.lewm_predictor = lewm_predictor
+        self.lewm_pred_projector = lewm_pred_projector or nn.Identity()
+        if future_decoder is not None and not isinstance(future_decoder, nn.Module):
+            self.future_decoder = future_decoder()
+        else:
+            self.future_decoder = future_decoder
+        self.future_query_tokens = None
+        self.future_query_init_std = float(future_query_init_std)
+        self.lewm_sigreg = lewm_sigreg
+        self.lewm_sigreg_weight = float(lewm_sigreg_weight)
+        if self.lewm_predictor is not None and self.future_decoder is not None:
+            raise ValueError('lewm_predictor and future_decoder are mutually exclusive')
+        if self.lewm_predictor is None and self.extra_condition_tokens > 0:
+            raise ValueError(
+                'extra_condition_tokens > 0 requires lewm_predictor to be provided'
+            )
+        if self.lewm_predictor is not None and self.extra_condition_tokens != inferred_extra_condition_tokens:
+            raise ValueError(
+                'extra_condition_tokens must equal len(lewm_query_offsets) when lewm_predictor is enabled'
+            )
+        if self.future_decoder is not None:
+            if inferred_extra_condition_tokens <= 0:
+                raise ValueError('future_decoder requires non-empty lewm_query_offsets')
+            if self.extra_condition_tokens != 0:
+                raise ValueError('future_decoder requires extra_condition_tokens=0')
+            self.future_query_tokens = nn.Parameter(
+                torch.randn(
+                    1,
+                    inferred_extra_condition_tokens,
+                    self.per_step_cond_dim,
+                ) * self.future_query_init_std
+            )
+        if lewm_pretrained_ckpt is not None:
+            self.load_lewm_pretrained_components(lewm_pretrained_ckpt)
 
     @staticmethod
     def _broadcast_batch_time(value: torch.Tensor, reference: torch.Tensor) -> torch.Tensor:
@@ -119,14 +195,251 @@ class IMFVLAAgent(VLAAgent):
         delta = self._broadcast_batch_time(t - r, z_t)
         return z_t - delta * u
 
+    def _normalize_qpos_for_lewm(self, qpos: torch.Tensor) -> torch.Tensor:
+        if not self.normalization.enabled:
+            return qpos
+
+        qpos_mean = getattr(self.normalization, 'qpos_mean', None)
+        qpos_std = getattr(self.normalization, 'qpos_std', None)
+        if qpos_mean is not None and qpos_std is not None:
+            return (qpos - qpos_mean) / qpos_std
+        if isinstance(self.dataset_stats, dict):
+            mean = self.dataset_stats.get('qpos_mean', None)
+            std = self.dataset_stats.get('qpos_std', None)
+            if mean is not None and std is not None:
+                mean = torch.as_tensor(mean, dtype=qpos.dtype, device=qpos.device)
+                std = torch.as_tensor(std, dtype=qpos.dtype, device=qpos.device)
+                return (qpos - mean) / std
+        return self.normalization.normalize_qpos(qpos)
+
+    def _project_lewm_future_tokens(self, predicted_tokens: torch.Tensor) -> torch.Tensor:
+        if predicted_tokens.ndim != 3:
+            raise ValueError(
+                f"expected predicted future tokens to be 3D, got rank {predicted_tokens.ndim}"
+            )
+        batch_size, token_count, token_dim = predicted_tokens.shape
+        flattened = predicted_tokens.reshape(batch_size * token_count, token_dim)
+        projected = self.lewm_pred_projector(flattened)
+        if projected.ndim != 2:
+            raise ValueError(
+                f"expected lewm_pred_projector to return rank-2 tensors, got rank {projected.ndim}"
+            )
+        return projected.reshape(batch_size, token_count, projected.shape[-1])
+
+    @staticmethod
+    def _load_checkpoint_payload(
+        checkpoint_or_path: str | Path | Mapping[str, Any],
+    ) -> Mapping[str, torch.Tensor]:
+        if isinstance(checkpoint_or_path, (str, Path)):
+            payload = torch.load(Path(checkpoint_or_path), map_location='cpu', weights_only=False)
+        else:
+            payload = checkpoint_or_path
+        state_dict = payload.get('state_dict', payload)
+        if not isinstance(state_dict, Mapping):
+            raise TypeError('checkpoint payload must contain a mapping state_dict')
+        return state_dict
+
+    @staticmethod
+    def _extract_prefixed_state_dict(
+        state_dict: Mapping[str, torch.Tensor],
+        prefix: str,
+    ) -> Dict[str, torch.Tensor]:
+        extracted = {
+            key[len(prefix):]: value
+            for key, value in state_dict.items()
+            if key.startswith(prefix)
+        }
+        if not extracted:
+            raise KeyError(f"checkpoint missing parameters with prefix {prefix!r}")
+        return extracted
+
+    @staticmethod
+    def _adapt_and_load_state_dict(
+        module: nn.Module,
+        incoming_state_dict: Mapping[str, torch.Tensor],
+        *,
+        query_key: str = 'query_tokens',
+        pos_key: str = 'pos_embedding',
+    ) -> Dict[str, Sequence[str]]:
+        current_state_dict = module.state_dict()
+        adapted_state_dict = dict(current_state_dict)
+        loaded_keys = []
+        mismatched_keys = []
+        missing_keys = []
+        for key, current_tensor in current_state_dict.items():
+            if key not in incoming_state_dict:
+                continue
+            source_tensor = incoming_state_dict[key]
+            if source_tensor.shape == current_tensor.shape:
+                adapted_state_dict[key] = source_tensor
+                loaded_keys.append(key)
+                continue
+
+            if key in {query_key, pos_key} and source_tensor.ndim == current_tensor.ndim:
+                patched = current_tensor.clone()
+                overlap_slices = tuple(
+                    slice(0, min(src_dim, cur_dim))
+                    for src_dim, cur_dim in zip(source_tensor.shape, current_tensor.shape)
+                )
+                patched[overlap_slices] = source_tensor[overlap_slices]
+                if key == query_key:
+                    copy_count = min(source_tensor.shape[1], current_tensor.shape[1])
+                    if copy_count < current_tensor.shape[1] and copy_count > 0:
+                        tail = source_tensor[:, copy_count - 1:copy_count, ...]
+                        feature_dim = min(tail.shape[-1], patched.shape[-1])
+                        patched[:, copy_count:, :feature_dim] = tail[:, :, :feature_dim]
+                else:
+                    copy_count = min(source_tensor.shape[1], current_tensor.shape[1])
+                    if copy_count < current_tensor.shape[1] and copy_count > 0:
+                        tail = source_tensor[:, copy_count - 1:copy_count, ...]
+                        feature_dim = min(tail.shape[-1], patched.shape[-1])
+                        patched[:, copy_count:, :feature_dim] = tail[:, :, :feature_dim]
+                adapted_state_dict[key] = patched
+                loaded_keys.append(key)
+                continue
+            mismatched_keys.append(key)
+
+        for key in incoming_state_dict.keys():
+            if key not in current_state_dict:
+                missing_keys.append(key)
+        module.load_state_dict(adapted_state_dict, strict=True)
+        return {
+            'loaded_keys': tuple(sorted(loaded_keys)),
+            'mismatched_keys': tuple(sorted(set(mismatched_keys))),
+            'missing_keys': tuple(sorted(set(missing_keys))),
+        }
+
+    @staticmethod
+    def _load_state_dict_ignoring_shape_mismatches(
+        module: nn.Module,
+        incoming_state_dict: Mapping[str, torch.Tensor],
+    ) -> Dict[str, Sequence[str]]:
+        current_state_dict = module.state_dict()
+        merged_state_dict = dict(current_state_dict)
+        loaded_keys = []
+        mismatched_keys = []
+        missing_keys = []
+
+        for key, value in incoming_state_dict.items():
+            if key not in current_state_dict:
+                missing_keys.append(key)
+                continue
+            if current_state_dict[key].shape != value.shape:
+                mismatched_keys.append(key)
+                continue
+            merged_state_dict[key] = value
+            loaded_keys.append(key)
+
+        module.load_state_dict(merged_state_dict, strict=True)
+        return {
+            'loaded_keys': tuple(sorted(loaded_keys)),
+            'mismatched_keys': tuple(sorted(mismatched_keys)),
+            'missing_keys': tuple(sorted(missing_keys)),
+        }
+
+    def load_lewm_pretrained_components(
+        self,
+        checkpoint_or_path: str | Path | Mapping[str, Any],
+    ) -> None:
+        state_dict = self._load_checkpoint_payload(checkpoint_or_path)
+
+        if hasattr(self.vision_encoder, 'load_lewm_checkpoint'):
+            try:
+                self.vision_encoder.load_lewm_checkpoint({'state_dict': state_dict})
+            except RuntimeError:
+                vision_state_dict = self._extract_prefixed_state_dict(state_dict, 'model.encoder.')
+                self._load_state_dict_ignoring_shape_mismatches(self.vision_encoder, vision_state_dict)
+        else:
+            vision_state_dict = self._extract_prefixed_state_dict(state_dict, 'model.encoder.')
+            self._load_state_dict_ignoring_shape_mismatches(self.vision_encoder, vision_state_dict)
+
+        state_encoder_state_dict = self._extract_prefixed_state_dict(state_dict, 'model.state_encoder.')
+        self._load_state_dict_ignoring_shape_mismatches(self.state_encoder, state_encoder_state_dict)
+
+        projector_state_dict = self._extract_prefixed_state_dict(state_dict, 'model.projector.proj.')
+        mapped_projector_state_dict = {
+            f'linear.{key}': value
+            for key, value in projector_state_dict.items()
+        }
+        self._load_state_dict_ignoring_shape_mismatches(self.cond_projector, mapped_projector_state_dict)
+
+        if self.lewm_predictor is not None:
+            predictor_state_dict = self._extract_prefixed_state_dict(state_dict, 'model.predictor.')
+            self._adapt_and_load_state_dict(self.lewm_predictor, predictor_state_dict)
+
+        if self.lewm_pred_projector is not None:
+            pred_projector_state_dict = self._extract_prefixed_state_dict(state_dict, 'model.pred_proj.')
+            self._load_state_dict_ignoring_shape_mismatches(
+                self.lewm_pred_projector,
+                pred_projector_state_dict,
+            )
+
+    def _predict_future_tokens_with_decoder(self, history_cond: torch.Tensor) -> torch.Tensor:
+        if self.future_decoder is None or self.future_query_tokens is None:
+            raise RuntimeError('future_decoder path requested but not initialized')
+        batch_size = history_cond.shape[0]
+        query_tokens = self.future_query_tokens.expand(batch_size, -1, -1)
+        r = torch.zeros(batch_size, device=history_cond.device, dtype=history_cond.dtype)
+        t = torch.ones(batch_size, device=history_cond.device, dtype=history_cond.dtype)
+        return self.future_decoder(query_tokens, r, t, cond=history_cond)
+
+    def _build_full_condition(
+        self,
+        images,
+        proprioception,
+        *,
+        lewm_images=None,
+        lewm_proprioception=None,
+    ):
+        normalized_proprioception = self.normalization.normalize_qpos(proprioception)
+        history_cond = self._build_cond(images, normalized_proprioception)
+        predicted_future_tokens = None
+        lewm_history_cond = None
+
+        if self.lewm_predictor is None and self.future_decoder is None:
+            return history_cond, predicted_future_tokens, lewm_history_cond
+
+        lewm_images = lewm_images if lewm_images is not None else images
+        lewm_proprioception = (
+            lewm_proprioception if lewm_proprioception is not None else proprioception
+        )
+        lewm_history_cond = self._build_cond(
+            lewm_images,
+            self._normalize_qpos_for_lewm(lewm_proprioception),
+        )
+        cond = history_cond
+        if self.lewm_predictor is not None:
+            predicted_future_tokens = self.lewm_predictor(lewm_history_cond)
+            predicted_future_tokens = self._project_lewm_future_tokens(predicted_future_tokens)
+            cond = torch.cat([history_cond, predicted_future_tokens], dim=1)
+            if cond.shape[1] != self.condition_sequence_length:
+                raise RuntimeError(
+                    f"完整条件序列长度不匹配: got {cond.shape[1]}, expected {self.condition_sequence_length}"
+                )
+            if cond.shape[-1] != self.per_step_cond_dim:
+                raise RuntimeError(
+                    f"完整条件维度不匹配: got {cond.shape[-1]}, expected {self.per_step_cond_dim}"
+                )
+        elif self.future_decoder is not None:
+            predicted_future_tokens = self._predict_future_tokens_with_decoder(lewm_history_cond)
+        return cond, predicted_future_tokens, lewm_history_cond
+
+    @staticmethod
+    def _masked_mse_loss(pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+        return F.mse_loss(pred, target)
+
     def compute_loss(self, batch):
         actions, states, images = batch['action'], batch['qpos'], batch['images']
         action_is_pad = batch.get('action_is_pad', None)
         batch_size = actions.shape[0]
 
-        states = self.normalization.normalize_qpos(states)
         actions = self.normalization.normalize_action(actions)
-        cond = self._build_cond(images, states)
+        cond, predicted_future_tokens, lewm_history_cond = self._build_full_condition(
+            images,
+            states,
+            lewm_images=batch.get('lewm_images', None),
+            lewm_proprioception=batch.get('lewm_qpos', None),
+        )
 
         x = actions
         e = torch.randn_like(x)
@@ -146,16 +459,109 @@ class IMFVLAAgent(VLAAgent):
         if action_is_pad is not None:
             mask = (~action_is_pad).unsqueeze(-1).to(loss.dtype)
             valid_count = mask.sum() * loss.shape[-1]
-            loss = (loss * mask).sum() / valid_count.clamp_min(1.0)
+            action_loss = (loss * mask).sum() / valid_count.clamp_min(1.0)
         else:
-            loss = loss.mean()
+            action_loss = loss.mean()
-        return loss
+
+        lewm_pred_loss = torch.zeros((), device=action_loss.device, dtype=action_loss.dtype)
+        lewm_sigreg_loss = torch.zeros((), device=action_loss.device, dtype=action_loss.dtype)
+        if predicted_future_tokens is not None:
+            lewm_future_images = batch.get('lewm_future_images', None)
+            lewm_future_qpos = batch.get('lewm_future_qpos', None)
+            if lewm_future_images is not None and lewm_future_qpos is not None:
+                future_target = self._build_cond(
+                    lewm_future_images,
+                    self._normalize_qpos_for_lewm(lewm_future_qpos),
+                )
+                lewm_pred_loss = self._masked_mse_loss(predicted_future_tokens, future_target)
+            if self.lewm_sigreg is not None and lewm_history_cond is not None:
+                lewm_sigreg_loss = self.lewm_sigreg(lewm_history_cond.transpose(0, 1))
+
+        lewm_loss = lewm_pred_loss + self.lewm_sigreg_weight * lewm_sigreg_loss
+        total_loss = action_loss + self.lewm_loss_weight * lewm_loss
+        self._last_loss_breakdown = {
+            'action_loss': float(action_loss.detach().item()),
+            'lewm_pred_loss': float(lewm_pred_loss.detach().item()),
+            'lewm_sigreg_loss': float(lewm_sigreg_loss.detach().item()),
+            'lewm_loss': float(lewm_loss.detach().item()),
+            'loss': float(total_loss.detach().item()),
+        }
+        return total_loss
+
+    def get_last_loss_breakdown(self) -> Dict[str, float]:
+        return dict(self._last_loss_breakdown)
+
+    def reset(self):
+        super().reset()
+        if self.lewm_predictor is not None:
+            self._queues['lewm_qpos'] = deque(maxlen=self.lewm_history_horizon)
+            self._queues['lewm_images'] = deque(maxlen=self.lewm_history_horizon)
+
+    def _populate_queues(self, observation: Dict[str, torch.Tensor]) -> None:
+        super()._populate_queues(observation)
+        if self.lewm_predictor is None:
+            return
+        if 'qpos' in observation:
+            self._queues['lewm_qpos'].append(observation['qpos'].clone())
+        if 'images' in observation:
+            ordered_images = self._order_images(observation['images'])
+            self._queues['lewm_images'].append({k: v.clone() for k, v in ordered_images.items()})
+
+    def _prepare_observation_batch(self) -> Dict[str, torch.Tensor]:
+        batch = super()._prepare_observation_batch()
+        if self.lewm_predictor is None:
+            return batch
+
+        qpos_list = list(self._queues['lewm_qpos'])
+        images_list = list(self._queues['lewm_images'])
+        if len(qpos_list) == 0 or len(images_list) == 0:
+            raise ValueError("LeWM 观测队列为空，请先调用 _populate_queues 添加观测")
+        while len(qpos_list) < self.lewm_history_horizon:
+            qpos_list.append(qpos_list[-1])
+        while len(images_list) < self.lewm_history_horizon:
+            images_list.append(images_list[-1])
+
+        batch['lewm_qpos'] = torch.stack(qpos_list, dim=0).unsqueeze(0)
+        batch['lewm_images'] = {}
+        camera_names = self.camera_names if self.camera_names is not None else tuple(sorted(images_list[0].keys()))
+        for cam_name in camera_names:
+            batch['lewm_images'][cam_name] = torch.stack(
+                [img[cam_name] for img in images_list],
+                dim=0,
+            ).unsqueeze(0)
+        return batch
 
     @torch.no_grad()
-    def predict_action(self, images, proprioception):
+    def predict_action_chunk(self, batch: Dict[str, torch.Tensor]) -> torch.Tensor:
+        return self.predict_action(
+            batch['images'],
+            batch['qpos'],
+            lewm_images=batch.get('lewm_images', None),
+            lewm_proprioception=batch.get('lewm_qpos', None),
+        )
+
+    @torch.no_grad()
+    def predict_action(
+        self,
+        images,
+        proprioception,
+        *,
+        lewm_images=None,
+        lewm_proprioception=None,
+    ):
         batch_size = proprioception.shape[0]
-        proprioception = self.normalization.normalize_qpos(proprioception)
+        if self.lewm_predictor is not None:
-        cond = self._build_cond(images, proprioception)
+            cond, _predicted_future_tokens, _lewm_history_cond = self._build_full_condition(
+                images,
+                proprioception,
+                lewm_images=lewm_images,
+                lewm_proprioception=lewm_proprioception,
+            )
+        else:
+            cond = self._build_cond(
+                images,
+                self.normalization.normalize_qpos(proprioception),
+            )
         z_t = torch.randn((batch_size, self.pred_horizon, self.action_dim), device=cond.device, dtype=cond.dtype)
         action = self._sample_one_step(z_t, cond=cond)
         return self.normalization.denormalize_action(action)

roboimi/vla/agent_smolvla_conditioned.py

@@ -1,277 +0,0 @@
-from __future__ import annotations
-
-from collections import deque
-from typing import Dict, Optional, Sequence
-
-import torch
-import torch.nn as nn
-
-from roboimi.vla.agent_imf import IMFVLAAgent
-from roboimi.vla.models.normalization import NormalizationModule
-
-
-class SmolVLAIMFAttnResAgent(IMFVLAAgent):
-    """IMF-AttnRes action expert conditioned by SmolVLA-style VLM prefix tokens.
-
-    Unlike ``VLAAgent`` this agent does not concatenate ResNet features and raw
-    state at every observation step.  A ``condition_encoder`` receives images,
-    normalized state, and variable task language, and returns a condition token
-    sequence `(B, S, D)` that is passed directly to the IMF head.
-    """
-
-    def __init__(
-        self,
-        condition_encoder: nn.Module,
-        action_encoder,
-        head,
-        action_dim: int,
-        obs_dim: int,
-        pred_horizon: int = 16,
-        obs_horizon: int = 2,
-        diffusion_steps: int = 100,
-        inference_steps: int = 1,
-        num_cams: int = 3,
-        camera_names: Optional[Sequence[str]] = None,
-        dataset_stats=None,
-        normalization_type: str = 'min_max',
-        num_action_steps: int = 8,
-        head_type: str = 'transformer',
-        condition_dim: int | None = None,
-        condition_sequence_length: int | None = None,
-        task_description: str | None = None,
-        **kwargs,
-    ) -> None:
-        # Intentionally bypass VLAAgent.__init__; its condition dimensions are
-        # tied to ResNet-style visual+state concatenation.
-        nn.Module.__init__(self)
-        if inference_steps != 1:
-            raise ValueError(
-                'SmolVLAIMFAttnResAgent only supports one-step IMF inference; '
-                f'inference_steps must be 1, got {inference_steps}.'
-            )
-        if head_type != 'transformer':
-            raise ValueError(f'SmolVLAIMFAttnResAgent requires head_type="transformer", got {head_type!r}')
-        del diffusion_steps, kwargs
-
-        self.action_dim = int(action_dim)
-        self.obs_dim = int(obs_dim)
-        self.pred_horizon = int(pred_horizon)
-        self.obs_horizon = int(obs_horizon)
-        self.num_cams = int(num_cams)
-        self.num_action_steps = int(num_action_steps)
-        self.inference_steps = 1
-        self.head_type = head_type
-        self.camera_names = tuple(camera_names) if camera_names is not None else None
-        if self.camera_names is not None and len(self.camera_names) != self.num_cams:
-            raise ValueError(f'camera_names length({len(self.camera_names)}) does not match num_cams({self.num_cams})')
-
-        self.normalization = NormalizationModule(stats=dataset_stats, normalization_type=normalization_type)
-        self.condition_encoder = condition_encoder
-        # Compatibility aliases used by some utilities/tests.
-        self.vision_encoder = condition_encoder
-        self.action_encoder = action_encoder
-        self.state_encoder = None
-        self.task_description = task_description
-
-        encoder_dim = getattr(condition_encoder, 'output_dim', None)
-        if encoder_dim is None:
-            encoder_dim = getattr(condition_encoder, 'joint_output_dim', None)
-        if condition_dim is None:
-            if encoder_dim is None:
-                raise ValueError('condition_dim must be provided when condition_encoder has no output_dim')
-            condition_dim = int(encoder_dim)
-        self.per_step_cond_dim = int(condition_dim)
-        self.raw_per_step_cond_dim = self.per_step_cond_dim
-
-        encoder_seq_len = getattr(condition_encoder, 'condition_sequence_length', None)
-        if condition_sequence_length is None:
-            if encoder_seq_len is None:
-                raise ValueError(
-                    'condition_sequence_length must be provided when condition_encoder has no condition_sequence_length'
-                )
-            condition_sequence_length = int(encoder_seq_len)
-        self.condition_sequence_length = int(condition_sequence_length)
-        self.condition_tokens_per_step = self.condition_sequence_length
-        self.global_cond_dim = self.per_step_cond_dim * self.condition_sequence_length
-
-        if isinstance(head, nn.Module):
-            self.noise_pred_net = head
-        else:
-            self.noise_pred_net = head(
-                input_dim=self.action_dim,
-                output_dim=self.action_dim,
-                horizon=self.pred_horizon,
-                n_obs_steps=self.condition_sequence_length,
-                cond_dim=self.per_step_cond_dim,
-            )
-        self.reset()
-
-    def _get_model_device(self) -> torch.device:
-        return next(self.parameters()).device
-
-    def _move_to_device(self, data, device: torch.device):
-        if torch.is_tensor(data):
-            return data.to(device)
-        if isinstance(data, dict):
-            return {k: self._move_to_device(v, device) for k, v in data.items()}
-        if isinstance(data, list):
-            return [self._move_to_device(v, device) for v in data]
-        if isinstance(data, tuple):
-            return tuple(self._move_to_device(v, device) for v in data)
-        return data
-
-    def _order_images(self, images: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
-        if self.camera_names is None:
-            names = tuple(sorted(images.keys()))
-            if len(names) != self.num_cams:
-                raise ValueError(f'image camera count({len(names)}) does not match num_cams({self.num_cams})')
-            return {name: images[name] for name in names}
-        missing = [name for name in self.camera_names if name not in images]
-        if missing:
-            raise ValueError(f'image condition missing required cameras. missing={missing}, expected={list(self.camera_names)}')
-        return {name: images[name] for name in self.camera_names}
-
-    def _resolve_task(self, task, batch_size: int):
-        def is_missing_task(item) -> bool:
-            return item is None or (isinstance(item, str) and item.strip().lower() in {'', 'unknown'})
-
-        def fallback_or(item):
-            if self.task_description is not None and is_missing_task(item):
-                return self.task_description
-            if item is None:
-                return ''
-            return item
-
-        if task is None:
-            task = self.task_description
-        if task is None:
-            return [''] * batch_size
-        if isinstance(task, str):
-            return [fallback_or(task)] * batch_size
-        task = list(task)
-        if len(task) != batch_size:
-            raise ValueError(f'task batch size ({len(task)}) must match batch size ({batch_size})')
-        return [fallback_or(item) for item in task]
-
-    def _build_cond(self, images: Dict[str, torch.Tensor], states: torch.Tensor, task=None) -> torch.Tensor:
-        ordered_images = self._order_images(images)
-        batch_size = states.shape[0]
-        tasks = self._resolve_task(task, batch_size=batch_size)
-        cond = self.condition_encoder(ordered_images, state=states, task=tasks)
-        if cond.ndim != 3:
-            raise RuntimeError(f'condition_encoder must return (B,S,D), got {tuple(cond.shape)}')
-        if cond.shape[0] != batch_size:
-            raise RuntimeError(f'condition batch mismatch: got {cond.shape[0]}, expected {batch_size}')
-        if cond.shape[1] != self.condition_sequence_length:
-            raise RuntimeError(
-                f'condition sequence length mismatch: got {cond.shape[1]}, expected {self.condition_sequence_length}'
-            )
-        if cond.shape[2] != self.per_step_cond_dim:
-            raise RuntimeError(f'condition dim mismatch: got {cond.shape[2]}, expected {self.per_step_cond_dim}')
-        head_dtype = next(self.noise_pred_net.parameters()).dtype
-        return cond.to(dtype=head_dtype)
-
-    def compute_loss(self, batch):
-        actions, states, images = batch['action'], batch['qpos'], batch['images']
-        action_is_pad = batch.get('action_is_pad', None)
-        batch_size = actions.shape[0]
-
-        states = self.normalization.normalize_qpos(states)
-        actions = self.normalization.normalize_action(actions)
-        cond = self._build_cond(images, states, task=batch.get('task', None))
-
-        x = actions
-        e = torch.randn_like(x)
-        t = torch.rand(batch_size, device=x.device, dtype=x.dtype)
-        r = torch.rand(batch_size, device=x.device, dtype=x.dtype)
-        t, r = torch.maximum(t, r), torch.minimum(t, r)
-
-        t_broadcast = self._broadcast_batch_time(t, x)
-        z_t = (1 - t_broadcast) * x + t_broadcast * e
-
-        v = self.fn(z_t, t, t, cond=cond)
-        u, du_dt = self._compute_u_and_du_dt(z_t, r, t, cond=cond, v=v)
-        V = self._compound_velocity(u, du_dt, r, t)
-        target = e - x
-
-        loss = nn.functional.mse_loss(V, target, reduction='none')
-        if action_is_pad is not None:
-            mask = (~action_is_pad).unsqueeze(-1).to(loss.dtype)
-            valid_count = mask.sum() * loss.shape[-1]
-            loss = (loss * mask).sum() / valid_count.clamp_min(1.0)
-        else:
-            loss = loss.mean()
-        return loss
-
-    @torch.no_grad()
-    def predict_action(self, images, proprioception, task=None):
-        batch_size = proprioception.shape[0]
-        proprioception = self.normalization.normalize_qpos(proprioception)
-        cond = self._build_cond(images, proprioception, task=task)
-        z_t = torch.randn((batch_size, self.pred_horizon, self.action_dim), device=cond.device, dtype=cond.dtype)
-        action = self._sample_one_step(z_t, cond=cond)
-        return self.normalization.denormalize_action(action)
-
-    @torch.no_grad()
-    def predict_action_chunk(self, batch: Dict[str, torch.Tensor]) -> torch.Tensor:
-        return self.predict_action(batch['images'], batch['qpos'], task=batch.get('task', None))
-
-    def reset(self):
-        self._queues = {
-            'qpos': deque(maxlen=self.obs_horizon),
-            'images': deque(maxlen=self.obs_horizon),
-            'task': deque(maxlen=self.obs_horizon),
-            'action': deque(maxlen=self.pred_horizon - self.obs_horizon + 1),
-        }
-
-    def _populate_queues(self, observation: Dict[str, torch.Tensor]) -> None:
-        if 'qpos' in observation:
-            self._queues['qpos'].append(observation['qpos'].clone())
-        if 'images' in observation:
-            ordered_images = self._order_images(observation['images'])
-            self._queues['images'].append({k: v.clone() for k, v in ordered_images.items()})
-        if 'task' in observation:
-            self._queues['task'].append(observation['task'])
-
-    def _prepare_observation_batch(self) -> Dict[str, torch.Tensor]:
-        qpos_list = list(self._queues['qpos'])
-        if not qpos_list:
-            raise ValueError('observation queue is empty.')
-        while len(qpos_list) < self.obs_horizon:
-            qpos_list.append(qpos_list[-1])
-        batch_qpos = torch.stack(qpos_list, dim=0).unsqueeze(0)
-
-        images_list = list(self._queues['images'])
-        if not images_list:
-            raise ValueError('image queue is empty.')
-        while len(images_list) < self.obs_horizon:
-            images_list.append(images_list[-1])
-        names = self.camera_names if self.camera_names is not None else tuple(sorted(images_list[0].keys()))
-        batch_images = {
-            name: torch.stack([item[name] for item in images_list], dim=0).unsqueeze(0)
-            for name in names
-        }
-        batch = {'qpos': batch_qpos, 'images': batch_images}
-        if self._queues['task']:
-            batch['task'] = [list(self._queues['task'])[-1]]
-        elif self.task_description is not None:
-            batch['task'] = [self.task_description]
-        return batch
-
-    @torch.no_grad()
-    def select_action(self, observation: Dict[str, torch.Tensor]) -> torch.Tensor:
-        device = self._get_model_device()
-        observation = self._move_to_device(observation, device)
-        self._populate_queues(observation)
-        if len(self._queues['action']) == 0:
-            batch = self._prepare_observation_batch()
-            actions = self.predict_action_chunk(batch)
-            start = self.obs_horizon - 1
-            end = start + self.num_action_steps
-            executable_actions = actions[:, start:end]
-            for i in range(executable_actions.shape[1]):
-                self._queues['action'].append(executable_actions[:, i].squeeze(0))
-        return self._queues['action'].popleft()
-
-    def get_normalization_stats(self):
-        return self.normalization.get_stats()

roboimi/vla/agent_smolvla_native.py

@@ -1,398 +0,0 @@
-from __future__ import annotations
-
-from collections import deque
-from dataclasses import fields
-from typing import Dict, Optional, Sequence
-
-import torch
-import torch.nn as nn
-
-from roboimi.vla.models.normalization import NormalizationModule
-
-
-class SmolVLANativeAgent(nn.Module):
-    """RoboIMI wrapper for the native SmolVLA flow-matching model.
-
-    The wrapper owns RoboIMI-facing concerns only: dataset normalization,
-    camera ordering, language fallback/tokenization, rollout queues, and action
-    denormalization.  The native SmolVLA model itself is imported lazily so unit
-    tests can inject fakes without downloading or loading a real VLM.
-    """
-
-    def __init__(
-        self,
-        model: Optional[nn.Module] = None,
-        tokenizer=None,
-        action_dim: int = 16,
-        obs_dim: int = 16,
-        chunk_size: int = 16,
-        n_action_steps: int = 8,
-        obs_horizon: int = 1,
-        action_chunk_start: int = 0,
-        num_cams: int = 3,
-        camera_names: Optional[Sequence[str]] = None,
-        dataset_stats=None,
-        normalization_type: str = 'min_max',
-        task_description: Optional[str] = None,
-        model_config: Optional[dict] = None,
-        tokenizer_name: Optional[str] = None,
-        **kwargs,
-    ) -> None:
-        super().__init__()
-        del kwargs
-        self.action_dim = int(action_dim)
-        self.obs_dim = int(obs_dim)
-        self.chunk_size = int(chunk_size)
-        self.pred_horizon = self.chunk_size
-        self.n_action_steps = int(n_action_steps)
-        self.num_action_steps = self.n_action_steps
-        self.obs_horizon = int(obs_horizon)
-        self.action_chunk_start = int(action_chunk_start)
-        self.num_cams = int(num_cams)
-        self.camera_names = tuple(camera_names) if camera_names is not None else None
-        if self.camera_names is not None and len(self.camera_names) != self.num_cams:
-            raise ValueError(f'camera_names length({len(self.camera_names)}) does not match num_cams({self.num_cams})')
-        if self.n_action_steps < 1:
-            raise ValueError('n_action_steps must be >= 1')
-        if self.action_chunk_start < 0:
-            raise ValueError('action_chunk_start must be >= 0')
-        if self.action_chunk_start + self.n_action_steps > self.chunk_size:
-            raise ValueError('action_chunk_start + n_action_steps must be <= chunk_size')
-
-        self.normalization = NormalizationModule(stats=dataset_stats, normalization_type=normalization_type)
-        self.task_description = task_description
-        self.model_config = dict(model_config or {})
-        self.max_state_dim = int(self.model_config.get('max_state_dim', self.obs_dim))
-        self.max_action_dim = int(self.model_config.get('max_action_dim', self.action_dim))
-        self.resize_imgs_with_padding = self._normalize_resize_shape(
-            self.model_config.get('resize_imgs_with_padding', self.model_config.get('image_resize_shape', (512, 512)))
-        )
-        self.tokenizer_max_length = int(self.model_config.get('tokenizer_max_length', 48))
-        self.pad_language_to = str(self.model_config.get('pad_language_to', 'longest'))
-        self.tokenizer_name = tokenizer_name
-        self.tokenizer = tokenizer if tokenizer is not None else self._build_tokenizer(tokenizer_name)
-        self.model = model if model is not None else self._build_model()
-        self.reset()
-
-    @staticmethod
-    def _normalize_resize_shape(shape):
-        if shape is None:
-            return None
-        normalized = tuple(int(v) for v in shape)
-        if len(normalized) != 2:
-            raise ValueError(f'resize_imgs_with_padding must contain exactly two values, got {normalized}')
-        return normalized
-
-    def _build_tokenizer(self, tokenizer_name):
-        if tokenizer_name is None:
-            tokenizer_name = self.model_config.get('tokenizer_name') or self.model_config.get('vlm_model_name')
-        if tokenizer_name is None:
-            raise ImportError('A tokenizer or tokenizer_name/model_config.vlm_model_name is required for SmolVLANativeAgent')
-        try:
-            from transformers import AutoTokenizer
-        except ImportError as exc:
-            raise ImportError('SmolVLANativeAgent requires transformers to construct the real tokenizer') from exc
-        return AutoTokenizer.from_pretrained(tokenizer_name)
-
-    def _native_config_kwargs(self) -> dict:
-        from roboimi.vla.models.smolvla.configuration import NativeSmolVLAConfig
-
-        cfg_kwargs = dict(self.model_config)
-
-        # Backwards-compatible aliases from earlier RoboIMI-facing drafts.  The
-        # native model core intentionally keeps only SmolVLA model fields.
-        if 'image_resize_shape' in cfg_kwargs and 'resize_imgs_with_padding' not in cfg_kwargs:
-            cfg_kwargs['resize_imgs_with_padding'] = cfg_kwargs['image_resize_shape']
-        if 'freeze_vlm' in cfg_kwargs and 'train_expert_only' not in cfg_kwargs:
-            cfg_kwargs['train_expert_only'] = bool(cfg_kwargs['freeze_vlm'])
-
-        for wrapper_only_key in (
-            'state_dim',
-            'action_dim',
-            'tokenizer_name',
-            'freeze_vlm',
-            'image_resize_shape',
-            'num_cameras',
-        ):
-            cfg_kwargs.pop(wrapper_only_key, None)
-
-        cfg_kwargs.setdefault('max_state_dim', self.max_state_dim)
-        cfg_kwargs.setdefault('max_action_dim', self.max_action_dim)
-        cfg_kwargs.setdefault('chunk_size', self.chunk_size)
-        cfg_kwargs.setdefault('n_action_steps', self.n_action_steps)
-        cfg_kwargs['resize_imgs_with_padding'] = self._normalize_resize_shape(
-            cfg_kwargs.get('resize_imgs_with_padding', self.resize_imgs_with_padding)
-        )
-
-        allowed_keys = {field.name for field in fields(NativeSmolVLAConfig)}
-        return {key: value for key, value in cfg_kwargs.items() if key in allowed_keys}
-
-    def _build_model(self):
-        try:
-            from roboimi.vla.models.smolvla.configuration import NativeSmolVLAConfig
-            from roboimi.vla.models.smolvla.modeling import VLAFlowMatching
-        except ImportError as exc:
-            raise ImportError(
-                'Native SmolVLA model modules are unavailable. Pass injected model/tokenizer fakes in tests, '
-                'or add roboimi.vla.models.smolvla.configuration/modeling for real construction.'
-            ) from exc
-        config = NativeSmolVLAConfig(**self._native_config_kwargs())
-        return VLAFlowMatching(config=config)
-
-    def _get_model_device(self) -> torch.device:
-        try:
-            return next(self.model.parameters()).device
-        except (StopIteration, AttributeError):
-            return torch.device('cpu')
-
-    def _move_to_device(self, data, device: torch.device):
-        if torch.is_tensor(data):
-            return data.to(device)
-        if isinstance(data, dict):
-            return {k: self._move_to_device(v, device) for k, v in data.items()}
-        if isinstance(data, list):
-            return [self._move_to_device(v, device) for v in data]
-        if isinstance(data, tuple):
-            return tuple(self._move_to_device(v, device) for v in data)
-        return data
-
-    def _order_images(self, images: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
-        if self.camera_names is None:
-            names = tuple(sorted(images.keys()))
-            if len(names) != self.num_cams:
-                raise ValueError(f'image camera count({len(names)}) does not match num_cams({self.num_cams})')
-            return {name: images[name] for name in names}
-        missing = [name for name in self.camera_names if name not in images]
-        if missing:
-            raise ValueError(f'image batch missing required cameras. missing={missing}, expected={list(self.camera_names)}')
-        return {name: images[name] for name in self.camera_names}
-
-    @staticmethod
-    def _is_missing_task(item) -> bool:
-        return item is None or (isinstance(item, str) and item.strip().lower() in {'', 'unknown'})
-
-    def _resolve_task(self, task, batch_size: int):
-        def fallback_or(item):
-            if self._is_missing_task(item):
-                if self.task_description is not None:
-                    return self.task_description
-                if item is None:
-                    return ''
-            return item
-
-        if task is None:
-            task = self.task_description
-        if task is None:
-            return [''] * batch_size
-        if isinstance(task, str):
-            return [fallback_or(task)] * batch_size
-        task = list(task)
-        if len(task) != batch_size:
-            raise ValueError(f'task batch size ({len(task)}) must match batch size ({batch_size})')
-        return [fallback_or(item) for item in task]
-
-    def _tokenize_tasks(self, task, batch_size: int, device: torch.device) -> dict:
-        tasks = []
-        for text in self._resolve_task(task, batch_size):
-            text = str(text)
-            tasks.append(text if text.endswith('\n') else f'{text}\n')
-        old_padding_side = getattr(self.tokenizer, 'padding_side', None)
-        if old_padding_side is not None:
-            self.tokenizer.padding_side = 'right'
-        try:
-            tokenized = self.tokenizer(
-                tasks,
-                padding=self.pad_language_to,
-                max_length=self.tokenizer_max_length,
-                truncation=True,
-                return_tensors='pt',
-            )
-        finally:
-            if old_padding_side is not None:
-                self.tokenizer.padding_side = old_padding_side
-        return {
-            'lang_tokens': tokenized['input_ids'].to(device=device),
-            'lang_masks': tokenized['attention_mask'].to(device=device, dtype=torch.bool),
-        }
-
-    @staticmethod
-    def _pad_vector(vector: torch.Tensor, new_dim: int) -> torch.Tensor:
-        current_dim = vector.shape[-1]
-        if current_dim == new_dim:
-            return vector
-        if current_dim > new_dim:
-            raise ValueError(f'cannot pad vector with dim {current_dim} to smaller dim {new_dim}')
-        padded_shape = list(vector.shape)
-        padded_shape[-1] = int(new_dim)
-        padded = torch.zeros(*padded_shape, dtype=vector.dtype, device=vector.device)
-        padded[..., :current_dim] = vector
-        return padded
-
-    @staticmethod
-    def _resize_with_pad(img: torch.Tensor, width: int, height: int, pad_value: float = 0.0) -> torch.Tensor:
-        if img.ndim != 4:
-            raise ValueError(f'expected image tensor shaped (B,C,H,W), got {tuple(img.shape)}')
-        import torch.nn.functional as F
-
-        cur_height, cur_width = img.shape[2:]
-        ratio = max(cur_width / width, cur_height / height)
-        resized_height = int(cur_height / ratio)
-        resized_width = int(cur_width / ratio)
-        resized = F.interpolate(img, size=(resized_height, resized_width), mode='bilinear', align_corners=False)
-        pad_height = max(0, int(height - resized_height))
-        pad_width = max(0, int(width - resized_width))
-        return F.pad(resized, (pad_width, 0, pad_height, 0), value=pad_value)
-
-    def _prepare_native_images(self, images: Dict[str, torch.Tensor]) -> tuple[list[torch.Tensor], list[torch.Tensor]]:
-        ordered_images = self._order_images(images)
-        prepared_images: list[torch.Tensor] = []
-        img_masks: list[torch.Tensor] = []
-        reference_batch_size: int | None = None
-        for camera_name, image in ordered_images.items():
-            if image.ndim == 5:
-                image = image[:, -1]
-            elif image.ndim != 4:
-                raise ValueError(
-                    f'image for camera {camera_name!r} must be shaped (B,T,C,H,W) or (B,C,H,W), got {tuple(image.shape)}'
-                )
-            if reference_batch_size is None:
-                reference_batch_size = int(image.shape[0])
-            elif int(image.shape[0]) != reference_batch_size:
-                raise ValueError(f'image batch size mismatch for camera {camera_name!r}')
-
-            image = image.contiguous().float().clamp(0.0, 1.0)
-            if self.resize_imgs_with_padding is not None:
-                image = self._resize_with_pad(image, *self.resize_imgs_with_padding, pad_value=0.0)
-            image = image * 2.0 - 1.0
-            prepared_images.append(image)
-            img_masks.append(torch.ones(image.shape[0], dtype=torch.bool, device=image.device))
-        return prepared_images, img_masks
-
-    def _prepare_native_state(self, states: torch.Tensor) -> torch.Tensor:
-        states = self.normalization.normalize_qpos(states)
-        if states.ndim > 2:
-            states = states[:, -1, :]
-        return self._pad_vector(states.float(), self.max_state_dim)
-
-    def _prepare_native_actions(self, actions: torch.Tensor) -> torch.Tensor:
-        actions = self.normalization.normalize_action(actions)
-        return self._pad_vector(actions.float(), self.max_action_dim)
-
-    def _prepare_model_inputs(self, batch: Dict[str, torch.Tensor], include_actions: bool = False) -> dict:
-        states = batch['qpos']
-        batch_size = states.shape[0]
-        device = states.device
-        images, img_masks = self._prepare_native_images(batch['images'])
-        inputs = {
-            'images': images,
-            'img_masks': img_masks,
-            'state': self._prepare_native_state(states),
-        }
-        inputs.update(self._tokenize_tasks(batch.get('task', None), batch_size, device))
-        if include_actions:
-            inputs['actions'] = self._prepare_native_actions(batch['action'])
-        return inputs
-
-    def _reduce_native_losses(self, losses: torch.Tensor, action_is_pad: torch.Tensor | None = None) -> torch.Tensor:
-        if losses.ndim == 0:
-            return losses
-        if losses.shape[-1] < self.action_dim:
-            raise RuntimeError(f'loss action dim mismatch: got {losses.shape[-1]}, expected at least {self.action_dim}')
-        losses = losses[..., : self.action_dim]
-        if action_is_pad is None:
-            return losses.mean()
-        if losses.ndim != 3:
-            raise RuntimeError(f'action padding mask requires per-element losses shaped (B,H,A), got {tuple(losses.shape)}')
-        if tuple(action_is_pad.shape) != tuple(losses.shape[:2]):
-            raise RuntimeError(
-                f'action_is_pad shape {tuple(action_is_pad.shape)} does not match loss prefix {tuple(losses.shape[:2])}'
-            )
-        mask = (~action_is_pad).to(device=losses.device, dtype=losses.dtype).unsqueeze(-1)
-        denom = (mask.sum() * losses.shape[-1]).clamp_min(1.0)
-        return (losses * mask).sum() / denom
-
-    def compute_loss(self, batch):
-        inputs = self._prepare_model_inputs(batch, include_actions=True)
-        action_is_pad = batch.get('action_is_pad', None)
-        output = self.model(**inputs)
-        if isinstance(output, dict):
-            if 'loss' not in output:
-                raise RuntimeError('SmolVLA model forward returned a dict without loss')
-            return output['loss']
-        if torch.is_tensor(output):
-            return self._reduce_native_losses(output, action_is_pad=action_is_pad)
-        if hasattr(output, 'loss'):
-            return output.loss
-        raise RuntimeError(f'Unsupported SmolVLA forward output type: {type(output)!r}')
-
-    @torch.no_grad()
-    def predict_action_chunk(self, batch: Dict[str, torch.Tensor]) -> torch.Tensor:
-        inputs = self._prepare_model_inputs(batch, include_actions=False)
-        if not hasattr(self.model, 'sample_actions'):
-            raise RuntimeError('SmolVLA model must implement sample_actions for inference')
-        actions = self.model.sample_actions(**inputs)
-        if not torch.is_tensor(actions) or actions.ndim != 3:
-            raise RuntimeError(f'sample_actions must return (B,H,A), got {type(actions)!r} {getattr(actions, "shape", None)}')
-        if actions.shape[-1] < self.action_dim:
-            raise RuntimeError(f'action dim mismatch: got {actions.shape[-1]}, expected at least {self.action_dim}')
-        actions = actions[..., : self.action_dim]
-        return self.normalization.denormalize_action(actions)
-
-    def reset(self):
-        self._queues = {
-            'qpos': deque(maxlen=self.obs_horizon),
-            'images': deque(maxlen=self.obs_horizon),
-            'task': deque(maxlen=self.obs_horizon),
-            'action': deque(maxlen=self.n_action_steps),
-        }
-
-    def _populate_queues(self, observation: Dict[str, torch.Tensor]) -> None:
-        if 'qpos' in observation:
-            self._queues['qpos'].append(observation['qpos'].clone())
-        if 'images' in observation:
-            ordered_images = self._order_images(observation['images'])
-            self._queues['images'].append({k: v.clone() for k, v in ordered_images.items()})
-        if 'task' in observation:
-            self._queues['task'].append(observation['task'])
-
-    def _prepare_observation_batch(self) -> Dict[str, torch.Tensor]:
-        qpos_list = list(self._queues['qpos'])
-        if not qpos_list:
-            raise ValueError('observation queue is empty.')
-        while len(qpos_list) < self.obs_horizon:
-            qpos_list.append(qpos_list[-1])
-        batch_qpos = torch.stack(qpos_list, dim=0).unsqueeze(0)
-
-        images_list = list(self._queues['images'])
-        if not images_list:
-            raise ValueError('image queue is empty.')
-        while len(images_list) < self.obs_horizon:
-            images_list.append(images_list[-1])
-        names = self.camera_names if self.camera_names is not None else tuple(sorted(images_list[0].keys()))
-        batch_images = {name: torch.stack([item[name] for item in images_list], dim=0).unsqueeze(0) for name in names}
-        batch = {'qpos': batch_qpos, 'images': batch_images}
-        if self._queues['task']:
-            batch['task'] = [list(self._queues['task'])[-1]]
-        elif self.task_description is not None:
-            batch['task'] = [self.task_description]
-        return batch
-
-    @torch.no_grad()
-    def select_action(self, observation: Dict[str, torch.Tensor]) -> torch.Tensor:
-        device = self._get_model_device()
-        observation = self._move_to_device(observation, device)
-        self._populate_queues(observation)
-        if len(self._queues['action']) == 0:
-            batch = self._prepare_observation_batch()
-            actions = self.predict_action_chunk(batch)
-            start = self.action_chunk_start
-            end = start + self.n_action_steps
-            if actions.shape[1] < end:
-                raise RuntimeError(f'action chunk too short: got {actions.shape[1]}, need at least {end}')
-            executable_actions = actions[:, start:end]
-            for i in range(executable_actions.shape[1]):
-                self._queues['action'].append(executable_actions[:, i].squeeze(0))
-        return self._queues['action'].popleft()
-
-    def get_normalization_stats(self):
-        return self.normalization.get_stats()

roboimi/vla/conf/agent/lewm_resnet_dual_decoder_imf_attnres.yaml

@@ -0,0 +1,74 @@
+# @package agent
+defaults:
+  - /backbone@vision_backbone: lewm_resnet_query_fusion
+  - /modules@state_encoder: lewm_state_encoder
+  - /modules@action_encoder: identity_action_encoder
+  - /modules@cond_projector: linear_condition_projector
+  - /head: imf_transformer1d
+  - /head@future_decoder: imf_transformer1d
+  - _self_
+
+_target_: roboimi.vla.agent_imf.IMFVLAAgent
+
+action_dim: 16
+obs_dim: 16
+normalization_type: "min_max"
+pred_horizon: 8
+obs_horizon: 2
+num_action_steps: 8
+camera_names: ${data.camera_names}
+num_cams: 3
+
+vision_backbone:
+  camera_names: ${agent.camera_names}
+  num_views: ${agent.num_cams}
+
+cond_projector:
+  output_dim: 288
+
+lewm_history_horizon: 3
+lewm_query_offsets: [8]
+extra_condition_tokens: 0
+lewm_loss_weight: 1.0
+lewm_sigreg_weight: 0.09
+lewm_pretrained_ckpt: null
+future_query_init_std: 0.02
+
+lewm_sigreg:
+  _target_: roboimi.vla.models.backbones.lewm_resnet_query_fusion.SIGReg
+  knots: 17
+  num_proj: 1024
+
+diffusion_steps: 100
+inference_steps: 1
+head_type: "transformer"
+
+head:
+  input_dim: ${agent.action_dim}
+  output_dim: ${agent.action_dim}
+  horizon: ${agent.pred_horizon}
+  n_obs_steps: ${agent.obs_horizon}
+  cond_dim: ${agent.cond_projector.output_dim}
+  n_emb: 384
+  causal_attn: false
+  time_as_cond: true
+  obs_as_cond: true
+  n_cond_layers: 0
+  backbone_type: attnres_full
+  n_head: 1
+  n_kv_head: 1
+
+future_decoder:
+  input_dim: ${agent.cond_projector.output_dim}
+  output_dim: ${agent.cond_projector.output_dim}
+  horizon: ${len:${agent.lewm_query_offsets}}
+  n_obs_steps: ${agent.lewm_history_horizon}
+  cond_dim: ${agent.cond_projector.output_dim}
+  n_emb: 384
+  causal_attn: false
+  time_as_cond: true
+  obs_as_cond: true
+  n_cond_layers: 0
+  backbone_type: attnres_full
+  n_head: 1
+  n_kv_head: 1

roboimi/vla/conf/agent/lewm_resnet_query_imf_attnres.yaml

@@ -0,0 +1,77 @@
+# @package agent
+defaults:
+  - /backbone@vision_backbone: lewm_resnet_query_fusion
+  - /modules@state_encoder: lewm_state_encoder
+  - /modules@action_encoder: identity_action_encoder
+  - /modules@cond_projector: linear_condition_projector
+  - /head: imf_transformer1d
+  - _self_
+
+_target_: roboimi.vla.agent_imf.IMFVLAAgent
+
+action_dim: 16
+obs_dim: 16
+normalization_type: "min_max"
+pred_horizon: 8
+obs_horizon: 2
+num_action_steps: 8
+camera_names: ${data.camera_names}
+num_cams: 3
+
+vision_backbone:
+  camera_names: ${agent.camera_names}
+  num_views: ${agent.num_cams}
+
+cond_projector:
+  output_dim: 288
+
+lewm_history_horizon: 3
+lewm_query_offsets: [8]
+extra_condition_tokens: ${len:${agent.lewm_query_offsets}}
+lewm_loss_weight: 1.0
+lewm_sigreg_weight: 0.09
+lewm_pretrained_ckpt: null
+
+lewm_sigreg:
+  _target_: roboimi.vla.models.backbones.lewm_resnet_query_fusion.SIGReg
+  knots: 17
+  num_proj: 1024
+
+lewm_predictor:
+  _target_: roboimi.vla.models.backbones.lewm_resnet_query_fusion.QueryTokenPredictor
+  num_frames: ${agent.lewm_history_horizon}
+  query_offsets: ${agent.lewm_query_offsets}
+  input_dim: ${agent.cond_projector.output_dim}
+  hidden_dim: ${agent.cond_projector.output_dim}
+  output_dim: ${agent.cond_projector.output_dim}
+  depth: 6
+  heads: 16
+  mlp_dim: 2048
+  dim_head: 64
+  dropout: 0.1
+  emb_dropout: 0.0
+
+lewm_pred_projector:
+  _target_: roboimi.vla.models.backbones.lewm_resnet_query_fusion.LeWMProjectorMLP
+  input_dim: ${agent.cond_projector.output_dim}
+  hidden_dim: 2048
+  output_dim: ${agent.cond_projector.output_dim}
+
+diffusion_steps: 100
+inference_steps: 1
+head_type: "transformer"
+
+head:
+  input_dim: ${agent.action_dim}
+  output_dim: ${agent.action_dim}
+  horizon: ${agent.pred_horizon}
+  n_obs_steps: ${agent.obs_horizon}
+  cond_dim: 288
+  n_emb: 384
+  causal_attn: false
+  time_as_cond: true
+  obs_as_cond: true
+  n_cond_layers: 0
+  backbone_type: attnres_full
+  n_head: 1
+  n_kv_head: 1

roboimi/vla/conf/agent/smolvla_imf_attnres.yaml

@@ -1,46 +0,0 @@
-# @package agent
-defaults:
-  - /backbone@condition_encoder: smolvla_prefix_encoder
-  - /modules@action_encoder: identity_action_encoder
-  - /head: imf_transformer1d
-  - _self_
-
-_target_: roboimi.vla.agent_smolvla_conditioned.SmolVLAIMFAttnResAgent
-
-action_dim: 16
-obs_dim: 16
-normalization_type: "min_max"
-pred_horizon: 16
-obs_horizon: 2
-num_action_steps: 8
-camera_names: ${data.camera_names}
-num_cams: ${len:${agent.camera_names}}
-
-# Optional fallback language instruction used only when a batch/observation
-# does not provide `task`. Keep null by default so this architecture remains
-# generic and dataset/eval can supply variable language.
-task_description: null
-condition_dim: 960
-condition_sequence_length: 241
-
-condition_encoder:
-  num_cameras: ${agent.num_cams}
-  camera_names: ${agent.camera_names}
-
-diffusion_steps: 100
-inference_steps: 1
-head_type: "transformer"
-
-head:
-  input_dim: ${agent.action_dim}
-  output_dim: ${agent.action_dim}
-  horizon: ${agent.pred_horizon}
-  n_obs_steps: ${agent.condition_sequence_length}
-  cond_dim: ${agent.condition_dim}
-  causal_attn: false
-  time_as_cond: true
-  obs_as_cond: true
-  n_cond_layers: 0
-  backbone_type: attnres_full
-  n_head: 1
-  n_kv_head: 1

roboimi/vla/conf/agent/smolvla_native.yaml

@@ -1,35 +0,0 @@
-# @package agent
-_target_: roboimi.vla.agent_smolvla_native.SmolVLANativeAgent
-
-model: null
-tokenizer: null
-tokenizer_name: ${agent.model_config.vlm_model_name}
-
-action_dim: 16
-obs_dim: 16
-normalization_type: "gaussian"
-chunk_size: 32
-pred_horizon: ${agent.chunk_size}
-obs_horizon: 2
-n_action_steps: 16
-num_action_steps: ${agent.n_action_steps}
-action_chunk_start: 0
-camera_names: ${data.camera_names}
-num_cams: ${len:${agent.camera_names}}
-task_description: null
-# SmolVLA performs its own SigLIP-style resize+pad inside the wrapper/core.
-# Keeping dataset/eval resize disabled avoids lossy double resizing.
-dataset_image_resize_shape: null
-eval_image_resize_shape: null
-
-model_config:
-  max_state_dim: 32
-  max_action_dim: 32
-  chunk_size: ${agent.chunk_size}
-  n_action_steps: ${agent.n_action_steps}
-  vlm_model_name: HuggingFaceTB/SmolVLM2-500M-Video-Instruct
-  load_vlm_weights: true
-  train_expert_only: true
-  freeze_vision_encoder: true
-  num_vlm_layers: 16
-  resize_imgs_with_padding: [512, 512]

roboimi/vla/conf/backbone/lewm_resnet_query_fusion.yaml

@@ -0,0 +1,7 @@
+_target_: roboimi.vla.models.backbones.lewm_resnet_query_fusion.LeWMMultiViewResNetBackbone
+
+view_feature_dim: 96
+num_views: ${agent.num_cams}
+view_encoder_mode: separate
+camera_names: ${agent.camera_names}
+checkpoint_path: null

roboimi/vla/conf/backbone/smolvla_prefix_encoder.yaml

@@ -1,19 +0,0 @@
-_target_: roboimi.vla.models.backbones.smolvla_prefix_encoder.SmolVLAPrefixEncoder
-
-model_name: HuggingFaceTB/SmolVLM2-500M-Video-Instruct
-load_vlm_weights: true
-local_files_only: false
-num_vlm_layers: 16
-freeze_vlm: true
-freeze_vision_encoder: true
-train_state_proj: true
-max_state_dim: 32
-resize_imgs_with_padding: [512, 512]
-tokenizer_max_length: 48
-pad_language_to: max_length
-run_text_model: true
-camera_names: [r_vis, top, front]
-num_cameras: 3
-
-dataset_image_resize_shape: null
-eval_image_resize_shape: null

roboimi/vla/conf/config.yaml

@@ -18,6 +18,8 @@ train:
   # 数据加载
   num_workers: 12     # DataLoader 工作进程数（调试时设为 0）
   val_split: 0.0      # 验证集比例；默认使用全量数据训练
+  val_episode_indices: null  # 显式按 episode 划出的验证集，例如 [100]
+  action_mse_val_freq_epochs: 0  # >0 时每隔多少个 epoch 在 held-out episode 上计算 action MSE
   seed: 42            # 随机种子（用于数据划分）
 
   # 日志和检查点

roboimi/vla/conf/data/simpe_robot_dataset.yaml

@@ -11,6 +11,8 @@ dataset_dir: "roboimi/demos/dataset/sim_transfer"
 # ====================
 pred_horizon: ${agent.pred_horizon}  # 预测步数
 obs_horizon: ${agent.obs_horizon}      # 观测步数
+lewm_history_horizon: ${oc.select:agent.lewm_history_horizon,null}
+lewm_query_offsets: ${oc.select:agent.lewm_query_offsets,null}
 
 # ====================
 # 相机配置

roboimi/vla/conf/eval/eval.yaml

@@ -9,7 +9,6 @@ server_startup_timeout_s: 300.0 # parent 等待 inference server 就绪的超时
 max_timesteps: 700        # 每回合最大时间步
 device: ${train.device}   # 与训练保持一致
 task_name: "sim_transfer" # 环境任务名称
-task_description: null    # 可选语言指令；环境 obs 没有 task 时注入给语言条件策略
 
 # ====================
 # 策略执行参数

roboimi/vla/conf/modules/lewm_state_encoder.yaml

@@ -0,0 +1,5 @@
+_target_: roboimi.vla.modules.encoders.LeWMStateEncoder
+
+input_dim: ${agent.obs_dim}
+hidden_dim: 256
+output_dim: 64

roboimi/vla/data/simpe_robot_dataset.py

@@ -4,7 +4,6 @@ from torch.utils.data import Dataset
 from typing import List, Dict, Union, Optional, Sequence
 from pathlib import Path
 from collections import OrderedDict
-import re
 
 
 class SimpleRobotDataset(Dataset):
@@ -25,6 +24,8 @@ class SimpleRobotDataset(Dataset):
         camera_names: List[str] = None,
         image_resize_shape: Optional[Sequence[int]] = (224, 224),
         max_open_files: int = 64,
+        lewm_history_horizon: Optional[int] = None,
+        lewm_query_offsets: Optional[Sequence[int]] = None,
         episode_indices: Optional[Sequence[int]] = None,
     ):
         """
@@ -35,7 +36,6 @@ class SimpleRobotDataset(Dataset):
             camera_names: 相机名称列表，如 ["r_vis", "top", "front"]
             image_resize_shape: 图像缩放尺寸 (W, H)；为 None 时保留原始分辨率
             max_open_files: 每个 worker 最多缓存的 HDF5 文件句柄数
-            episode_indices: 可选的原始 episode 编号子集
 
         HDF5 文件格式:
             - action: [T, action_dim]
@@ -45,6 +45,13 @@ class SimpleRobotDataset(Dataset):
         self.obs_horizon = obs_horizon
         self.pred_horizon = pred_horizon
         self.camera_names = camera_names or []
+        self.lewm_history_horizon = (
+            int(lewm_history_horizon) if lewm_history_horizon is not None else None
+        )
+        self.lewm_query_offsets = (
+            tuple(int(offset) for offset in lewm_query_offsets)
+            if lewm_query_offsets is not None else ()
+        )
         self.image_resize_shape = (
             tuple(int(v) for v in image_resize_shape)
             if image_resize_shape is not None else None
@@ -65,13 +72,19 @@ class SimpleRobotDataset(Dataset):
             self.hdf5_files = sorted(self.dataset_dir.glob("episode_*.hdf5"))
         if not self.hdf5_files:
             raise FileNotFoundError(f"在 {dataset_dir} 中未找到 HDF5 文件")
+
         if self.requested_episode_indices is not None:
             requested = set(self.requested_episode_indices)
             filtered = []
             for hdf5_path in self.hdf5_files:
-                match = re.search(r'episode_(\d+)$', hdf5_path.stem)
+                stem = hdf5_path.stem
-                if match and int(match.group(1)) in requested:
+                if stem.startswith("episode_"):
-                    filtered.append(hdf5_path)
+                    try:
+                        idx = int(stem.split("_")[-1])
+                    except ValueError:
+                        continue
+                    if idx in requested:
+                        filtered.append(hdf5_path)
             self.hdf5_files = filtered
             if not self.hdf5_files:
                 raise FileNotFoundError(
@@ -85,9 +98,12 @@ class SimpleRobotDataset(Dataset):
             with h5py.File(hdf5_path, 'r') as f:
                 T = f['action'].shape[0]
                 dataset_episode_idx = ep_idx
-                match = re.search(r'episode_(\d+)$', hdf5_path.stem)
+                stem = hdf5_path.stem
-                if match:
+                if stem.startswith("episode_"):
-                    dataset_episode_idx = int(match.group(1))
+                    try:
+                        dataset_episode_idx = int(stem.split("_")[-1])
+                    except ValueError:
+                        pass
                 start_idx = len(self.frame_meta)
                 for t in range(T):
                     self.frame_meta.append({
@@ -242,6 +258,60 @@ class SimpleRobotDataset(Dataset):
         for cam_name in self.camera_names:
             result[f"observation.{cam_name}"] = torch.stack(observations[f"observation.{cam_name}"])
 
+        if self.lewm_history_horizon is not None and self.lewm_history_horizon > 0:
+            lewm_observations = {
+                "state": [],
+            }
+            for cam_name in self.camera_names:
+                lewm_observations[f"observation.{cam_name}"] = []
+
+            for delta in range(-self.lewm_history_horizon + 1, 1):
+                target_idx = idx + delta
+                if ep_start <= target_idx <= ep_end:
+                    target_frame = self._load_frame(target_idx)
+                else:
+                    boundary_idx = ep_start if target_idx < ep_start else ep_end
+                    target_frame = self._load_frame(boundary_idx)
+
+                lewm_observations["state"].append(target_frame["observation.state"])
+                for cam_name in self.camera_names:
+                    lewm_observations[f"observation.{cam_name}"].append(
+                        target_frame[f"observation.{cam_name}"]
+                    )
+
+            result["lewm.observation.state"] = torch.stack(lewm_observations["state"])
+            for cam_name in self.camera_names:
+                result[f"lewm.observation.{cam_name}"] = torch.stack(
+                    lewm_observations[f"observation.{cam_name}"]
+                )
+
+        if self.lewm_query_offsets:
+            lewm_future = {
+                "state": [],
+            }
+            for cam_name in self.camera_names:
+                lewm_future[f"observation.{cam_name}"] = []
+
+            for offset in self.lewm_query_offsets:
+                target_idx = idx + offset
+                if ep_start <= target_idx <= ep_end:
+                    target_frame = self._load_frame(target_idx)
+                else:
+                    boundary_idx = ep_start if target_idx < ep_start else ep_end
+                    target_frame = self._load_frame(boundary_idx)
+
+                lewm_future["state"].append(target_frame["observation.state"])
+                for cam_name in self.camera_names:
+                    lewm_future[f"observation.{cam_name}"].append(
+                        target_frame[f"observation.{cam_name}"]
+                    )
+
+            result["lewm.future.state"] = torch.stack(lewm_future["state"])
+            for cam_name in self.camera_names:
+                result[f"lewm.future.{cam_name}"] = torch.stack(
+                    lewm_future[f"observation.{cam_name}"]
+                )
+
         return result
 
     @property

roboimi/vla/models/backbones/__init__.py

@@ -1,5 +1,14 @@
 # Backbone models
-__all__ = ["LEWMViTBackbone", "ResNetBackbone", "ResNetDiffusionBackbone", "SigLIP2DiffusionBackbone"]
+__all__ = [
+    "LEWMViTBackbone",
+    "LeWMMultiViewResNetBackbone",
+    "QueryTokenPredictor",
+    "LeWMProjectorMLP",
+    "SIGReg",
+    "ResNetBackbone",
+    "ResNetDiffusionBackbone",
+    "SigLIP2DiffusionBackbone",
+]
 
 
 def __getattr__(name):
@@ -9,6 +18,19 @@ def __getattr__(name):
     if name == "SigLIP2DiffusionBackbone":
         from .siglip2_diffusion_backbone import SigLIP2DiffusionBackbone
         return SigLIP2DiffusionBackbone
+    if name in {"LeWMMultiViewResNetBackbone", "QueryTokenPredictor", "LeWMProjectorMLP", "SIGReg"}:
+        from .lewm_resnet_query_fusion import (
+            LeWMMultiViewResNetBackbone,
+            QueryTokenPredictor,
+            LeWMProjectorMLP,
+            SIGReg,
+        )
+        return {
+            "LeWMMultiViewResNetBackbone": LeWMMultiViewResNetBackbone,
+            "QueryTokenPredictor": QueryTokenPredictor,
+            "LeWMProjectorMLP": LeWMProjectorMLP,
+            "SIGReg": SIGReg,
+        }[name]
     if name in {"ResNetBackbone", "ResNetDiffusionBackbone"}:
         from .resnet_diffusion import ResNetDiffusionBackbone
         return ResNetDiffusionBackbone

roboimi/vla/models/backbones/lewm_resnet_query_fusion.py

@@ -0,0 +1,409 @@
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any, Dict, Mapping, Optional, Sequence
+
+import torch
+from einops import rearrange
+from torch import nn
+import torch.nn.functional as F
+from torchvision import models
+
+from roboimi.vla.core.interfaces import VLABackbone
+
+
+class SpatialSoftmax2D(nn.Module):
+    """Convert a feature map into expected 2D keypoint coordinates per channel."""
+
+    def forward(self, feature_map):
+        if feature_map.ndim != 4:
+            raise ValueError(
+                f"SpatialSoftmax2D expects a 4D tensor, got rank {feature_map.ndim}"
+            )
+
+        batch, channels, height, width = feature_map.shape
+        scores = feature_map.reshape(batch, channels, height * width)
+        attention = F.softmax(scores, dim=-1)
+
+        ys = torch.linspace(-1.0, 1.0, height, device=feature_map.device, dtype=feature_map.dtype)
+        xs = torch.linspace(-1.0, 1.0, width, device=feature_map.device, dtype=feature_map.dtype)
+        grid_y, grid_x = torch.meshgrid(ys, xs, indexing="ij")
+        grid_x = grid_x.reshape(1, 1, height * width)
+        grid_y = grid_y.reshape(1, 1, height * width)
+
+        expected_x = (attention * grid_x).sum(dim=-1)
+        expected_y = (attention * grid_y).sum(dim=-1)
+        return torch.cat([expected_x, expected_y], dim=-1)
+
+
+class ResNet18SpatialEncoder(nn.Module):
+    """Encode one camera view into a fixed-dimensional spatial-softmax embedding."""
+
+    def __init__(self, view_feature_dim=96):
+        super().__init__()
+        if view_feature_dim % 2 != 0:
+            raise ValueError("view_feature_dim must be even for spatial softmax features")
+
+        backbone = models.resnet18(weights=None)
+        if all(
+            hasattr(backbone, name)
+            for name in ("conv1", "bn1", "relu", "maxpool", "layer1", "layer2", "layer3", "layer4")
+        ):
+            self.backbone = nn.Sequential(
+                backbone.conv1,
+                backbone.bn1,
+                backbone.relu,
+                backbone.maxpool,
+                backbone.layer1,
+                backbone.layer2,
+                backbone.layer3,
+                backbone.layer4,
+            )
+            feature_channels = 512
+        else:
+            children = list(backbone.children())
+            if len(children) < 1:
+                raise ValueError("resnet18 backbone must expose child modules")
+            truncated = children[:-2] if len(children) > 2 else children
+            self.backbone = nn.Sequential(*truncated)
+            with torch.no_grad():
+                dummy = torch.zeros(1, 3, 16, 16)
+                feature_channels = int(self.backbone(dummy).shape[1])
+
+        self.proj = nn.Conv2d(feature_channels, view_feature_dim // 2, kernel_size=1)
+        self.spatial_softmax = SpatialSoftmax2D()
+        self.output_dim = int(view_feature_dim)
+
+    def forward(self, pixels):
+        if pixels.ndim not in (4, 5):
+            raise ValueError(
+                f"ResNet18SpatialEncoder expects a 4D or 5D tensor, got rank {pixels.ndim}"
+            )
+
+        needs_unflatten = pixels.ndim == 5
+        if needs_unflatten:
+            batch, steps, channels, height, width = pixels.shape
+            pixels = rearrange(pixels, "b t c h w -> (b t) c h w")
+
+        features = self.backbone(pixels.float())
+        features = self.proj(features)
+        embeddings = self.spatial_softmax(features)
+
+        if needs_unflatten:
+            embeddings = rearrange(embeddings, "(b t) d -> b t d", b=batch, t=steps)
+        return embeddings
+
+
+class LeWMMultiViewResNetBackbone(VLABackbone):
+    """RoboIMI-side LeWM multiview ResNet spatial-softmax encoder."""
+
+    def __init__(
+        self,
+        view_feature_dim: int = 96,
+        num_views: int = 3,
+        view_encoder_mode: str = "shared",
+        camera_names: Sequence[str] = ("r_vis", "top", "front"),
+        checkpoint_path: str | Path | None = None,
+    ) -> None:
+        super().__init__()
+        if view_encoder_mode not in {"shared", "separate"}:
+            raise ValueError(
+                f"view_encoder_mode must be 'shared' or 'separate', got {view_encoder_mode}"
+            )
+
+        self.view_feature_dim = int(view_feature_dim)
+        self.num_views = int(num_views)
+        self.view_encoder_mode = view_encoder_mode
+        self.camera_names = tuple(camera_names)
+        if len(self.camera_names) != self.num_views:
+            raise ValueError(
+                f"camera_names length({len(self.camera_names)}) must equal num_views({self.num_views})"
+            )
+        self.output_dim = self.view_feature_dim * self.num_views
+        self.joint_output_dim = self.output_dim
+        self.tokens_per_step = 1
+
+        if view_encoder_mode == "shared":
+            self.single_view_encoder = ResNet18SpatialEncoder(
+                view_feature_dim=view_feature_dim
+            )
+            self.view_encoders = None
+        else:
+            self.single_view_encoder = None
+            self.view_encoders = nn.ModuleList(
+                [
+                    ResNet18SpatialEncoder(view_feature_dim=view_feature_dim)
+                    for _ in range(num_views)
+                ]
+            )
+
+        if checkpoint_path is not None:
+            self.load_lewm_checkpoint(checkpoint_path)
+
+    @staticmethod
+    def _unwrap_state_dict(payload: Mapping[str, Any]) -> Mapping[str, torch.Tensor]:
+        state_dict = payload.get("state_dict", payload)
+        if not isinstance(state_dict, Mapping):
+            raise TypeError("checkpoint payload must contain a mapping state_dict")
+        return state_dict
+
+    @staticmethod
+    def _extract_prefixed_state_dict(
+        state_dict: Mapping[str, torch.Tensor],
+        prefix: str,
+    ) -> Dict[str, torch.Tensor]:
+        extracted = {
+            key[len(prefix):]: value
+            for key, value in state_dict.items()
+            if key.startswith(prefix)
+        }
+        if not extracted:
+            raise KeyError(f"checkpoint missing parameters with prefix {prefix!r}")
+        return extracted
+
+    def load_lewm_checkpoint(self, checkpoint_or_path: str | Path | Mapping[str, Any]) -> None:
+        if isinstance(checkpoint_or_path, (str, Path)):
+            payload = torch.load(Path(checkpoint_or_path), map_location="cpu", weights_only=False)
+        else:
+            payload = checkpoint_or_path
+        state_dict = self._unwrap_state_dict(payload)
+        encoder_state_dict = self._extract_prefixed_state_dict(state_dict, "model.encoder.")
+        self.load_state_dict(encoder_state_dict, strict=True)
+
+    def forward(self, images):
+        missing = [camera_name for camera_name in self.camera_names if camera_name not in images]
+        if missing:
+            raise ValueError(
+                f"image input missing required cameras. missing={missing}, expected={list(self.camera_names)}"
+            )
+
+        first_image = images[self.camera_names[0]]
+        batch_size, steps = first_image.shape[:2]
+        view_embeddings = []
+        if self.view_encoder_mode == "shared":
+            for camera_name in self.camera_names:
+                view_embeddings.append(self.single_view_encoder(images[camera_name]))
+        else:
+            for single_view_encoder, camera_name in zip(self.view_encoders, self.camera_names):
+                view_embeddings.append(single_view_encoder(images[camera_name]))
+
+        embeddings = torch.cat(view_embeddings, dim=-1)
+        return embeddings.reshape(batch_size, steps, self.output_dim)
+
+
+class FeedForward(nn.Module):
+    def __init__(self, dim, hidden_dim, dropout=0.0):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.LayerNorm(dim),
+            nn.Linear(dim, hidden_dim),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(hidden_dim, dim),
+            nn.Dropout(dropout),
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+
+class Attention(nn.Module):
+    def __init__(self, dim, heads=8, dim_head=64, dropout=0.0):
+        super().__init__()
+        inner_dim = dim_head * heads
+        project_out = not (heads == 1 and dim_head == dim)
+        self.heads = heads
+        self.dropout = dropout
+        self.norm = nn.LayerNorm(dim)
+        self.to_qkv = nn.Linear(dim, inner_dim * 3, bias=False)
+        self.to_out = (
+            nn.Sequential(nn.Linear(inner_dim, dim), nn.Dropout(dropout))
+            if project_out
+            else nn.Identity()
+        )
+
+    def forward(self, x, causal=True):
+        x = self.norm(x)
+        drop = self.dropout if self.training else 0.0
+        qkv = self.to_qkv(x).chunk(3, dim=-1)
+        q, k, v = (rearrange(t, "b t (h d) -> b h t d", h=self.heads) for t in qkv)
+        out = F.scaled_dot_product_attention(q, k, v, dropout_p=drop, is_causal=causal)
+        out = rearrange(out, "b h t d -> b t (h d)")
+        return self.to_out(out)
+
+
+class Block(nn.Module):
+    def __init__(self, dim, heads, dim_head, mlp_dim, dropout=0.0):
+        super().__init__()
+        self.attn = Attention(dim, heads=heads, dim_head=dim_head, dropout=dropout)
+        self.mlp = FeedForward(dim, mlp_dim, dropout=dropout)
+        self.norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+
+    def forward(self, x):
+        x = x + self.attn(self.norm1(x))
+        x = x + self.mlp(self.norm2(x))
+        return x
+
+
+class Transformer(nn.Module):
+    def __init__(
+        self,
+        input_dim,
+        hidden_dim,
+        output_dim,
+        depth,
+        heads,
+        dim_head,
+        mlp_dim,
+        dropout=0.0,
+        block_class=Block,
+    ):
+        super().__init__()
+        self.norm = nn.LayerNorm(hidden_dim)
+        self.layers = nn.ModuleList([])
+
+        self.input_proj = (
+            nn.Linear(input_dim, hidden_dim)
+            if input_dim != hidden_dim
+            else nn.Identity()
+        )
+        self.cond_proj = (
+            nn.Linear(input_dim, hidden_dim)
+            if input_dim != hidden_dim
+            else nn.Identity()
+        )
+        self.output_proj = (
+            nn.Linear(hidden_dim, output_dim)
+            if hidden_dim != output_dim
+            else nn.Identity()
+        )
+
+        for _ in range(depth):
+            self.layers.append(block_class(hidden_dim, heads, dim_head, mlp_dim, dropout))
+
+    def forward(self, x, c=None):
+        x = self.input_proj(x)
+        if c is not None:
+            c = self.cond_proj(c)
+        for block in self.layers:
+            x = block(x)
+        x = self.norm(x)
+        return self.output_proj(x)
+
+
+class QueryTokenPredictor(nn.Module):
+    """History-only transformer predictor that decodes learned query tokens."""
+
+    def __init__(
+        self,
+        *,
+        num_frames,
+        query_offsets,
+        depth,
+        heads,
+        mlp_dim,
+        input_dim,
+        hidden_dim,
+        output_dim=None,
+        dim_head=64,
+        dropout=0.0,
+        emb_dropout=0.0,
+    ):
+        super().__init__()
+        if num_frames <= 0:
+            raise ValueError(f"num_frames must be positive, got {num_frames}")
+
+        query_offsets = tuple(query_offsets)
+        if not query_offsets:
+            raise ValueError("query_offsets must contain at least one offset")
+        if any(offset <= 0 for offset in query_offsets):
+            raise ValueError(f"query_offsets must be positive, got {query_offsets}")
+
+        self.num_frames = int(num_frames)
+        self.query_offsets = query_offsets
+        self.num_query_tokens = len(query_offsets)
+        self.pos_embedding = nn.Parameter(
+            torch.randn(1, self.num_frames + self.num_query_tokens, input_dim)
+        )
+        self.query_tokens = nn.Parameter(
+            torch.randn(1, self.num_query_tokens, input_dim)
+        )
+        self.dropout = nn.Dropout(emb_dropout)
+        self.transformer = Transformer(
+            input_dim,
+            hidden_dim,
+            output_dim or input_dim,
+            depth,
+            heads,
+            dim_head,
+            mlp_dim,
+            dropout,
+            block_class=Block,
+        )
+
+    def forward(self, x):
+        if x.ndim != 3:
+            raise ValueError(
+                f"QueryTokenPredictor expects a 3D tensor, got rank {x.ndim}"
+            )
+
+        T = x.size(1)
+        if T > self.num_frames:
+            raise ValueError(
+                f"input sequence length {T} exceeds configured num_frames {self.num_frames}"
+            )
+
+        query_tokens = self.query_tokens.expand(x.size(0), -1, -1)
+        tokens = torch.cat([x, query_tokens], dim=1)
+        tokens = tokens + self.pos_embedding[:, : tokens.size(1)]
+        tokens = self.dropout(tokens)
+        tokens = self.transformer(tokens)
+        return tokens[:, -self.num_query_tokens :]
+
+
+class LeWMProjectorMLP(nn.Module):
+    def __init__(
+        self,
+        input_dim: int = 288,
+        hidden_dim: int = 2048,
+        output_dim: int = 288,
+    ) -> None:
+        super().__init__()
+        self.output_dim = int(output_dim)
+        self.net = nn.Sequential(
+            nn.Linear(int(input_dim), int(hidden_dim)),
+            nn.BatchNorm1d(int(hidden_dim)),
+            nn.GELU(),
+            nn.Linear(int(hidden_dim), self.output_dim),
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+
+class SIGReg(nn.Module):
+    """Sketch Isotropic Gaussian Regularizer, matching the original LeWM design."""
+
+    def __init__(self, knots: int = 17, num_proj: int = 1024) -> None:
+        super().__init__()
+        self.num_proj = int(num_proj)
+        t = torch.linspace(0, 3, int(knots), dtype=torch.float32)
+        dt = 3 / (int(knots) - 1)
+        weights = torch.full((int(knots),), 2 * dt, dtype=torch.float32)
+        weights[[0, -1]] = dt
+        window = torch.exp(-t.square() / 2.0)
+        self.register_buffer("t", t)
+        self.register_buffer("phi", window)
+        self.register_buffer("weights", weights * window)
+
+    def forward(self, proj: torch.Tensor) -> torch.Tensor:
+        """
+        proj: (T, B, D)
+        """
+        A = torch.randn(proj.size(-1), self.num_proj, device=proj.device)
+        A = A.div_(A.norm(p=2, dim=0))
+        x_t = (proj @ A).unsqueeze(-1) * self.t
+        err = (x_t.cos().mean(-3) - self.phi).square() + x_t.sin().mean(-3).square()
+        statistic = (err @ self.weights) * proj.size(-2)
+        return statistic.mean()

roboimi/vla/models/backbones/smolvla_prefix_encoder.py

@@ -1,410 +0,0 @@
-from __future__ import annotations
-
-import warnings
-from typing import Dict, Sequence
-
-import torch
-import torch.nn.functional as F
-from torch import nn
-
-from roboimi.vla.core.interfaces import VLABackbone
-
-try:  # pragma: no cover - exercised by tests via monkeypatch/fakes
-    from transformers import AutoModelForImageTextToText, AutoTokenizer
-except Exception:  # pragma: no cover
-    AutoModelForImageTextToText = None
-    AutoTokenizer = None
-
-
-def _resize_with_pad(img: torch.Tensor, width: int, height: int, pad_value: float = 0.0) -> torch.Tensor:
-    if img.ndim != 4:
-        raise ValueError(f'expected image tensor shaped (B,C,H,W), got {tuple(img.shape)}')
-    cur_height, cur_width = img.shape[2:]
-    ratio = max(cur_width / width, cur_height / height)
-    resized_height = int(cur_height / ratio)
-    resized_width = int(cur_width / ratio)
-    resized = F.interpolate(img, size=(resized_height, resized_width), mode='bilinear', align_corners=False)
-    pad_height = max(0, int(height - resized_height))
-    pad_width = max(0, int(width - resized_width))
-    return F.pad(resized, (pad_width, 0, pad_height, 0), value=pad_value)
-
-
-def _pad_vector(vector: torch.Tensor, new_dim: int) -> torch.Tensor:
-    if vector.shape[-1] == new_dim:
-        return vector
-    if vector.shape[-1] > new_dim:
-        raise ValueError(f'cannot pad vector with dim {vector.shape[-1]} to smaller dim {new_dim}')
-    shape = list(vector.shape)
-    shape[-1] = int(new_dim)
-    padded = torch.zeros(*shape, dtype=vector.dtype, device=vector.device)
-    padded[..., : vector.shape[-1]] = vector
-    return padded
-
-
-class SmolVLAPrefixEncoder(VLABackbone):
-    """SmolVLA-compatible VLM prefix encoder for RoboIMI action experts.
-
-    This module intentionally extracts only the conditioning path from SmolVLA:
-    multiview image tokens, variable language task tokens, and one projected
-    state token.  It freezes the pretrained VLM by default and returns a fixed
-    condition token sequence `(B, S, D)` that can be consumed by existing IMF /
-    transformer-style action heads.
-    """
-
-    def __init__(
-        self,
-        model_name: str = 'HuggingFaceTB/SmolVLM2-500M-Video-Instruct',
-        *,
-        model_name_or_path: str | None = None,
-        vlm: nn.Module | None = None,
-        tokenizer=None,
-        load_vlm_weights: bool = True,
-        local_files_only: bool = False,
-        num_vlm_layers: int = 16,
-        freeze_vlm: bool = True,
-        freeze_vision_encoder: bool = True,
-        train_state_proj: bool = True,
-        max_state_dim: int = 32,
-        resize_imgs_with_padding: Sequence[int] | None = (512, 512),
-        dataset_image_resize_shape: Sequence[int] | None = None,
-        eval_image_resize_shape: Sequence[int] | None = None,
-        tokenizer_max_length: int = 48,
-        pad_language_to: str = 'max_length',
-        run_text_model: bool = True,
-        camera_names: Sequence[str] = ('r_vis', 'top', 'front'),
-        num_cameras: int | None = None,
-    ) -> None:
-        super().__init__()
-        if model_name_or_path is not None:
-            model_name = model_name_or_path
-        self.model_name = str(model_name)
-        self.camera_names = tuple(camera_names)
-        self.num_cameras = int(num_cameras) if num_cameras is not None else len(self.camera_names)
-        if len(self.camera_names) != self.num_cameras:
-            raise ValueError(
-                f'camera_names length ({len(self.camera_names)}) must match num_cameras ({self.num_cameras})'
-            )
-        self.load_vlm_weights = bool(load_vlm_weights)
-        self.local_files_only = bool(local_files_only)
-        self.num_vlm_layers_requested = int(num_vlm_layers)
-        self.freeze_vlm = bool(freeze_vlm)
-        self.freeze_vision_encoder = bool(freeze_vision_encoder)
-        self.max_state_dim = int(max_state_dim)
-        self.resize_imgs_with_padding = self._normalize_resize_shape(resize_imgs_with_padding)
-        self.dataset_image_resize_shape = self._normalize_resize_shape(dataset_image_resize_shape)
-        self.eval_image_resize_shape = self._normalize_resize_shape(eval_image_resize_shape)
-        self.tokenizer_max_length = int(tokenizer_max_length)
-        self.pad_language_to = str(pad_language_to)
-        self.run_text_model = bool(run_text_model)
-        self._warned_trainable_frozen_text_model = False
-
-        if vlm is None:
-            if AutoModelForImageTextToText is None:
-                raise ImportError('transformers AutoModelForImageTextToText is required for SmolVLAPrefixEncoder')
-            if not self.load_vlm_weights:
-                raise ValueError('SmolVLAPrefixEncoder currently requires load_vlm_weights=True')
-            vlm = AutoModelForImageTextToText.from_pretrained(
-                self.model_name,
-                torch_dtype='bfloat16',
-                low_cpu_mem_usage=True,
-                local_files_only=self.local_files_only,
-            )
-        self.vlm = vlm
-
-        if tokenizer is None:
-            if AutoTokenizer is None:
-                raise ImportError('transformers AutoTokenizer is required for SmolVLAPrefixEncoder')
-            tokenizer = AutoTokenizer.from_pretrained(self.model_name, local_files_only=self.local_files_only)
-        self.tokenizer = tokenizer
-
-        if self.num_vlm_layers_requested > 0:
-            text_layers = self.vlm.model.text_model.layers
-            if self.num_vlm_layers_requested > len(text_layers):
-                raise ValueError(
-                    f'num_vlm_layers ({self.num_vlm_layers_requested}) exceeds available text layers '
-                    f'({len(text_layers)})'
-                )
-            self.vlm.model.text_model.layers = text_layers[: self.num_vlm_layers_requested]
-        self.num_vlm_layers = len(self.vlm.model.text_model.layers)
-        text_config = getattr(self.vlm.config, 'text_config', None)
-        if text_config is not None and hasattr(text_config, 'num_hidden_layers'):
-            text_config.num_hidden_layers = self.num_vlm_layers
-
-        hidden_size = int(self.vlm.config.text_config.hidden_size)
-        self._output_dim = hidden_size
-        self.state_proj = nn.Linear(self.max_state_dim, hidden_size)
-        for param in self.state_proj.parameters():
-            param.requires_grad = bool(train_state_proj)
-
-        self.last_prefix_pad_mask: torch.Tensor | None = None
-        self.last_prefix_att_mask: torch.Tensor | None = None
-        self.last_attention_2d_mask: torch.Tensor | None = None
-        self.last_position_ids: torch.Tensor | None = None
-        self._configured_condition_sequence_length = self._infer_configured_condition_sequence_length()
-
-        self.set_requires_grad()
-
-    @staticmethod
-    def _normalize_resize_shape(shape: Sequence[int] | None) -> tuple[int, int] | None:
-        if shape is None:
-            return None
-        normalized = tuple(int(v) for v in shape)
-        if len(normalized) != 2:
-            raise ValueError(f'resize_imgs_with_padding must contain exactly two values, got {normalized}')
-        return normalized
-
-    @property
-    def output_dim(self) -> int:
-        return self._output_dim
-
-    @property
-    def joint_output_dim(self) -> int:
-        return self._output_dim
-
-    def _infer_configured_condition_sequence_length(self) -> int:
-        image_tokens_per_camera = 0
-        config = getattr(self.vlm, 'config', None)
-        vision_config = getattr(config, 'vision_config', None)
-        if vision_config is not None:
-            image_size = int(getattr(vision_config, 'image_size', 0) or 0)
-            patch_size = int(getattr(vision_config, 'patch_size', 0) or 0)
-            scale_factor = int(getattr(config, 'scale_factor', 1) or 1)
-            if image_size > 0 and patch_size > 0 and scale_factor > 0:
-                image_tokens_per_camera = int(((image_size // patch_size) ** 2) / (scale_factor**2))
-        return self.num_cameras * image_tokens_per_camera + self.tokenizer_max_length + 1
-
-    @property
-    def tokens_per_step(self) -> int:
-        if self.last_prefix_pad_mask is not None:
-            return int(self.last_prefix_pad_mask.shape[1])
-        return int(self._configured_condition_sequence_length)
-
-    @property
-    def condition_sequence_length(self) -> int:
-        return self.tokens_per_step
-
-    @staticmethod
-    def _freeze_module(module) -> None:
-        if module is None:
-            return
-        if hasattr(module, 'eval'):
-            module.eval()
-        parameters = getattr(module, 'parameters', None)
-        if callable(parameters):
-            for param in parameters():
-                param.requires_grad = False
-
-    def set_requires_grad(self) -> None:
-        if self.freeze_vision_encoder:
-            self._freeze_module(self.vlm.model.vision_model)
-        if self.freeze_vlm:
-            self._freeze_module(self.vlm)
-            self._freeze_module(getattr(self.vlm.model, 'vision_model', None))
-            self._freeze_module(getattr(self.vlm.model, 'connector', None))
-            self._freeze_module(getattr(self.vlm.model, 'text_model', None))
-
-    def train(self, mode: bool = True):
-        super().train(mode)
-        if self.freeze_vlm:
-            self.vlm.eval()
-        elif self.freeze_vision_encoder:
-            self.vlm.model.vision_model.eval()
-        return self
-
-    def _ordered_camera_names(self, images: Dict[str, torch.Tensor]) -> tuple[str, ...]:
-        missing = [name for name in self.camera_names if name not in images]
-        if missing:
-            raise ValueError(f'image input missing required cameras. missing={missing}, expected={list(self.camera_names)}')
-        return self.camera_names
-
-    def _batch_size_from_images(self, images: Dict[str, torch.Tensor]) -> int:
-        return int(next(iter(images.values())).shape[0])
-
-    def _normalize_tasks(self, task, batch_size: int) -> list[str]:
-        if task is None:
-            tasks = [''] * batch_size
-        elif isinstance(task, str):
-            tasks = [task] * batch_size
-        elif isinstance(task, tuple):
-            tasks = list(task)
-        elif isinstance(task, list):
-            tasks = task
-        else:
-            raise TypeError(f'task must be str/list/tuple/None, got {type(task)!r}')
-        if len(tasks) != batch_size:
-            raise ValueError(f'task batch size ({len(tasks)}) must match image batch size ({batch_size})')
-        return [item if item.endswith('\n') else f'{item}\n' for item in tasks]
-
-    def tokenize_task(self, task, batch_size: int, device: torch.device) -> tuple[torch.Tensor, torch.Tensor]:
-        tasks = self._normalize_tasks(task, batch_size)
-        old_padding_side = getattr(self.tokenizer, 'padding_side', None)
-        if old_padding_side is not None:
-            self.tokenizer.padding_side = 'right'
-        try:
-            tokenized = self.tokenizer(
-                tasks,
-                padding=self.pad_language_to,
-                max_length=self.tokenizer_max_length,
-                return_tensors='pt',
-                truncation=True,
-            )
-        finally:
-            if old_padding_side is not None:
-                self.tokenizer.padding_side = old_padding_side
-        tokens = tokenized['input_ids'].to(device=device)
-        masks = tokenized['attention_mask'].to(device=device, dtype=torch.bool)
-        return tokens, masks
-
-    def prepare_images(self, images: Dict[str, torch.Tensor]) -> tuple[list[torch.Tensor], list[torch.Tensor]]:
-        camera_names = self._ordered_camera_names(images)
-        reference = images[camera_names[0]]
-        if reference.ndim != 5:
-            raise ValueError(f'expected image tensor shaped (B,T,C,H,W), got {tuple(reference.shape)}')
-        batch_size = reference.shape[0]
-        prepared: list[torch.Tensor] = []
-        masks: list[torch.Tensor] = []
-        for camera_name in camera_names:
-            image = images[camera_name]
-            if image.shape[:2] != reference.shape[:2] or image.shape[2] != reference.shape[2]:
-                raise ValueError(f'camera {camera_name!r} shape {tuple(image.shape)} does not match reference {tuple(reference.shape)}')
-            image = image[:, -1].contiguous().float().clamp(0.0, 1.0)
-            if self.resize_imgs_with_padding is not None:
-                image = _resize_with_pad(image, *self.resize_imgs_with_padding, pad_value=0.0)
-            image = image * 2.0 - 1.0
-            prepared.append(image)
-            masks.append(torch.ones(batch_size, dtype=torch.bool, device=image.device))
-        return prepared, masks
-
-    def prepare_state(self, state: torch.Tensor) -> torch.Tensor:
-        if state.ndim > 2:
-            state = state[:, -1, :]
-        return _pad_vector(state.float(), self.max_state_dim)
-
-    def embed_image(self, image: torch.Tensor) -> torch.Tensor:
-        if hasattr(self.vlm.model, 'get_image_features'):
-            pixel_values = image[:, None, ...]
-            pixel_attention_mask = torch.ones(
-                image.shape[0],
-                1,
-                image.shape[2],
-                image.shape[3],
-                dtype=torch.bool,
-                device=image.device,
-            )
-            with torch.set_grad_enabled(
-                torch.is_grad_enabled() and not self.freeze_vlm and not self.freeze_vision_encoder
-            ):
-                hidden = self.vlm.model.get_image_features(
-                    pixel_values=pixel_values,
-                    pixel_attention_mask=pixel_attention_mask,
-                    return_dict=True,
-                ).pooler_output
-        else:
-            vision_model = self.vlm.model.vision_model
-            with torch.set_grad_enabled(
-                torch.is_grad_enabled() and not self.freeze_vlm and not self.freeze_vision_encoder
-            ):
-                patch_attention_mask = torch.ones(
-                    image.shape[0],
-                    image.shape[2] // self.vlm.config.vision_config.patch_size,
-                    image.shape[3] // self.vlm.config.vision_config.patch_size,
-                    dtype=torch.bool,
-                    device=image.device,
-                ) if hasattr(self.vlm.config, 'vision_config') and hasattr(self.vlm.config.vision_config, 'patch_size') else None
-                hidden = vision_model(
-                    pixel_values=image.to(dtype=vision_model.dtype),
-                    patch_attention_mask=patch_attention_mask,
-                ).last_hidden_state
-                hidden = self.vlm.model.connector(hidden)
-        return hidden
-
-    def embed_language_tokens(self, tokens: torch.Tensor) -> torch.Tensor:
-        return self.vlm.model.text_model.get_input_embeddings()(tokens)
-
-    @staticmethod
-    def make_att_2d_masks(pad_masks: torch.Tensor, att_masks: torch.Tensor) -> torch.Tensor:
-        cumsum = torch.cumsum(att_masks, dim=1)
-        att_2d = cumsum[:, None, :] <= cumsum[:, :, None]
-        pad_2d = pad_masks[:, None, :] * pad_masks[:, :, None]
-        return att_2d & pad_2d
-
-    def embed_prefix(self, images: Dict[str, torch.Tensor], state: torch.Tensor, task=None) -> torch.Tensor:
-        prepared_images, image_masks = self.prepare_images(images)
-        batch_size = prepared_images[0].shape[0]
-        device = prepared_images[0].device
-
-        embs: list[torch.Tensor] = []
-        pad_masks: list[torch.Tensor] = []
-        att_masks: list[int] = []
-
-        for image, image_mask in zip(prepared_images, image_masks, strict=False):
-            img_emb = self.embed_image(image)
-            img_emb = img_emb * torch.tensor(
-                img_emb.shape[-1] ** 0.5,
-                dtype=img_emb.dtype,
-                device=img_emb.device,
-            )
-            num_img_tokens = img_emb.shape[1]
-            embs.append(img_emb)
-            pad_masks.append(image_mask[:, None].expand(batch_size, num_img_tokens))
-            att_masks += [0] * num_img_tokens
-
-        lang_tokens, lang_masks = self.tokenize_task(task, batch_size=batch_size, device=device)
-        lang_emb = self.embed_language_tokens(lang_tokens)
-        lang_emb = lang_emb * (lang_emb.shape[-1] ** 0.5)
-        embs.append(lang_emb)
-        pad_masks.append(lang_masks)
-        att_masks += [0] * lang_emb.shape[1]
-
-        state = self.prepare_state(state).to(device=device)
-        state_emb = self.state_proj(state).unsqueeze(1)
-        embs.append(state_emb)
-        pad_masks.append(torch.ones(batch_size, 1, dtype=torch.bool, device=device))
-        att_masks += [1]
-
-        prefix = torch.cat(embs, dim=1)
-        prefix_pad_mask = torch.cat(pad_masks, dim=1)
-        prefix_att_mask = torch.tensor(att_masks, dtype=torch.bool, device=device)[None, :].expand(batch_size, -1)
-
-        self.last_prefix_pad_mask = prefix_pad_mask
-        self.last_prefix_att_mask = prefix_att_mask
-        self.last_attention_2d_mask = self.make_att_2d_masks(prefix_pad_mask, prefix_att_mask)
-        self.last_position_ids = torch.cumsum(prefix_pad_mask, dim=1) - 1
-        return prefix
-
-    def encode_prefix(self, prefix: torch.Tensor) -> torch.Tensor:
-        if not self.run_text_model:
-            return prefix
-        if self.last_attention_2d_mask is None or self.last_position_ids is None:
-            raise RuntimeError('encode_prefix requires masks from embed_prefix')
-        text_model = self.vlm.model.text_model
-        text_dtype = getattr(text_model, 'dtype', prefix.dtype)
-        if (
-            self.freeze_vlm
-            and not self._warned_trainable_frozen_text_model
-            and any(param.requires_grad for param in text_model.parameters())
-        ):
-            warnings.warn(
-                'freeze_vlm=True but text_model has trainable parameters; temporarily disabling '
-                'their gradients while keeping gradients for trainable prefix inputs.',
-                RuntimeWarning,
-            )
-            self._warned_trainable_frozen_text_model = True
-        with torch.set_grad_enabled(torch.is_grad_enabled()):
-            outputs = text_model(
-                inputs_embeds=prefix.to(dtype=text_dtype),
-                attention_mask=self.last_attention_2d_mask[:, None, :, :],
-                position_ids=self.last_position_ids,
-                use_cache=False,
-                return_dict=True,
-            )
-        return outputs.last_hidden_state
-
-    def forward(self, images: Dict[str, torch.Tensor], state: torch.Tensor | None = None, task=None) -> torch.Tensor:
-        if state is None:
-            raise ValueError('SmolVLAPrefixEncoder.forward requires `state` for SmolVLA-compatible state token')
-        prefix = self.embed_prefix(images=images, state=state, task=task)
-        return self.encode_prefix(prefix)
-
-
-SmolVLMPrefixEncoder = SmolVLAPrefixEncoder

roboimi/vla/models/smolvla/__init__.py

@@ -1,13 +0,0 @@
-"""Native SmolVLA model core."""
-
-from .configuration import NativeSmolVLAConfig, SmolVLAConfig
-from .modeling import VLAFlowMatching, make_att_2d_masks, pad_vector, resize_with_pad
-
-__all__ = [
-    "NativeSmolVLAConfig",
-    "SmolVLAConfig",
-    "VLAFlowMatching",
-    "make_att_2d_masks",
-    "pad_vector",
-    "resize_with_pad",
-]

roboimi/vla/models/smolvla/configuration.py

@@ -1,75 +0,0 @@
-"""Native SmolVLA configuration, independent of LeRobot."""
-
-from __future__ import annotations
-
-from dataclasses import dataclass
-
-
-@dataclass
-class NativeSmolVLAConfig:
-    """Lightweight configuration for the native SmolVLA model core.
-
-    This intentionally keeps only model-core fields needed by
-    :class:`VLAFlowMatching`; policy/dataset/optimizer concerns stay outside of
-    the native model package.
-    """
-
-    # Input / output structure.
-    n_obs_steps: int = 1
-    chunk_size: int = 50
-    n_action_steps: int = 50
-
-    # Shorter state and action vectors are padded before entering the model.
-    max_state_dim: int = 32
-    max_action_dim: int = 32
-
-    # Image preprocessing.
-    resize_imgs_with_padding: tuple[int, int] = (512, 512)
-
-    # Tokenizer / decoding.
-    tokenizer_max_length: int = 48
-    num_steps: int = 10
-
-    # Attention utils.
-    use_cache: bool = True
-
-    # Finetuning settings.
-    freeze_vision_encoder: bool = True
-    train_expert_only: bool = True
-    train_state_proj: bool = True
-
-    # VLM / expert construction settings.
-    vlm_model_name: str = "HuggingFaceTB/SmolVLM2-500M-Video-Instruct"
-    load_vlm_weights: bool = False
-    add_image_special_tokens: bool = False
-    attention_mode: str = "cross_attn"
-    prefix_length: int = -1
-    pad_language_to: str = "longest"
-    num_expert_layers: int = -1
-    num_vlm_layers: int = 16
-    self_attn_every_n_layers: int = 2
-    expert_width_multiplier: float = 0.75
-
-    # Flow-matching timestep embedding.
-    min_period: float = 4e-3
-    max_period: float = 4.0
-
-    # Runtime settings.
-    device: str | None = None
-    compile_model: bool = False
-    compile_mode: str = "max-autotune"
-
-    # RTC hook placeholder; the native core does not implement RTC, but keeping
-    # the field lets migrated call sites pass configs through unchanged.
-    rtc_config: object | None = None
-
-    def __post_init__(self) -> None:
-        if self.n_action_steps > self.chunk_size:
-            raise ValueError(
-                "The chunk size is the upper bound for the number of action steps per model invocation. "
-                f"Got {self.n_action_steps} for `n_action_steps` and {self.chunk_size} for `chunk_size`."
-            )
-
-
-# Compatibility alias for migrated code that still imports SmolVLAConfig.
-SmolVLAConfig = NativeSmolVLAConfig

roboimi/vla/models/smolvla/modeling.py

@@ -1,421 +0,0 @@
-"""Native SmolVLA flow-matching core.
-
-This module ports the lightweight helpers and ``VLAFlowMatching`` from the
-LeRobot SmolVLA implementation while avoiding any LeRobot imports.  The heavy
-Transformers-backed VLM/expert is optional and can be injected for tests.
-"""
-
-from __future__ import annotations
-
-import math
-from typing import TypedDict
-
-try:
-    from typing import Unpack
-except ImportError:  # Python 3.10
-    from typing_extensions import Unpack
-
-import torch
-import torch.nn.functional as F
-from torch import Tensor, nn
-
-from .configuration import NativeSmolVLAConfig
-
-
-class ActionSelectKwargs(TypedDict, total=False):
-    inference_delay: int | None
-    prev_chunk_left_over: Tensor | None
-    execution_horizon: int | None
-
-
-def create_sinusoidal_pos_embedding(
-    time: torch.Tensor,
-    dimension: int,
-    min_period: float,
-    max_period: float,
-    device: torch.device | str = "cpu",
-) -> Tensor:
-    """Compute sine/cosine positional embeddings for scalar positions."""
-    if dimension % 2 != 0:
-        raise ValueError(f"dimension ({dimension}) must be divisible by 2")
-    if time.ndim != 1:
-        raise ValueError("The time tensor is expected to be of shape `(batch_size, )`.")
-
-    fraction = torch.linspace(0.0, 1.0, dimension // 2, dtype=torch.float64, device=device)
-    period = min_period * (max_period / min_period) ** fraction
-    scaling_factor = 1.0 / period * 2 * math.pi
-    sin_input = scaling_factor[None, :] * time[:, None].to(dtype=torch.float64)
-    pos_emb = torch.cat([torch.sin(sin_input), torch.cos(sin_input)], dim=1)
-    return pos_emb
-
-
-def make_att_2d_masks(pad_masks: Tensor, att_masks: Tensor) -> Tensor:
-    """Build Big Vision-style 2D prefix-LM attention masks.
-
-    ``pad_masks`` is ``bool[B, N]`` and marks valid tokens. ``att_masks`` is
-    ``bool/int[B, N]`` where cumulative increments begin new causal groups.  A
-    query token can attend to valid key tokens whose cumulative attention group
-    is less than or equal to the query's group.
-    """
-    if att_masks.ndim != 2:
-        raise ValueError(att_masks.ndim)
-    if pad_masks.ndim != 2:
-        raise ValueError(pad_masks.ndim)
-
-    cumsum = torch.cumsum(att_masks.to(dtype=torch.long), dim=1)
-    att_2d_masks = cumsum[:, None, :] <= cumsum[:, :, None]
-    pad_2d_masks = pad_masks[:, None, :].bool() & pad_masks[:, :, None].bool()
-    return att_2d_masks & pad_2d_masks
-
-
-def resize_with_pad(img: Tensor, width: int, height: int, pad_value: float = -1) -> Tensor:
-    """Resize a BCHW image batch preserving aspect ratio, then top/left pad."""
-    if img.ndim != 4:
-        raise ValueError(f"(b,c,h,w) expected, but {img.shape}")
-
-    cur_height, cur_width = img.shape[2:]
-    ratio = max(cur_width / width, cur_height / height)
-    resized_height = int(cur_height / ratio)
-    resized_width = int(cur_width / ratio)
-    resized_img = F.interpolate(
-        img,
-        size=(resized_height, resized_width),
-        mode="bilinear",
-        align_corners=False,
-    )
-
-    pad_height = max(0, int(height - resized_height))
-    pad_width = max(0, int(width - resized_width))
-    return F.pad(resized_img, (pad_width, 0, pad_height, 0), value=pad_value)
-
-
-def pad_vector(vector: Tensor, new_dim: int) -> Tensor:
-    """Pad a vector-like tensor's last dimension with zeros to ``new_dim``."""
-    current_dim = vector.shape[-1]
-    if current_dim == new_dim:
-        return vector
-    if current_dim > new_dim:
-        raise ValueError(f"Cannot pad vector with current dimension {current_dim} to smaller target dimension {new_dim}.")
-
-    shape = list(vector.shape)
-    shape[-1] = new_dim
-    new_vector = torch.zeros(*shape, dtype=vector.dtype, device=vector.device)
-    new_vector[..., :current_dim] = vector
-    return new_vector
-
-
-def pad_tensor(tensor: Tensor, max_len: int, pad_value: int | float) -> Tensor:
-    """Pad a tensor along sequence dimension to ``max_len``."""
-    bsize, seq_len = tensor.shape[:2]
-    if seq_len >= max_len:
-        return tensor
-    padded_tensor = torch.full(
-        (bsize, max_len, *tensor.shape[2:]),
-        pad_value,
-        dtype=tensor.dtype,
-        device=tensor.device,
-    )
-    padded_tensor[:, :seq_len] = tensor
-    return padded_tensor
-
-
-class VLAFlowMatching(nn.Module):
-    """SmolVLA flow-matching action head around a VLM plus action expert."""
-
-    def __init__(
-        self,
-        config: NativeSmolVLAConfig,
-        rtc_processor: object | None = None,
-        vlm_with_expert: nn.Module | None = None,
-    ):
-        super().__init__()
-        self.config = config
-
-        if vlm_with_expert is None:
-            from .smolvlm_with_expert import SmolVLMWithExpertModel
-
-            vlm_with_expert = SmolVLMWithExpertModel(
-                model_id=self.config.vlm_model_name,
-                freeze_vision_encoder=self.config.freeze_vision_encoder,
-                train_expert_only=self.config.train_expert_only,
-                load_vlm_weights=self.config.load_vlm_weights,
-                attention_mode=self.config.attention_mode,
-                num_expert_layers=self.config.num_expert_layers,
-                num_vlm_layers=self.config.num_vlm_layers,
-                self_attn_every_n_layers=self.config.self_attn_every_n_layers,
-                expert_width_multiplier=self.config.expert_width_multiplier,
-                device=self.config.device if self.config.device is not None else "auto",
-            )
-        self.vlm_with_expert = vlm_with_expert
-
-        vlm_hidden_size = self.vlm_with_expert.config.text_config.hidden_size
-        expert_hidden_size = self.vlm_with_expert.expert_hidden_size
-        self.state_proj = nn.Linear(self.config.max_state_dim, vlm_hidden_size)
-        self.action_in_proj = nn.Linear(self.config.max_action_dim, expert_hidden_size)
-        self.action_out_proj = nn.Linear(expert_hidden_size, self.config.max_action_dim)
-        self.action_time_mlp_in = nn.Linear(expert_hidden_size * 2, expert_hidden_size)
-        self.action_time_mlp_out = nn.Linear(expert_hidden_size, expert_hidden_size)
-
-        self.set_requires_grad()
-        tokenizer = self.vlm_with_expert.processor.tokenizer
-        self.fake_image_token = tokenizer.fake_image_token_id
-        self.global_image_token = tokenizer.global_image_token_id
-        self.global_image_start_token = torch.tensor(
-            [self.fake_image_token, self.global_image_token], dtype=torch.long
-        )
-        self.add_image_special_tokens = self.config.add_image_special_tokens
-        self.image_end_token = torch.tensor([self.fake_image_token], dtype=torch.long)
-        self.prefix_length = self.config.prefix_length
-        self.rtc_processor = rtc_processor
-
-        if config.compile_model:
-            torch.set_float32_matmul_precision("high")
-            self.sample_actions = torch.compile(self.sample_actions, mode=config.compile_mode)
-            self.forward = torch.compile(self.forward, mode=config.compile_mode)
-
-    def _rtc_enabled(self) -> bool:
-        return bool(self.config.rtc_config is not None and getattr(self.config.rtc_config, "enabled", False))
-
-    def set_requires_grad(self) -> None:
-        for params in self.state_proj.parameters():
-            params.requires_grad = self.config.train_state_proj
-
-    def sample_noise(self, shape: tuple[int, ...] | torch.Size, device: torch.device | str) -> Tensor:
-        return torch.normal(mean=0.0, std=1.0, size=shape, dtype=torch.float32, device=device)
-
-    def sample_time(self, bsize: int, device: torch.device | str) -> Tensor:
-        beta_dist = torch.distributions.Beta(concentration1=1.5, concentration0=1.0)
-        time_beta = beta_dist.sample((bsize,)).to(device=device, dtype=torch.float32)
-        return time_beta * 0.999 + 0.001
-
-    def _vlm_device(self) -> torch.device:
-        vlm = getattr(self.vlm_with_expert, "vlm", None)
-        return getattr(vlm, "device", next(self.parameters()).device)
-
-    def embed_prefix(
-        self,
-        images: list[Tensor],
-        img_masks: list[Tensor],
-        lang_tokens: Tensor,
-        lang_masks: Tensor,
-        state: Tensor,
-    ) -> tuple[Tensor, Tensor, Tensor]:
-        """Embed images, language tokens, and robot state as prefix tokens."""
-        embs: list[Tensor] = []
-        pad_masks: list[Tensor] = []
-        att_masks: list[int] = []
-
-        for img, img_mask in zip(images, img_masks, strict=False):
-            if self.add_image_special_tokens:
-                image_start_token = (
-                    self.vlm_with_expert.embed_language_tokens(self.global_image_start_token.to(device=self._vlm_device()))
-                    .unsqueeze(0)
-                    .expand(img.shape[0], -1, -1)
-                )
-                image_start_mask = torch.ones_like(image_start_token[:, :, 0], dtype=torch.bool)
-                embs.append(image_start_token)
-                pad_masks.append(image_start_mask)
-                att_masks += [0] * image_start_mask.shape[-1]
-
-            img_emb = self.vlm_with_expert.embed_image(img)
-            img_emb = img_emb * torch.tensor(img_emb.shape[-1] ** 0.5, dtype=img_emb.dtype, device=img_emb.device)
-            bsize, num_img_embs = img_emb.shape[:2]
-            img_mask = img_mask.to(device=img_emb.device, dtype=torch.bool)[:, None].expand(bsize, num_img_embs)
-            embs.append(img_emb)
-            pad_masks.append(img_mask)
-            att_masks += [0] * num_img_embs
-
-            if self.add_image_special_tokens:
-                image_end_token = (
-                    self.vlm_with_expert.embed_language_tokens(self.image_end_token.to(device=self._vlm_device()))
-                    .unsqueeze(0)
-                    .expand(img.shape[0], -1, -1)
-                )
-                image_end_mask = torch.ones_like(image_end_token[:, :, 0], dtype=torch.bool)
-                embs.append(image_end_token)
-                pad_masks.append(image_end_mask)
-                att_masks += [0] * image_end_mask.shape[1]
-
-        lang_emb = self.vlm_with_expert.embed_language_tokens(lang_tokens)
-        lang_emb = lang_emb * math.sqrt(lang_emb.shape[-1])
-        embs.append(lang_emb)
-        pad_masks.append(lang_masks.to(device=lang_emb.device, dtype=torch.bool))
-        att_masks += [0] * lang_emb.shape[1]
-
-        state_emb = self.state_proj(state)
-        state_emb = state_emb[:, None, :] if state_emb.ndim == 2 else state_emb
-        embs.append(state_emb)
-        state_mask = torch.ones(state_emb.shape[:2], dtype=torch.bool, device=state_emb.device)
-        pad_masks.append(state_mask)
-        att_masks += [1] * state_emb.shape[1]
-
-        all_embs = torch.cat(embs, dim=1)
-        all_pad_masks = torch.cat(pad_masks, dim=1)
-        all_att_masks = torch.tensor(att_masks, dtype=torch.bool, device=all_pad_masks.device)[None, :]
-
-        if self.prefix_length > 0 and all_pad_masks.shape[1] < self.prefix_length:
-            all_embs = pad_tensor(all_embs, self.prefix_length, pad_value=0)
-            all_pad_masks = pad_tensor(all_pad_masks, self.prefix_length, pad_value=0)
-            all_att_masks = pad_tensor(all_att_masks, self.prefix_length, pad_value=0)
-
-        all_att_masks = all_att_masks.expand(all_pad_masks.shape[0], -1)
-        return all_embs, all_pad_masks, all_att_masks
-
-    def embed_suffix(self, noisy_actions: Tensor, timestep: Tensor) -> tuple[Tensor, Tensor, Tensor]:
-        """Embed noisy action tokens and timestep for the expert suffix."""
-        action_emb = self.action_in_proj(noisy_actions)
-        device = action_emb.device
-        bsize = action_emb.shape[0]
-        dtype = action_emb.dtype
-
-        time_emb = create_sinusoidal_pos_embedding(
-            timestep,
-            self.vlm_with_expert.expert_hidden_size,
-            self.config.min_period,
-            self.config.max_period,
-            device=device,
-        ).to(dtype=dtype)
-        time_emb = time_emb[:, None, :].expand_as(action_emb)
-        action_time_emb = torch.cat([action_emb, time_emb], dim=2)
-        action_time_emb = self.action_time_mlp_in(action_time_emb)
-        action_time_emb = F.silu(action_time_emb)
-        action_time_emb = self.action_time_mlp_out(action_time_emb)
-
-        action_time_mask = torch.ones(action_time_emb.shape[:2], dtype=torch.bool, device=device)
-        att_masks = torch.ones(bsize, self.config.chunk_size, dtype=torch.bool, device=device)
-        return action_time_emb, action_time_mask, att_masks
-
-    def forward(
-        self,
-        images: list[Tensor],
-        img_masks: list[Tensor],
-        lang_tokens: Tensor,
-        lang_masks: Tensor,
-        state: Tensor,
-        actions: Tensor,
-        noise: Tensor | None = None,
-        time: Tensor | None = None,
-    ) -> Tensor:
-        """Run a training forward pass and return per-element flow loss."""
-        if noise is None:
-            noise = self.sample_noise(actions.shape, actions.device)
-        if time is None:
-            time = self.sample_time(actions.shape[0], actions.device)
-
-        time_expanded = time[:, None, None]
-        x_t = time_expanded * noise + (1 - time_expanded) * actions
-        u_t = noise - actions
-
-        prefix_embs, prefix_pad_masks, prefix_att_masks = self.embed_prefix(
-            images, img_masks, lang_tokens, lang_masks, state=state
-        )
-        suffix_embs, suffix_pad_masks, suffix_att_masks = self.embed_suffix(x_t, time)
-
-        pad_masks = torch.cat([prefix_pad_masks, suffix_pad_masks], dim=1)
-        att_masks = torch.cat([prefix_att_masks, suffix_att_masks], dim=1)
-        att_2d_masks = make_att_2d_masks(pad_masks, att_masks)
-        position_ids = torch.cumsum(pad_masks, dim=1) - 1
-
-        (_, suffix_out), _ = self.vlm_with_expert.forward(
-            attention_mask=att_2d_masks,
-            position_ids=position_ids,
-            past_key_values=None,
-            inputs_embeds=[prefix_embs, suffix_embs],
-            use_cache=False,
-            fill_kv_cache=False,
-        )
-        suffix_out = suffix_out[:, -self.config.chunk_size :].to(dtype=torch.float32)
-        v_t = self.action_out_proj(suffix_out)
-        return F.mse_loss(u_t, v_t, reduction="none")
-
-    def sample_actions(
-        self,
-        images: list[Tensor],
-        img_masks: list[Tensor],
-        lang_tokens: Tensor,
-        lang_masks: Tensor,
-        state: Tensor,
-        noise: Tensor | None = None,
-        **kwargs: Unpack[ActionSelectKwargs],
-    ) -> Tensor:
-        """Sample an action chunk with Euler integration over the flow field."""
-        bsize = state.shape[0]
-        device = state.device
-        if noise is None:
-            actions_shape = (bsize, self.config.chunk_size, self.config.max_action_dim)
-            noise = self.sample_noise(actions_shape, device)
-
-        prefix_embs, prefix_pad_masks, prefix_att_masks = self.embed_prefix(
-            images, img_masks, lang_tokens, lang_masks, state=state
-        )
-        prefix_att_2d_masks = make_att_2d_masks(prefix_pad_masks, prefix_att_masks)
-        prefix_position_ids = torch.cumsum(prefix_pad_masks, dim=1) - 1
-        _, past_key_values = self.vlm_with_expert.forward(
-            attention_mask=prefix_att_2d_masks,
-            position_ids=prefix_position_ids,
-            past_key_values=None,
-            inputs_embeds=[prefix_embs, None],
-            use_cache=self.config.use_cache,
-            fill_kv_cache=True,
-        )
-
-        dt = -1.0 / self.config.num_steps
-        x_t = noise
-        for step in range(self.config.num_steps):
-            time = 1.0 + step * dt
-            time_tensor = torch.tensor(time, dtype=torch.float32, device=device).expand(bsize)
-
-            def denoise_step_partial_call(input_x_t: Tensor, current_timestep: Tensor = time_tensor) -> Tensor:
-                return self.denoise_step(
-                    x_t=input_x_t,
-                    prefix_pad_masks=prefix_pad_masks,
-                    past_key_values=past_key_values,
-                    timestep=current_timestep,
-                )
-
-            if self._rtc_enabled() and self.rtc_processor is not None:
-                v_t = self.rtc_processor.denoise_step(
-                    x_t=x_t,
-                    prev_chunk_left_over=kwargs.get("prev_chunk_left_over"),
-                    inference_delay=kwargs.get("inference_delay"),
-                    time=time,
-                    original_denoise_step_partial=denoise_step_partial_call,
-                    execution_horizon=kwargs.get("execution_horizon"),
-                )
-            else:
-                v_t = denoise_step_partial_call(x_t)
-            x_t = x_t + dt * v_t
-
-            if self.rtc_processor is not None and getattr(self.rtc_processor, "is_debug_enabled", lambda: False)():
-                self.rtc_processor.track(time=time, x_t=x_t, v_t=v_t)
-
-        return x_t
-
-    def denoise_step(
-        self,
-        prefix_pad_masks: Tensor,
-        past_key_values: object,
-        x_t: Tensor,
-        timestep: Tensor,
-    ) -> Tensor:
-        """Apply one denoising step at a given timestep."""
-        suffix_embs, suffix_pad_masks, suffix_att_masks = self.embed_suffix(x_t, timestep)
-        suffix_len = suffix_pad_masks.shape[1]
-        batch_size = prefix_pad_masks.shape[0]
-        prefix_len = prefix_pad_masks.shape[1]
-        prefix_pad_2d_masks = prefix_pad_masks[:, None, :].expand(batch_size, suffix_len, prefix_len)
-        suffix_att_2d_masks = make_att_2d_masks(suffix_pad_masks, suffix_att_masks)
-        full_att_2d_masks = torch.cat([prefix_pad_2d_masks, suffix_att_2d_masks], dim=2)
-        prefix_offsets = torch.sum(prefix_pad_masks, dim=-1)[:, None]
-        position_ids = prefix_offsets + torch.cumsum(suffix_pad_masks, dim=1) - 1
-
-        outputs_embeds, _ = self.vlm_with_expert.forward(
-            attention_mask=full_att_2d_masks,
-            position_ids=position_ids,
-            past_key_values=past_key_values,
-            inputs_embeds=[None, suffix_embs],
-            use_cache=self.config.use_cache,
-            fill_kv_cache=False,
-        )
-        suffix_out = outputs_embeds[1][:, -self.config.chunk_size :].to(dtype=torch.float32)
-        return self.action_out_proj(suffix_out)

roboimi/vla/models/smolvla/smolvlm_with_expert.py

@@ -1,571 +0,0 @@
-# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import copy
-from typing import TYPE_CHECKING
-
-import torch
-from torch import nn
-
-if TYPE_CHECKING:
-    from transformers import (
-        AutoConfig,
-        AutoModel,
-        AutoModelForImageTextToText,
-        AutoProcessor,
-        SmolVLMForConditionalGeneration,
-    )
-
-
-def _require_transformers():
-    try:
-        from transformers import (
-            AutoConfig,
-            AutoModel,
-            AutoModelForImageTextToText,
-            AutoProcessor,
-            SmolVLMForConditionalGeneration,
-        )
-    except ImportError as exc:
-        raise ImportError(
-            "Native SmolVLA requires the optional `transformers` package to construct "
-            "SmolVLMWithExpertModel. Install transformers or inject `vlm_with_expert` "
-            "when constructing VLAFlowMatching."
-        ) from exc
-    return AutoConfig, AutoModel, AutoModelForImageTextToText, AutoProcessor, SmolVLMForConditionalGeneration
-
-
-def apply_rope(x, positions, max_wavelength=10_000):
-    """
-    Applies RoPE positions [B, L] to x [B, L, H, D].
-    """
-    d_half = x.shape[-1] // 2
-    device = x.device
-    dtype = x.dtype
-    x = x.to(torch.float32)
-
-    freq_exponents = (2.0 / x.shape[-1]) * torch.arange(d_half, dtype=torch.float32, device=device)
-    timescale = max_wavelength**freq_exponents
-    radians = positions[..., None].to(torch.float32) / timescale[None, None, :].to(torch.float32)
-
-    radians = radians[..., None, :]
-
-    sin = torch.sin(radians)  # .to(dtype=dtype)
-    cos = torch.cos(radians)  # .to(dtype=dtype)
-
-    x1, x2 = x.split(d_half, dim=-1)
-    res = torch.empty_like(x)
-    res[..., :d_half] = x1 * cos - x2 * sin
-    res[..., d_half:] = x2 * cos + x1 * sin
-
-    return res.to(dtype)
-
-
-def get_intermediate_size(hidden_dim, ffn_dim_multiplier=4, multiple_of=256):
-    hidden_dim = int(2 * hidden_dim / 3)
-    hidden_dim = int(ffn_dim_multiplier * hidden_dim)
-    hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)
-    return hidden_dim
-
-
-class SmolVLMWithExpertModel(nn.Module):
-    def __init__(
-        self,
-        model_id: str = "HuggingFaceTB/SmolVLM2-500M-Video-Instruct",
-        load_vlm_weights: bool = True,
-        train_expert_only: bool = True,
-        freeze_vision_encoder: bool = False,
-        attention_mode: str = "self_attn",
-        num_expert_layers: int = -1,
-        num_vlm_layers: int = -1,
-        self_attn_every_n_layers: int = -1,
-        expert_width_multiplier: float = 0.5,
-        device: str = "auto",
-    ):
-        super().__init__()
-        AutoConfig, AutoModel, AutoModelForImageTextToText, AutoProcessor, SmolVLMForConditionalGeneration = _require_transformers()
-        if load_vlm_weights:
-            print(f"Loading  {model_id} weights ...")
-            self.vlm = AutoModelForImageTextToText.from_pretrained(
-                model_id,
-                torch_dtype="bfloat16",
-                low_cpu_mem_usage=True,
-            )
-            config = self.vlm.config
-        else:
-            config = AutoConfig.from_pretrained(model_id)
-            self.vlm = SmolVLMForConditionalGeneration(config=config)
-        self.processor = AutoProcessor.from_pretrained(model_id)
-        if num_vlm_layers > 0:
-            print(f"Reducing the number of VLM layers to {num_vlm_layers} ...")
-            self.get_vlm_model().text_model.layers = self.get_vlm_model().text_model.layers[:num_vlm_layers]
-        self.num_vlm_layers = len(self.get_vlm_model().text_model.layers)
-        self.config = config
-        # Smaller lm expert
-        lm_expert_config = copy.deepcopy(config.text_config)
-        hidden_size = lm_expert_config.hidden_size
-        lm_expert_config.hidden_size = int(hidden_size * expert_width_multiplier)  # hidden_size // 2
-        lm_expert_config.intermediate_size = get_intermediate_size(int(hidden_size * expert_width_multiplier))
-        lm_expert_config.num_hidden_layers = self.num_vlm_layers
-        if num_expert_layers > 0:
-            assert len(self.get_vlm_model().text_model.layers) % num_expert_layers == 0, (
-                f"Number of layers in the VLM {len(self.get_vlm_model().text_model.layers)} are not multiple of num_expert_layers {num_expert_layers}"
-            )
-            lm_expert_config.num_hidden_layers = num_expert_layers
-        self.lm_expert = AutoModel.from_config(lm_expert_config)
-
-        self.num_expert_layers = len(self.lm_expert.layers)
-        self.self_attn_every_n_layers = self_attn_every_n_layers
-        if "cross" in attention_mode:
-            # Reshape qkv projections to have the same input dimension as the vlm
-            for layer_idx in range(len(self.lm_expert.layers)):
-                if self.self_attn_every_n_layers > 0 and layer_idx % self.self_attn_every_n_layers == 0:
-                    continue
-                self.lm_expert.layers[layer_idx].self_attn.k_proj = nn.Linear(
-                    config.text_config.num_key_value_heads * config.text_config.head_dim,
-                    lm_expert_config.num_key_value_heads * lm_expert_config.head_dim,
-                    bias=lm_expert_config.attention_bias,
-                )
-                self.lm_expert.layers[layer_idx].self_attn.v_proj = nn.Linear(
-                    config.text_config.num_key_value_heads * config.text_config.head_dim,
-                    lm_expert_config.num_key_value_heads * lm_expert_config.head_dim,
-                    bias=lm_expert_config.attention_bias,
-                )
-        # Remove unused embed_tokens
-        self.lm_expert.embed_tokens = None
-
-        self.num_attention_heads = self.config.text_config.num_attention_heads
-        self.num_key_value_heads = self.config.text_config.num_key_value_heads
-
-        self.freeze_vision_encoder = freeze_vision_encoder
-        self.train_expert_only = train_expert_only
-        self.attention_mode = attention_mode
-        self.expert_hidden_size = lm_expert_config.hidden_size
-        self.set_requires_grad()
-
-    def get_vlm_model(self):
-        return self.vlm.model
-
-    def set_requires_grad(self):
-        if self.freeze_vision_encoder:
-            self.get_vlm_model().vision_model.eval()
-            for params in self.get_vlm_model().vision_model.parameters():
-                params.requires_grad = False
-        if self.train_expert_only:
-            self.vlm.eval()
-            for params in self.vlm.parameters():
-                params.requires_grad = False
-        else:
-            # To avoid unused params issue with distributed training
-            last_layers = [self.num_vlm_layers - 1]
-            if (
-                self.num_vlm_layers != self.num_expert_layers
-                and self.num_vlm_layers % self.num_expert_layers == 0
-            ):
-                last_layers.append(self.num_vlm_layers - 2)
-            frozen_layers = [
-                "lm_head",
-                "text_model.model.norm.weight",
-            ]
-            for layer in last_layers:
-                frozen_layers.append(f"text_model.model.layers.{layer}.")
-
-            for name, params in self.vlm.named_parameters():
-                if any(k in name for k in frozen_layers):
-                    params.requires_grad = False
-        # To avoid unused params issue with distributed training
-        for name, params in self.lm_expert.named_parameters():
-            if "lm_head" in name:
-                params.requires_grad = False
-
-    def train(self, mode: bool = True):
-        super().train(mode)
-
-        if self.freeze_vision_encoder:
-            self.get_vlm_model().vision_model.eval()
-
-        if self.train_expert_only:
-            self.vlm.eval()
-
-    def embed_image(self, image: torch.Tensor):
-        patch_attention_mask = None
-        # Get sequence from the vision encoder
-        image_hidden_states = (
-            self.get_vlm_model()
-            .vision_model(
-                pixel_values=image.to(dtype=self.get_vlm_model().vision_model.dtype),
-                patch_attention_mask=patch_attention_mask,
-            )
-            .last_hidden_state
-        )
-        # Modality projection & resampling
-        image_hidden_states = self.get_vlm_model().connector(image_hidden_states)
-        return image_hidden_states
-
-    def embed_language_tokens(self, tokens: torch.Tensor):
-        return self.get_vlm_model().text_model.get_input_embeddings()(tokens)
-
-    def forward_attn_layer(
-        self,
-        model_layers,
-        inputs_embeds,
-        layer_idx,
-        position_ids,
-        attention_mask,
-        batch_size,
-        head_dim,
-        use_cache: bool = True,
-        fill_kv_cache: bool = True,
-        past_key_values=None,
-    ) -> list[torch.Tensor]:
-        query_states = []
-        key_states = []
-        value_states = []
-        for i, hidden_states in enumerate(inputs_embeds):
-            layer = model_layers[i][layer_idx]
-            if hidden_states is None or layer is None:
-                continue
-            hidden_states = layer.input_layernorm(hidden_states)
-
-            input_shape = hidden_states.shape[:-1]
-            hidden_shape = (*input_shape, -1, layer.self_attn.head_dim)
-
-            hidden_states = hidden_states.to(dtype=layer.self_attn.q_proj.weight.dtype)
-            query_state = layer.self_attn.q_proj(hidden_states).view(hidden_shape)
-            key_state = layer.self_attn.k_proj(hidden_states).view(hidden_shape)
-            value_state = layer.self_attn.v_proj(hidden_states).view(hidden_shape)
-
-            query_states.append(query_state)
-            key_states.append(key_state)
-            value_states.append(value_state)
-
-        # B,L,H,D with L sequence length, H number of heads, D head dim
-        # concatenate on the number of embeddings/tokens
-        query_states = torch.cat(query_states, dim=1)
-        key_states = torch.cat(key_states, dim=1)
-        value_states = torch.cat(value_states, dim=1)
-        seq_len = query_states.shape[1]
-        if seq_len < position_ids.shape[1]:
-            _position_ids = position_ids[:, :seq_len]
-            _attention_mask = attention_mask[:, :seq_len, :seq_len]
-        else:
-            _position_ids = position_ids
-            _attention_mask = attention_mask
-
-        attention_mask_ = _attention_mask
-        position_ids_ = _position_ids
-
-        query_states = apply_rope(query_states, position_ids_)
-        key_states = apply_rope(key_states, position_ids_)
-
-        if use_cache and past_key_values is None:
-            past_key_values = {}
-
-        if use_cache:
-            if fill_kv_cache:
-                past_key_values[layer_idx] = {
-                    "key_states": key_states,
-                    "value_states": value_states,
-                }
-            else:
-                # TODO here, some optimization can be done - similar to a `StaticCache` we can declare the `max_len` before.
-                # so we create an empty cache, with just one cuda malloc, and if (in autoregressive case) we reach
-                # the max len, then we (for instance) double the cache size. This implementation already exists
-                # in `transformers`. (molbap)
-                key_states = torch.cat([past_key_values[layer_idx]["key_states"], key_states], dim=1)
-                value_states = torch.cat([past_key_values[layer_idx]["value_states"], value_states], dim=1)
-
-        attention_interface = self.get_attention_interface()
-
-        att_output = attention_interface(
-            attention_mask_, batch_size, head_dim, query_states, key_states, value_states
-        )
-        return [att_output], past_key_values
-
-    def forward_cross_attn_layer(
-        self,
-        model_layers,
-        inputs_embeds,
-        layer_idx,
-        position_ids,
-        attention_mask,
-        batch_size,
-        head_dim,
-        use_cache: bool = True,
-        fill_kv_cache: bool = True,
-        past_key_values=None,
-    ) -> list[torch.Tensor]:
-        attention_interface = self.get_attention_interface()
-
-        att_outputs = []
-        assert len(inputs_embeds) == 2 or (use_cache and past_key_values is not None and not fill_kv_cache), (
-            f"Both len(inputs_embeds) == {len(inputs_embeds)} and past_key_values is {past_key_values}"
-        )
-
-        if len(inputs_embeds) == 2 and not past_key_values:
-            # Prefix attention
-            seq_len = inputs_embeds[0].shape[1]
-            position_id, expert_position_id = position_ids[:, :seq_len], position_ids[:, seq_len:]
-            prefix_attention_mask = attention_mask[:, :seq_len, :seq_len]
-
-            layer = model_layers[0][layer_idx]
-
-            hidden_states = layer.input_layernorm(inputs_embeds[0])
-
-            input_shape = hidden_states.shape[:-1]
-            hidden_shape = (*input_shape, -1, layer.self_attn.head_dim)
-
-            hidden_states = hidden_states.to(dtype=layer.self_attn.q_proj.weight.dtype)
-            query_state = layer.self_attn.q_proj(hidden_states).view(hidden_shape)
-            key_state = layer.self_attn.k_proj(hidden_states).view(hidden_shape)
-            value_states = layer.self_attn.v_proj(hidden_states).view(hidden_shape)
-
-            # B,L,H,D with L sequence length, H number of heads, D head dim
-            query_states = apply_rope(query_state, position_id)
-            key_states = apply_rope(key_state, position_id)
-
-            att_output = attention_interface(
-                prefix_attention_mask, batch_size, head_dim, query_states, key_states, value_states
-            )
-            att_outputs.append(att_output)
-        else:
-            expert_position_id = position_ids
-
-        if use_cache and past_key_values is None:
-            past_key_values = {}
-
-        if use_cache:
-            if fill_kv_cache:
-                past_key_values[layer_idx] = {
-                    "key_states": key_states,
-                    "value_states": value_states,
-                }
-            else:
-                # TODO here, some optimization can be done - similar to a `StaticCache` we can declare the `max_len` before.
-                # so we create an empty cache, with just one cuda malloc, and if (in autoregressive case) we reach
-                # the max len, then we (for instance) double the cache size. This implementation already exists
-                # in `transformers`. (molbap)
-                key_states = past_key_values[layer_idx]["key_states"]
-                value_states = past_key_values[layer_idx]["value_states"]
-
-        # Expert
-        expert_layer = model_layers[1][layer_idx]
-        if expert_layer is not None:
-            expert_hidden_states = expert_layer.input_layernorm(inputs_embeds[1])
-
-            expert_input_shape = expert_hidden_states.shape[:-1]
-            expert_hidden_shape = (*expert_input_shape, -1, expert_layer.self_attn.head_dim)
-
-            expert_hidden_states = expert_hidden_states.to(dtype=expert_layer.self_attn.q_proj.weight.dtype)
-            expert_query_state = expert_layer.self_attn.q_proj(expert_hidden_states).view(expert_hidden_shape)
-
-            _key_states = key_states.to(dtype=expert_layer.self_attn.k_proj.weight.dtype).view(
-                *key_states.shape[:2], -1
-            )
-            expert_key_states = expert_layer.self_attn.k_proj(_key_states).view(
-                *_key_states.shape[:-1], -1, expert_layer.self_attn.head_dim
-            )  # k_proj should have same dim as kv
-
-            _value_states = value_states.to(dtype=expert_layer.self_attn.v_proj.weight.dtype).view(
-                *value_states.shape[:2], -1
-            )
-            expert_value_states = expert_layer.self_attn.v_proj(_value_states).view(
-                *_value_states.shape[:-1], -1, expert_layer.self_attn.head_dim
-            )
-
-            expert_position_id = (
-                expert_position_id - torch.min(expert_position_id, dim=1, keepdim=True).values
-            )  # start from 0
-            expert_attention_mask = attention_mask[
-                :, -inputs_embeds[1].shape[1] :, : expert_key_states.shape[1] :
-            ]  # take into account kv
-
-            expert_query_states = apply_rope(expert_query_state, expert_position_id)
-
-            att_output = attention_interface(
-                expert_attention_mask,
-                batch_size,
-                head_dim,
-                expert_query_states,
-                expert_key_states,
-                expert_value_states,
-            )
-            att_outputs.append(att_output)
-        else:
-            att_outputs.append(None)
-
-        # att_output = att_output.to(dtype=models[i].dtype)
-        return att_outputs, past_key_values
-
-    def get_model_layers(self, models: list) -> list:
-        vlm_layers = []
-        expert_layers = []
-        multiple_of = self.num_vlm_layers // self.num_expert_layers
-        for i in range(self.num_vlm_layers):
-            if multiple_of > 0 and i > 0 and i % multiple_of != 0:
-                expert_layer = None
-            else:
-                expert_layer_index = i // multiple_of if multiple_of > 0 else i
-                expert_layer = models[1].layers[expert_layer_index]
-            vlm_layers.append(models[0].layers[i])
-            expert_layers.append(expert_layer)
-        return [vlm_layers, expert_layers]
-
-    def forward(
-        self,
-        attention_mask: torch.Tensor | None = None,
-        position_ids: torch.LongTensor | None = None,
-        past_key_values: list[torch.FloatTensor] | None = None,
-        inputs_embeds: list[torch.FloatTensor] = None,
-        use_cache: bool | None = None,
-        fill_kv_cache: bool | None = None,
-    ):
-        models = [self.get_vlm_model().text_model, self.lm_expert]
-        model_layers = self.get_model_layers(models)
-        for hidden_states in inputs_embeds:
-            # TODO this is very inefficient
-            # dtype is always the same, batch size too (if > 1 len)
-            # device could be trickier in multi gpu edge cases but that's it
-            if hidden_states is None:
-                continue
-            batch_size = hidden_states.shape[0]
-
-        # RMSNorm
-        num_layers = self.num_vlm_layers
-        head_dim = self.vlm.config.text_config.head_dim
-        for layer_idx in range(num_layers):
-            if (
-                fill_kv_cache
-                or "cross" not in self.attention_mode
-                or (self.self_attn_every_n_layers > 0 and layer_idx % self.self_attn_every_n_layers == 0)
-            ):
-                att_outputs, past_key_values = self.forward_attn_layer(
-                    model_layers,
-                    inputs_embeds,
-                    layer_idx,
-                    position_ids,
-                    attention_mask,
-                    batch_size,
-                    head_dim,
-                    use_cache=use_cache,
-                    fill_kv_cache=fill_kv_cache,
-                    past_key_values=past_key_values,
-                )
-            else:
-                att_outputs, past_key_values = self.forward_cross_attn_layer(
-                    model_layers,
-                    inputs_embeds,
-                    layer_idx,
-                    position_ids,
-                    attention_mask,
-                    batch_size,
-                    head_dim,
-                    use_cache=use_cache,
-                    fill_kv_cache=fill_kv_cache,
-                    past_key_values=past_key_values,
-                )
-            outputs_embeds = []
-            start = 0
-            for i, hidden_states in enumerate(inputs_embeds):
-                layer = model_layers[i][layer_idx]
-                att_output = (
-                    att_outputs[i] if i < len(att_outputs) else att_outputs[0]
-                )  # in case of self_attn
-                if hidden_states is not None:
-                    if layer is None:
-                        outputs_embeds.append(hidden_states)
-                        continue
-                    end = start + hidden_states.shape[1]
-
-                    if att_output.dtype != layer.self_attn.o_proj.weight.dtype:
-                        att_output = att_output.to(layer.self_attn.o_proj.weight.dtype)
-                    att_out = att_output[:, start:end]
-                    out_emb = layer.self_attn.o_proj(att_out)
-
-                    out_emb += hidden_states
-                    after_first_residual = out_emb.clone()
-
-                    out_emb = layer.post_attention_layernorm(out_emb)
-                    out_emb = layer.mlp(out_emb)
-
-                    out_emb += after_first_residual
-
-                    outputs_embeds.append(out_emb)
-
-                    start = end if len(att_outputs) == 1 else 0
-                else:
-                    outputs_embeds.append(None)
-
-            inputs_embeds = outputs_embeds
-
-        # final norm
-        outputs_embeds = []
-        for i, hidden_states in enumerate(inputs_embeds):
-            if hidden_states is not None:
-                out_emb = models[i].norm(hidden_states)
-                outputs_embeds.append(out_emb)
-            else:
-                outputs_embeds.append(None)
-        return outputs_embeds, past_key_values
-
-    def get_attention_interface(self):
-        attention_interface = self.eager_attention_forward
-        return attention_interface
-
-    def eager_attention_forward(
-        self, attention_mask, batch_size, head_dim, query_states, key_states, value_states
-    ):
-        num_att_heads = self.num_attention_heads
-        num_key_value_heads = self.num_key_value_heads
-        num_key_value_groups = num_att_heads // num_key_value_heads
-
-        sequence_length = key_states.shape[1]
-
-        key_states = key_states[:, :, :, None, :].expand(
-            batch_size, sequence_length, num_key_value_heads, num_key_value_groups, head_dim
-        )
-        key_states = key_states.reshape(
-            batch_size, sequence_length, num_key_value_heads * num_key_value_groups, head_dim
-        )
-
-        value_states = value_states[:, :, :, None, :].expand(
-            batch_size, sequence_length, num_key_value_heads, num_key_value_groups, head_dim
-        )
-        value_states = value_states.reshape(
-            batch_size, sequence_length, num_key_value_heads * num_key_value_groups, head_dim
-        )
-
-        # Attention here is upcasted to float32 to match the original eager implementation.
-        query_states = query_states.to(dtype=torch.float32)
-        key_states = key_states.to(dtype=torch.float32)
-
-        query_states = query_states.transpose(1, 2)
-        key_states = key_states.transpose(1, 2)
-
-        att_weights = torch.matmul(query_states, key_states.transpose(2, 3))
-        att_weights *= head_dim**-0.5
-
-        att_weights = att_weights.to(dtype=torch.float32)
-        big_neg = torch.finfo(att_weights.dtype).min  # -2.3819763e38  # See gemma/modules.py
-        masked_att_weights = torch.where(attention_mask[:, None, :, :], att_weights, big_neg)
-        probs = nn.functional.softmax(masked_att_weights, dim=-1)
-        probs = probs.to(dtype=value_states.dtype)
-
-        att_output = torch.matmul(probs, value_states.permute(0, 2, 1, 3))
-
-        att_output = att_output.permute(0, 2, 1, 3)
-        # we use -1 because sequence length can change
-        att_output = att_output.reshape(batch_size, -1, num_key_value_heads * num_key_value_groups * head_dim)
-
-        return att_output

roboimi/vla/modules/encoders.py

@@ -16,3 +16,23 @@ class IdentityActionEncoder(nn.Module):
 
     def forward(self, action):
         return action
+
+
+class LeWMStateEncoder(nn.Module):
+    def __init__(
+        self,
+        input_dim: int = 16,
+        hidden_dim: int = 256,
+        output_dim: int = 64,
+    ):
+        super().__init__()
+        self.output_dim = int(output_dim)
+        self.net = nn.Sequential(
+            nn.Linear(int(input_dim), int(hidden_dim)),
+            nn.LayerNorm(int(hidden_dim)),
+            nn.GELU(),
+            nn.Linear(int(hidden_dim), self.output_dim),
+        )
+
+    def forward(self, state):
+        return self.net(state)

tests/test_air_insert_env.py

@@ -1,510 +0,0 @@
-import importlib
-import inspect
-import pathlib
-import unittest
-from unittest import mock
-import xml.etree.ElementTree as ET
-
-import numpy as np
-
-from roboimi.envs.double_pos_ctrl_env import make_sim_env
-from roboimi.utils import act_ex_utils
-from roboimi.utils.constants import SIM_TASK_CONFIGS
-
-
-TASK_NAME = "sim_air_insert_socket_peg"
-
-
-class AirInsertTaskRegistrationTest(unittest.TestCase):
-    def test_sim_task_configs_registers_air_insert_socket_peg(self):
-        self.assertIn(TASK_NAME, SIM_TASK_CONFIGS)
-        self.assertNotIn("sim_air_insert_ring_bar", SIM_TASK_CONFIGS)
-        self.assertEqual(SIM_TASK_CONFIGS[TASK_NAME]["episode_len"], 750)
-        self.assertEqual(SIM_TASK_CONFIGS[TASK_NAME]["camera_names"], ["l_vis", "r_vis", "front"])
-        self.assertTrue(SIM_TASK_CONFIGS[TASK_NAME]["dataset_dir"].endswith("/sim_air_insert_socket_peg"))
-
-    def test_sample_air_insert_socket_peg_state_returns_explicit_named_mapping(self):
-        sampler = getattr(act_ex_utils, "sample_air_insert_socket_peg_state", None)
-        self.assertIsNotNone(
-            sampler,
-            "Expected roboimi.utils.act_ex_utils.sample_air_insert_socket_peg_state()",
-        )
-        self.assertFalse(
-            hasattr(act_ex_utils, "sample_air_insert_ring_bar_state"),
-            "air insert sampler should use socket/peg naming after the task rename",
-        )
-
-        task_state = sampler()
-
-        self.assertEqual(
-            list(task_state.keys()),
-            ["socket_pos", "socket_quat", "peg_pos", "peg_quat"],
-        )
-        self.assertEqual(task_state["socket_pos"].shape, (3,))
-        self.assertEqual(task_state["socket_quat"].shape, (4,))
-        self.assertEqual(task_state["peg_pos"].shape, (3,))
-        self.assertEqual(task_state["peg_quat"].shape, (4,))
-
-    def test_sample_air_insert_socket_peg_state_uses_fixed_quats_and_left_right_planar_ranges(self):
-        sampler = getattr(act_ex_utils, "sample_air_insert_socket_peg_state", None)
-        self.assertIsNotNone(sampler)
-
-        task_state = sampler()
-
-        np.testing.assert_array_equal(task_state["socket_quat"], np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32))
-        np.testing.assert_array_equal(task_state["peg_quat"], np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32))
-        self.assertGreaterEqual(task_state["socket_pos"][0], -0.20)
-        self.assertLessEqual(task_state["socket_pos"][0], -0.05)
-        self.assertGreaterEqual(task_state["socket_pos"][1], 0.70)
-        self.assertLessEqual(task_state["socket_pos"][1], 1.00)
-        self.assertAlmostEqual(float(task_state["socket_pos"][2]), 0.472)
-        self.assertGreaterEqual(task_state["peg_pos"][0], 0.05)
-        self.assertLessEqual(task_state["peg_pos"][0], 0.20)
-        self.assertGreaterEqual(task_state["peg_pos"][1], 0.70)
-        self.assertLessEqual(task_state["peg_pos"][1], 1.00)
-        self.assertAlmostEqual(float(task_state["peg_pos"][2]), 0.46)
-
-    def test_make_sim_env_dispatches_air_insert_socket_peg_headless(self):
-        air_insert_env = importlib.import_module("roboimi.envs.double_air_insert_env")
-        air_insert_cls = getattr(air_insert_env, "DualDianaMed_Air_Insert", None)
-        self.assertIsNotNone(air_insert_cls)
-
-        diana_med = importlib.import_module("roboimi.assets.robots.diana_med")
-        socket_peg_robot_cls = getattr(diana_med, "BiDianaMedSocketPeg", None)
-        self.assertIsNotNone(
-            socket_peg_robot_cls,
-            "Expected roboimi.assets.robots.diana_med.BiDianaMedSocketPeg",
-        )
-
-        fake_env = object()
-        with mock.patch.object(
-            diana_med,
-            "BiDianaMedSocketPeg",
-            return_value="robot",
-        ), mock.patch.object(
-            air_insert_env,
-            "DualDianaMed_Air_Insert",
-            return_value=fake_env,
-        ) as env_cls:
-            env = make_sim_env(TASK_NAME, headless=True)
-
-        self.assertIs(env, fake_env)
-        env_cls.assert_called_once_with(
-            robot="robot",
-            is_render=False,
-            control_freq=30,
-            is_interpolate=True,
-            cam_view="front",
-        )
-
-    def test_diana_table_scene_exposes_only_top_and_front_scene_cameras(self):
-        xml_path = (
-            pathlib.Path(__file__).resolve().parents[1]
-            / "roboimi/assets/models/manipulators/DianaMed/table_square.xml"
-        )
-        root = ET.parse(xml_path).getroot()
-        cameras = {camera.attrib["name"]: camera.attrib for camera in root.findall(".//camera")}
-
-        self.assertNotIn("angle", cameras, "DianaMed scene should stop exposing the old angle camera")
-        self.assertNotIn("left_side", cameras, "DianaMed scene should no longer expose left_side")
-        self.assertIn("top", cameras)
-        self.assertIn("front", cameras)
-        self.assertEqual(cameras["top"].get("mode"), "targetbody")
-        self.assertEqual(cameras["top"].get("target"), "table")
-
-
-class AirInsertResetAndStateHelpersTest(unittest.TestCase):
-    def test_set_socket_peg_task_state_writes_free_joint_qpos(self):
-        air_insert_env = importlib.import_module("roboimi.envs.double_air_insert_env")
-        setter = getattr(air_insert_env, "set_socket_peg_task_state", None)
-        self.assertIsNotNone(
-            setter,
-            "Expected roboimi.envs.double_air_insert_env.set_socket_peg_task_state",
-        )
-
-        socket_qpos = np.zeros(7, dtype=np.float64)
-        peg_qpos = np.zeros(7, dtype=np.float64)
-
-        class _FakeJoint:
-            def __init__(self, qpos):
-                self.qpos = qpos
-
-        class _FakeData:
-            def joint(self, name):
-                if name == "blue_socket_joint":
-                    return _FakeJoint(socket_qpos)
-                if name == "red_peg_joint":
-                    return _FakeJoint(peg_qpos)
-                raise AssertionError(f"Unexpected joint name: {name}")
-
-        task_state = {
-            "socket_pos": np.array([-0.12, 0.90, 0.472], dtype=np.float64),
-            "socket_quat": np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float64),
-            "peg_pos": np.array([0.12, 0.91, 0.46], dtype=np.float64),
-            "peg_quat": np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float64),
-        }
-
-        setter(_FakeData(), task_state)
-
-        np.testing.assert_array_equal(
-            socket_qpos,
-            np.array([-0.12, 0.90, 0.472, 1.0, 0.0, 0.0, 0.0], dtype=np.float64),
-        )
-        np.testing.assert_array_equal(
-            peg_qpos,
-            np.array([0.12, 0.91, 0.46, 1.0, 0.0, 0.0, 0.0], dtype=np.float64),
-        )
-
-    def test_get_socket_peg_env_state_returns_stable_14d_vector(self):
-        air_insert_env = importlib.import_module("roboimi.envs.double_air_insert_env")
-        getter = getattr(air_insert_env, "get_socket_peg_env_state", None)
-        self.assertIsNotNone(
-            getter,
-            "Expected roboimi.envs.double_air_insert_env.get_socket_peg_env_state",
-        )
-
-        socket_qpos = np.array([-0.12, 0.90, 0.472, 1.0, 0.0, 0.0, 0.0], dtype=np.float64)
-        peg_qpos = np.array([0.12, 0.91, 0.46, 1.0, 0.0, 0.0, 0.0], dtype=np.float64)
-
-        class _FakeJoint:
-            def __init__(self, qpos):
-                self.qpos = qpos
-
-        class _FakeData:
-            def joint(self, name):
-                if name == "blue_socket_joint":
-                    return _FakeJoint(socket_qpos)
-                if name == "red_peg_joint":
-                    return _FakeJoint(peg_qpos)
-                raise AssertionError(f"Unexpected joint name: {name}")
-
-        env_state = getter(_FakeData())
-
-        self.assertEqual(env_state.shape, (14,))
-        np.testing.assert_array_equal(
-            env_state,
-            np.array(
-                [-0.12, 0.90, 0.472, 1.0, 0.0, 0.0, 0.0, 0.12, 0.91, 0.46, 1.0, 0.0, 0.0, 0.0],
-                dtype=np.float64,
-            ),
-        )
-
-    def test_air_insert_env_does_not_script_attach_or_assist_objects_after_reset(self):
-        air_insert_env = importlib.import_module("roboimi.envs.double_air_insert_env")
-        env_cls = getattr(air_insert_env, "DualDianaMed_Air_Insert", None)
-        self.assertIsNotNone(env_cls)
-
-        source = inspect.getsource(env_cls)
-
-        self.assertNotIn("_update_scripted_grasped_objects", source)
-        self.assertNotIn("_scripted_", source)
-        self.assertNotIn("_stabilize_ring_grasp", source)
-        self.assertNotIn("_ring_grasp_locked", source)
-        get_reward_source = inspect.getsource(env_cls._get_reward)
-        self.assertNotIn("ring_block", get_reward_source)
-        self.assertNotIn("bar_block", get_reward_source)
-
-    def test_socket_peg_xml_defines_active_socket_and_peg_objects(self):
-        asset_dir = pathlib.Path(__file__).resolve().parents[1] / "roboimi/assets/models/manipulators/DianaMed"
-        xml_path = asset_dir / "socket_peg_objects.xml"
-        self.assertTrue(xml_path.exists(), "socket/peg objects should live in socket_peg_objects.xml")
-        self.assertFalse((asset_dir / "ring_bar_objects.xml").exists(), "old ring_bar_objects.xml should be renamed")
-
-        root = ET.parse(xml_path).getroot()
-        body_names = {body.attrib.get("name") for body in root.findall(".//body")}
-        geom_names = {geom.attrib.get("name") for geom in root.findall(".//geom")}
-        joint_names = {joint.attrib.get("name") for joint in root.findall(".//joint")}
-
-        self.assertIn("socket", body_names)
-        self.assertIn("peg", body_names)
-        self.assertNotIn("ring_block", body_names)
-        self.assertNotIn("bar_block", body_names)
-        self.assertIn("blue_socket_joint", joint_names)
-        self.assertIn("red_peg_joint", joint_names)
-        for geom_name in ("socket-1", "socket-2", "socket-3", "socket-4", "pin", "red_peg"):
-            self.assertIn(geom_name, geom_names)
-
-    def test_socket_peg_wrapper_includes_socket_peg_objects(self):
-        xml_path = (
-            pathlib.Path(__file__).resolve().parents[1]
-            / "roboimi/assets/models/manipulators/DianaMed/bi_diana_socket_peg_ee.xml"
-        )
-        self.assertTrue(xml_path.exists(), "socket/peg wrapper XML should use the new task name")
-        root = ET.parse(xml_path).getroot()
-        includes = [include.attrib.get("file") for include in root.findall(".//include")]
-        self.assertIn("./socket_peg_objects.xml", includes)
-        self.assertNotIn("./ring_bar_objects.xml", includes)
-
-
-class AirInsertRewardAndSuccessTest(unittest.TestCase):
-    @staticmethod
-    def _make_env_state(
-        socket_pos=(0.0, 0.0, 0.472),
-        socket_quat=(1.0, 0.0, 0.0, 0.0),
-        peg_pos=(0.0, 0.0, 0.46),
-        peg_quat=(1.0, 0.0, 0.0, 0.0),
-    ):
-        return np.array([*socket_pos, *socket_quat, *peg_pos, *peg_quat], dtype=np.float64)
-
-    def test_compute_air_insert_reward_counts_left_contact_stage(self):
-        air_insert_env = importlib.import_module("roboimi.envs.double_air_insert_env")
-        reward_fn = getattr(air_insert_env, "compute_air_insert_reward", None)
-        self.assertIsNotNone(reward_fn)
-
-        reward = reward_fn(
-            contact_pairs=[
-                ("socket-1", "l_finger_left"),
-                ("socket-1", "table"),
-                ("red_peg", "table"),
-            ],
-            env_state=self._make_env_state(),
-        )
-
-        self.assertEqual(reward, 1)
-
-    def test_compute_air_insert_reward_counts_right_contact_stage(self):
-        air_insert_env = importlib.import_module("roboimi.envs.double_air_insert_env")
-        reward_fn = getattr(air_insert_env, "compute_air_insert_reward", None)
-
-        reward = reward_fn(
-            contact_pairs=[
-                ("socket-1", "l_finger_left"),
-                ("red_peg", "l_finger_right"),
-                ("socket-1", "table"),
-                ("red_peg", "table"),
-            ],
-            env_state=self._make_env_state(),
-        )
-
-        self.assertEqual(reward, 2)
-
-    def test_compute_air_insert_reward_counts_lift_stages(self):
-        air_insert_env = importlib.import_module("roboimi.envs.double_air_insert_env")
-        reward_fn = getattr(air_insert_env, "compute_air_insert_reward", None)
-
-        reward = reward_fn(
-            contact_pairs=[
-                ("socket-1", "l_finger_left"),
-                ("red_peg", "l_finger_right"),
-            ],
-            env_state=self._make_env_state(),
-        )
-
-        self.assertEqual(reward, 4)
-
-    def test_compute_air_insert_reward_counts_visual_fingertip_contacts_as_gripper_contacts(self):
-        air_insert_env = importlib.import_module("roboimi.envs.double_air_insert_env")
-        reward_fn = getattr(air_insert_env, "compute_air_insert_reward", None)
-
-        reward = reward_fn(
-            contact_pairs=[
-                ("socket-3", "r_fingertip_g0_vis_left"),
-                ("red_peg", "l_fingertip_g0_vis_right"),
-            ],
-            env_state=self._make_env_state(),
-        )
-
-        self.assertEqual(
-            reward,
-            4,
-            "visual fingertip geoms are collidable in the Diana XML and should count as gripper-object contacts",
-        )
-
-    def test_peg_inserted_into_socket_uses_pin_contact(self):
-        air_insert_env = importlib.import_module("roboimi.envs.double_air_insert_env")
-        success_fn = getattr(air_insert_env, "peg_inserted_into_socket", None)
-        self.assertIsNotNone(
-            success_fn,
-            "Expected roboimi.envs.double_air_insert_env.peg_inserted_into_socket",
-        )
-
-        self.assertTrue(success_fn([("red_peg", "pin")]))
-        self.assertTrue(success_fn([("pin", "red_peg")]))
-        self.assertFalse(success_fn([("red_peg", "socket-1")]))
-
-    def test_compute_air_insert_reward_requires_airborne_success_for_final_point(self):
-        air_insert_env = importlib.import_module("roboimi.envs.double_air_insert_env")
-        reward_fn = getattr(air_insert_env, "compute_air_insert_reward", None)
-
-        reward = reward_fn(
-            contact_pairs=[
-                ("socket-1", "l_finger_left"),
-                ("red_peg", "l_finger_right"),
-                ("socket-1", "table"),
-                ("red_peg", "pin"),
-            ],
-            env_state=self._make_env_state(),
-        )
-
-        self.assertEqual(reward, 3)
-
-    def test_compute_air_insert_reward_returns_full_score_on_true_airborne_insert(self):
-        air_insert_env = importlib.import_module("roboimi.envs.double_air_insert_env")
-        reward_fn = getattr(air_insert_env, "compute_air_insert_reward", None)
-
-        reward = reward_fn(
-            contact_pairs=[
-                ("socket-1", "l_finger_left"),
-                ("red_peg", "l_finger_right"),
-                ("red_peg", "pin"),
-            ],
-            env_state=self._make_env_state(),
-        )
-
-        self.assertEqual(reward, 5)
-
-
-class AirInsertPolicyAndSmokeTest(unittest.TestCase):
-    @staticmethod
-    def _canonical_task_state():
-        return {
-            "socket_pos": np.array([-0.12, 0.90, 0.472], dtype=np.float32),
-            "socket_quat": np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32),
-            "peg_pos": np.array([0.12, 0.90, 0.46], dtype=np.float32),
-            "peg_quat": np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32),
-        }
-
-    def test_air_insert_policy_emits_valid_16d_action(self):
-        policy_module = importlib.import_module("roboimi.demos.diana_air_insert_policy")
-        policy_cls = getattr(policy_module, "TestAirInsertPolicy", None)
-        self.assertIsNotNone(policy_cls)
-
-        task_state = act_ex_utils.sample_air_insert_socket_peg_state()
-        policy = policy_cls(inject_noise=False)
-        action = policy.predict(task_state, 0)
-
-        self.assertEqual(action.shape, (16,))
-        np.testing.assert_array_equal(action[-2:], np.array([100, 100]))
-
-    def test_air_insert_policy_inserts_peg_front_view_right_to_left_along_world_x(self):
-        policy_module = importlib.import_module("roboimi.demos.diana_air_insert_policy")
-        policy_cls = getattr(policy_module, "TestAirInsertPolicy", None)
-        self.assertIsNotNone(policy_cls)
-
-        task_state = self._canonical_task_state()
-        policy = policy_cls(inject_noise=False)
-        policy.generate_trajectory(task_state)
-
-        start_waypoint = next(wp for wp in policy.right_trajectory if wp["t"] == policy.INSERT_START_T)
-        end_waypoint = next(wp for wp in policy.right_trajectory if wp["t"] == policy.INSERT_END_T)
-
-        self.assertLess(
-            end_waypoint["xyz"][0],
-            start_waypoint["xyz"][0] - 0.10,
-            "front-view right-to-left peg insertion should decrease world x substantially",
-        )
-        self.assertAlmostEqual(float(end_waypoint["xyz"][1]), float(start_waypoint["xyz"][1]), delta=0.02)
-        expected_insert_end_x = float(task_state["socket_pos"][0] + 0.168)
-        self.assertAlmostEqual(float(end_waypoint["xyz"][0]), expected_insert_end_x, delta=0.02)
-        self.assertGreater(float(start_waypoint["xyz"][2]), 0.70)
-
-    def test_air_insert_policy_default_left_grasps_socket_and_right_grasps_peg(self):
-        policy_module = importlib.import_module("roboimi.demos.diana_air_insert_policy")
-        policy_cls = getattr(policy_module, "TestAirInsertPolicy", None)
-        self.assertIsNotNone(policy_cls)
-
-        task_state = {
-            "socket_pos": np.array([-0.18, 0.78, 0.472], dtype=np.float32),
-            "socket_quat": np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32),
-            "peg_pos": np.array([0.16, 0.98, 0.46], dtype=np.float32),
-            "peg_quat": np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32),
-        }
-
-        policy = policy_cls(inject_noise=False)
-        policy.generate_trajectory(task_state)
-        left_close = next(wp for wp in policy.left_trajectory if wp["t"] == 180)
-        right_close = next(wp for wp in policy.right_trajectory if wp["t"] == 180)
-        action_z_offset = getattr(policy_cls, "ACTION_OBJECT_Z_OFFSET", 0.11)
-        expected_socket_pick = task_state["socket_pos"] + np.array([-0.078, 0.0, action_z_offset])
-        expected_peg_pick = task_state["peg_pos"] + np.array([0.078, 0.0, action_z_offset + 0.01])
-
-        np.testing.assert_allclose(left_close["xyz"], expected_socket_pick, atol=1e-6)
-        np.testing.assert_allclose(right_close["xyz"], expected_peg_pick, atol=1e-6)
-        self.assertLess(left_close["gripper"], 0, "default policy should close the left gripper on the socket")
-        self.assertLess(right_close["gripper"], 0, "default policy should close the right gripper on the peg")
-
-    def test_air_insert_policy_socket_hold_tracks_socket_xy_without_sweeping_laterally(self):
-        policy_module = importlib.import_module("roboimi.demos.diana_air_insert_policy")
-        policy_cls = getattr(policy_module, "TestAirInsertPolicy", None)
-        self.assertIsNotNone(policy_cls)
-
-        base_state = {
-            "socket_pos": np.array([-0.20, 0.72, 0.472], dtype=np.float32),
-            "socket_quat": np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32),
-            "peg_pos": np.array([0.14, 0.76, 0.46], dtype=np.float32),
-            "peg_quat": np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32),
-        }
-        shifted_state = dict(base_state)
-        shifted_state["socket_pos"] = np.array([-0.06, 0.99, 0.472], dtype=np.float32)
-
-        base_policy = policy_cls(inject_noise=False)
-        base_policy.generate_trajectory(base_state)
-        shifted_policy = policy_cls(inject_noise=False)
-        shifted_policy.generate_trajectory(shifted_state)
-
-        base_hold = next(wp for wp in base_policy.left_trajectory if wp["t"] == 450)
-        shifted_hold = next(wp for wp in shifted_policy.left_trajectory if wp["t"] == 450)
-        np.testing.assert_allclose(
-            base_hold["xyz"][:2],
-            base_state["socket_pos"][:2] + np.array([-0.078, 0.0]),
-            atol=1e-6,
-        )
-        np.testing.assert_allclose(
-            shifted_hold["xyz"][:2],
-            shifted_state["socket_pos"][:2] + np.array([-0.078, 0.0]),
-            atol=1e-6,
-        )
-
-    def test_air_insert_policy_predicts_through_full_episode_without_exhausting_waypoints(self):
-        policy_module = importlib.import_module("roboimi.demos.diana_air_insert_policy")
-        policy_cls = getattr(policy_module, "TestAirInsertPolicy", None)
-        self.assertIsNotNone(policy_cls)
-
-        task_state = self._canonical_task_state()
-        policy = policy_cls(inject_noise=False)
-
-        for step in range(SIM_TASK_CONFIGS[TASK_NAME]["episode_len"]):
-            action = policy.predict(task_state, step)
-            self.assertEqual(action.shape, (16,))
-
-    def test_scripted_rollout_entrypoint_selects_socket_peg_sampler_and_policy(self):
-        rollout_module = importlib.import_module("roboimi.demos.diana_record_sim_episodes")
-        sampler_fn = getattr(rollout_module, "sample_task_state", None)
-        policy_factory = getattr(rollout_module, "make_policy", None)
-        self.assertIsNotNone(sampler_fn)
-        self.assertIsNotNone(policy_factory)
-
-        task_state = sampler_fn(TASK_NAME)
-        self.assertEqual(list(task_state.keys()), ["socket_pos", "socket_quat", "peg_pos", "peg_quat"])
-
-        policy = policy_factory(TASK_NAME, inject_noise=False)
-        self.assertEqual(policy.__class__.__name__, "TestAirInsertPolicy")
-
-    def test_real_headless_smoke_instantiates_resets_and_steps_new_task_once(self):
-        policy_module = importlib.import_module("roboimi.demos.diana_air_insert_policy")
-        policy_cls = getattr(policy_module, "TestAirInsertPolicy", None)
-        self.assertIsNotNone(policy_cls)
-
-        task_state = act_ex_utils.sample_air_insert_socket_peg_state()
-        env = make_sim_env(TASK_NAME, headless=True)
-        policy = policy_cls(inject_noise=False)
-
-        try:
-            env.reset(task_state)
-            action = policy.predict(task_state, 0)
-            env.step(action)
-            self.assertIsNotNone(env.obs)
-            self.assertIn("qpos", env.obs)
-            self.assertIn("images", env.obs)
-        finally:
-            env.exit_flag = True
-            cam_thread = getattr(env, "cam_thread", None)
-            if cam_thread is not None:
-                cam_thread.join(timeout=1.0)
-            viewer = getattr(env, "viewer", None)
-            if viewer is not None:
-                viewer.close()
-
-
-if __name__ == "__main__":
-    unittest.main()

tests/test_eval_vla_execution.py

@@ -129,55 +129,6 @@ class EvalVLAExecutionTest(unittest.TestCase):
             ["r_vis", "top", "front"],
         )
 
-    def test_resolve_eval_image_resize_shape_prefers_agent_top_level_override(self):
-        cfg = OmegaConf.create(
-            {
-                "agent": {
-                    "eval_image_resize_shape": None,
-                    "condition_encoder": {
-                        "eval_image_resize_shape": [256, 256],
-                    },
-                },
-                "data": {
-                    "image_resize_shape": [224, 224],
-                },
-            }
-        )
-
-        self.assertIsNone(eval_vla._resolve_eval_image_resize_shape(cfg))
-
-    def test_resolve_eval_image_resize_shape_prefers_condition_encoder_override(self):
-        cfg = OmegaConf.create(
-            {
-                "agent": {
-                    "condition_encoder": {
-                        "eval_image_resize_shape": None,
-                    },
-                },
-                "data": {
-                    "image_resize_shape": [224, 224],
-                },
-            }
-        )
-
-        self.assertIsNone(eval_vla._resolve_eval_image_resize_shape(cfg))
-
-    def test_resolve_eval_image_resize_shape_prefers_vision_backbone_override(self):
-        cfg = OmegaConf.create(
-            {
-                "agent": {
-                    "vision_backbone": {
-                        "eval_image_resize_shape": [256, 256],
-                    },
-                },
-                "data": {
-                    "image_resize_shape": [224, 224],
-                },
-            }
-        )
-
-        self.assertEqual(eval_vla._resolve_eval_image_resize_shape(cfg), (256, 256))
-
     def test_build_episode_plans_without_box_poses_keeps_serial_sampling_lazy(self):
         plans = eval_vla._build_episode_plans(num_episodes=3)
 
@@ -217,58 +168,6 @@ class EvalVLAExecutionTest(unittest.TestCase):
             np.array([[[[1.0]]], [[[1.0]]], [[[1.0]]]], dtype=np.float32),
         )
 
-    def test_prepare_local_policy_batch_keeps_latest_variable_task_when_present(self):
-        queues = eval_vla._new_local_policy_queues(obs_horizon=2)
-        first_observation = {
-            "qpos": torch.tensor([1.0, 2.0], dtype=torch.float32),
-            "images": {"front": torch.tensor([[[1.0]]], dtype=torch.float32)},
-            "task": "pick the red cube",
-        }
-        second_observation = {
-            "qpos": torch.tensor([3.0, 4.0], dtype=torch.float32),
-            "images": {"front": torch.tensor([[[2.0]]], dtype=torch.float32)},
-            "task": "insert the peg into the socket",
-        }
-
-        eval_vla._populate_local_policy_queues(queues, first_observation)
-        eval_vla._populate_local_policy_queues(queues, second_observation)
-        batch = eval_vla._prepare_local_policy_batch(
-            queues,
-            obs_horizon=2,
-            camera_names=["front"],
-        )
-
-        self.assertEqual(batch["task"], ["insert the peg into the socket"])
-
-    def test_prepare_local_policy_batch_omits_task_for_legacy_observations(self):
-        queues = eval_vla._new_local_policy_queues(obs_horizon=2)
-        observation = {
-            "qpos": torch.tensor([1.0, 2.0], dtype=torch.float32),
-            "images": {"front": torch.tensor([[[1.0]]], dtype=torch.float32)},
-        }
-
-        eval_vla._populate_local_policy_queues(queues, observation)
-        batch = eval_vla._prepare_local_policy_batch(
-            queues,
-            obs_horizon=2,
-            camera_names=["front"],
-        )
-
-        self.assertNotIn("task", batch)
-
-    def test_serialize_deserialize_policy_batch_preserves_task(self):
-        batch = {
-            "qpos": torch.zeros(1, 2, 2, dtype=torch.float32),
-            "images": {"front": torch.zeros(1, 2, 1, 1, 1, dtype=torch.float32)},
-            "task": ["pick the red cube", "insert the peg into the socket"],
-        }
-
-        serialized = eval_vla._serialize_policy_batch(batch)
-        deserialized = eval_vla._deserialize_policy_batch(serialized, device="cpu")
-
-        self.assertEqual(serialized["task"], batch["task"])
-        self.assertEqual(deserialized["task"], batch["task"])
-
     def test_enqueue_predicted_actions_uses_executable_slice(self):
         queues = eval_vla._new_local_policy_queues(obs_horizon=2)
         predicted_actions = torch.tensor(
@@ -287,25 +186,6 @@ class EvalVLAExecutionTest(unittest.TestCase):
         np.testing.assert_array_equal(queues["action"].popleft().numpy(), np.array([20.0], dtype=np.float32))
         np.testing.assert_array_equal(queues["action"].popleft().numpy(), np.array([30.0], dtype=np.float32))
 
-    def test_enqueue_predicted_actions_honors_explicit_chunk_start(self):
-        queues = eval_vla._new_local_policy_queues(obs_horizon=2)
-        predicted_actions = torch.tensor(
-            [[[10.0], [20.0], [30.0], [40.0]]],
-            dtype=torch.float32,
-        )
-
-        eval_vla._enqueue_predicted_actions(
-            queues,
-            predicted_actions=predicted_actions,
-            obs_horizon=2,
-            num_action_steps=2,
-            action_chunk_start=0,
-        )
-
-        self.assertEqual(len(queues["action"]), 2)
-        np.testing.assert_array_equal(queues["action"].popleft().numpy(), np.array([10.0], dtype=np.float32))
-        np.testing.assert_array_equal(queues["action"].popleft().numpy(), np.array([20.0], dtype=np.float32))
-
     def test_remote_policy_runner_only_requests_server_inference_when_local_action_queue_is_empty(self):
         request_queue = _FakeQueue()
         response_queue = _FakeQueue(
@@ -324,7 +204,6 @@ class EvalVLAExecutionTest(unittest.TestCase):
             camera_names=["front"],
             obs_horizon=2,
             num_action_steps=2,
-            action_chunk_start=0,
         )
         first_observation = {
             "qpos": torch.tensor([1.0, 2.0], dtype=torch.float32),
@@ -352,46 +231,8 @@ class EvalVLAExecutionTest(unittest.TestCase):
         self.assertEqual(request_queue.put_calls[0]["type"], "predict_chunk")
         self.assertEqual(request_queue.put_calls[0]["worker_index"], 3)
         self.assertEqual(request_queue.put_calls[0]["server_index"], 1)
-        np.testing.assert_array_equal(first_action.numpy(), np.array([10.0], dtype=np.float32))
+        np.testing.assert_array_equal(first_action.numpy(), np.array([20.0], dtype=np.float32))
-        np.testing.assert_array_equal(second_action.numpy(), np.array([20.0], dtype=np.float32))
+        np.testing.assert_array_equal(second_action.numpy(), np.array([30.0], dtype=np.float32))
-
-    def test_remote_eval_worker_passes_agent_action_chunk_start_to_remote_runner(self):
-        cfg = OmegaConf.create(
-            {
-                "agent": {
-                    "obs_horizon": 2,
-                    "num_action_steps": 2,
-                    "action_chunk_start": 0,
-                    "camera_names": ["front"],
-                },
-                "eval": {
-                    "obs_horizon": 2,
-                    "num_queries": 2,
-                    "response_timeout_s": 3.0,
-                    "camera_names": ["front"],
-                },
-            }
-        )
-        captured = {}
-
-        class CapturingRunner:
-            def __init__(self, **kwargs):
-                captured.update(kwargs)
-
-        with mock.patch.object(eval_vla, "_RemotePolicyRunner", CapturingRunner), \
-             mock.patch.object(eval_vla, "_run_eval_episode_plans", return_value={"ok": True}) as run_plans:
-            result = eval_vla._run_remote_eval_worker(
-                cfg,
-                episode_plans=[{"episode_index": 0}],
-                worker_index=1,
-                server_index=2,
-                request_queue=_FakeQueue(),
-                response_queue=_FakeQueue(),
-            )
-
-        self.assertEqual(result, {"ok": True})
-        self.assertEqual(captured["action_chunk_start"], 0)
-        run_plans.assert_called_once()
 
     def test_merge_worker_summaries_sorts_episodes_and_recomputes_aggregates(self):
         worker_summaries = [

tests/test_eval_vla_headless.py

@@ -15,7 +15,6 @@ class _FakeAgent:
     def __init__(self):
         self.reset_calls = 0
         self.last_observation = None
-        self.observation_shapes = []
 
     def eval(self):
         return self
@@ -28,7 +27,6 @@ class _FakeAgent:
 
     def select_action(self, observation):
         self.last_observation = observation
-        self.observation_shapes.append(tuple(observation["images"]["front"].shape))
         return torch.zeros(16)
 
 
@@ -38,8 +36,8 @@ class _FakeEnv:
         self.render_calls = 0
         self.reset_calls = []
 
-    def reset(self, task_state):
+    def reset(self, box_pos):
-        self.reset_calls.append(task_state)
+        self.reset_calls.append(np.array(box_pos))
 
     def _get_image_obs(self):
         self.image_obs_calls += 1
@@ -76,7 +74,7 @@ class _FakeRenderer:
         self._env = env
         self._frames = [
             np.full((4, 4, 3), fill_value=index, dtype=np.uint8)
-            for index in range(8)
+            for index in range(5)
         ]
         self._index = 0
 
@@ -110,23 +108,6 @@ class EvalVLAHeadlessTest(unittest.TestCase):
         self.assertEqual(tuple(prepared["images"]["front"].shape), (3, 8, 8))
         self.assertEqual(tuple(prepared["qpos"].shape), (16,))
 
-    def test_prepare_observation_preserves_task_when_present(self):
-        obs = {
-            "images": {
-                "front": np.zeros((4, 4, 3), dtype=np.uint8),
-            },
-            "qpos": np.zeros(16, dtype=np.float32),
-            "task": "insert the peg into the socket",
-        }
-
-        prepared = eval_vla.prepare_observation(
-            obs,
-            ["front"],
-            image_resize_shape=None,
-        )
-
-        self.assertEqual(prepared["task"], "insert the peg into the socket")
-
     def test_headless_eval_sets_mujoco_gl_to_egl_when_display_missing(self):
         cfg = OmegaConf.create({"eval": {"headless": True}})
         with mock.patch.dict(eval_vla.os.environ, {}, clear=True):
@@ -144,8 +125,26 @@ class EvalVLAHeadlessTest(unittest.TestCase):
 
         self.assertIn("headless", eval_cfg)
         self.assertFalse(eval_cfg.headless)
-        self.assertIn("task_description", eval_cfg)
+
-        self.assertIsNone(eval_cfg.task_description)
+    def test_eval_config_exposes_num_workers_default(self):
+        eval_cfg = OmegaConf.load(Path("roboimi/vla/conf/eval/eval.yaml"))
+
+        self.assertIn("num_workers", eval_cfg)
+        self.assertEqual(eval_cfg.num_workers, 1)
+
+    def test_eval_config_exposes_cuda_devices_default(self):
+        eval_cfg = OmegaConf.load(Path("roboimi/vla/conf/eval/eval.yaml"))
+
+        self.assertIn("cuda_devices", eval_cfg)
+        self.assertIsNone(eval_cfg.cuda_devices)
+
+    def test_eval_config_exposes_parallel_timeout_defaults(self):
+        eval_cfg = OmegaConf.load(Path("roboimi/vla/conf/eval/eval.yaml"))
+
+        self.assertIn("response_timeout_s", eval_cfg)
+        self.assertIn("server_startup_timeout_s", eval_cfg)
+        self.assertEqual(eval_cfg.response_timeout_s, 300.0)
+        self.assertEqual(eval_cfg.server_startup_timeout_s, 300.0)
 
     def test_make_sim_env_accepts_headless_and_disables_render(self):
         fake_env = object()
@@ -165,7 +164,7 @@ class EvalVLAHeadlessTest(unittest.TestCase):
             is_render=False,
             control_freq=30,
             is_interpolate=True,
-            cam_view="top",
+            cam_view="angle",
         )
 
     def test_headless_sync_camera_capture_populates_images_without_gui_calls(self):
@@ -174,10 +173,11 @@ class EvalVLAHeadlessTest(unittest.TestCase):
         env.mj_data = object()
         env.exit_flag = False
         env.is_render = False
-        env.cam = 'top'
+        env.cam = 'angle'
         env.r_vis = None
         env.l_vis = None
         env.top = None
+        env.angle = None
         env.front = None
         env._offscreen_renderer = None
 
@@ -196,6 +196,7 @@ class EvalVLAHeadlessTest(unittest.TestCase):
         self.assertIsNotNone(env.r_vis)
         self.assertIsNotNone(env.l_vis)
         self.assertIsNotNone(env.top)
+        self.assertIsNotNone(env.angle)
         self.assertIsNotNone(env.front)
 
     def test_cam_start_skips_background_thread_when_headless(self):
@@ -215,10 +216,11 @@ class EvalVLAHeadlessTest(unittest.TestCase):
         env.mj_data = object()
         env.exit_flag = False
         env.is_render = False
-        env.cam = "top"
+        env.cam = "angle"
         env.r_vis = None
         env.l_vis = None
         env.top = None
+        env.angle = None
         env.front = None
 
         with mock.patch(
@@ -235,33 +237,9 @@ class EvalVLAHeadlessTest(unittest.TestCase):
         self.assertIsNotNone(env.r_vis)
         self.assertIsNotNone(env.l_vis)
         self.assertIsNotNone(env.top)
+        self.assertIsNotNone(env.angle)
         self.assertIsNotNone(env.front)
 
-    def test_dual_diana_step_refreshes_obs_after_physics_step(self):
-        env = DualDianaMed.__new__(DualDianaMed)
-        env.compute_qpos = np.zeros(16)
-        env.interpolator_left = None
-        env.interpolator_right = None
-        env.control_timestep = 0.001
-        env.model_timestep = 0.001
-        env.base_time = 0.0
-        events = []
-
-        def fake_get_obs():
-            events.append("obs")
-            return {"images": {}, "qpos": np.zeros(16, dtype=np.float32)}
-
-        env._get_obs = fake_get_obs
-
-        with mock.patch(
-            "roboimi.envs.double_base.MujocoEnv.step",
-            autospec=True,
-            side_effect=lambda _self, _action: events.append("physics"),
-        ):
-            env.step(np.zeros(16))
-
-        self.assertEqual(events, ["physics", "obs"])
-
     def test_eval_main_headless_skips_render_and_still_executes_policy(self):
         fake_env = _FakeEnv()
         fake_agent = _FakeAgent()
@@ -312,74 +290,6 @@ class EvalVLAHeadlessTest(unittest.TestCase):
         self.assertIsNotNone(fake_agent.last_observation)
         self.assertIn("front", fake_agent.last_observation["images"])
 
-    def test_eval_main_uses_condition_encoder_resize_override(self):
-        fake_env = _FakeEnv()
-        fake_agent = _FakeAgent()
-        cfg = OmegaConf.create(
-            {
-                "agent": {
-                    "condition_encoder": {
-                        "eval_image_resize_shape": None,
-                    },
-                },
-                "data": {
-                    "image_resize_shape": [224, 224],
-                },
-                "eval": {
-                    "ckpt_path": "checkpoints/vla_model_best.pt",
-                    "num_episodes": 1,
-                    "max_timesteps": 1,
-                    "device": "cpu",
-                    "task_name": "sim_transfer",
-                    "camera_names": ["front"],
-                    "use_smoothing": False,
-                    "smooth_alpha": 0.3,
-                    "verbose_action": False,
-                    "headless": True,
-                },
-            }
-        )
-
-        with mock.patch.object(eval_vla, "load_checkpoint", return_value=(fake_agent, None)), \
-             mock.patch.object(eval_vla, "make_sim_env", return_value=fake_env), \
-             mock.patch.object(eval_vla, "sample_transfer_pose", return_value=np.array([0.1, 0.2, 0.3])), \
-             mock.patch.object(eval_vla, "execute_policy_action"), \
-             mock.patch.object(eval_vla, "tqdm", side_effect=lambda iterable, **kwargs: iterable):
-            eval_vla.main.__wrapped__(cfg)
-
-        self.assertEqual(fake_agent.observation_shapes, [(3, 8, 8)])
-
-    def test_eval_main_injects_configured_task_description_when_env_omits_task(self):
-        fake_env = _FakeEnv()
-        fake_agent = _FakeAgent()
-        cfg = OmegaConf.create(
-            {
-                "agent": {},
-                "eval": {
-                    "ckpt_path": "checkpoints/vla_model_best.pt",
-                    "num_episodes": 1,
-                    "max_timesteps": 1,
-                    "device": "cpu",
-                    "task_name": "sim_transfer",
-                    "task_description": "pick the red cube",
-                    "camera_names": ["front"],
-                    "use_smoothing": False,
-                    "smooth_alpha": 0.3,
-                    "verbose_action": False,
-                    "headless": True,
-                },
-            }
-        )
-
-        with mock.patch.object(eval_vla, "load_checkpoint", return_value=(fake_agent, None)), \
-             mock.patch.object(eval_vla, "make_sim_env", return_value=fake_env), \
-             mock.patch.object(eval_vla, "sample_transfer_pose", return_value=np.array([0.1, 0.2, 0.3])), \
-             mock.patch.object(eval_vla, "execute_policy_action"), \
-             mock.patch.object(eval_vla, "tqdm", side_effect=lambda iterable, **kwargs: iterable):
-            eval_vla.main.__wrapped__(cfg)
-
-        self.assertEqual(fake_agent.last_observation["task"], "pick the red cube")
-
     def test_run_eval_returns_average_reward_summary(self):
         reward_sequences = [
             [1.0, 2.0],
@@ -437,27 +347,6 @@ class EvalVLAHeadlessTest(unittest.TestCase):
         self.assertAlmostEqual(summary["avg_reward"], 3.75)
         self.assertEqual(summary["num_episodes"], 2)
 
-
-    def test_eval_config_exposes_num_workers_default(self):
-        eval_cfg = OmegaConf.load(Path("roboimi/vla/conf/eval/eval.yaml"))
-
-        self.assertIn("num_workers", eval_cfg)
-        self.assertEqual(eval_cfg.num_workers, 1)
-
-    def test_eval_config_exposes_cuda_devices_default(self):
-        eval_cfg = OmegaConf.load(Path("roboimi/vla/conf/eval/eval.yaml"))
-
-        self.assertIn("cuda_devices", eval_cfg)
-        self.assertIsNone(eval_cfg.cuda_devices)
-
-    def test_eval_config_exposes_parallel_timeout_defaults(self):
-        eval_cfg = OmegaConf.load(Path("roboimi/vla/conf/eval/eval.yaml"))
-
-        self.assertIn("response_timeout_s", eval_cfg)
-        self.assertIn("server_startup_timeout_s", eval_cfg)
-        self.assertEqual(eval_cfg.response_timeout_s, 300.0)
-        self.assertEqual(eval_cfg.server_startup_timeout_s, 300.0)
-
     def test_run_eval_uses_serial_path_when_num_workers_is_one(self):
         cfg = OmegaConf.create(
             {
@@ -624,69 +513,6 @@ class EvalVLAHeadlessTest(unittest.TestCase):
         with self.assertRaisesRegex(ValueError, "cuda_devices"):
             eval_vla._resolve_cuda_devices(cfg)
 
-    def test_run_eval_uses_air_insert_sampler_for_socket_peg_task(self):
-        self.assertTrue(
-            hasattr(eval_vla, "sample_air_insert_socket_peg_state"),
-            "Expected eval_vla to expose the new socket/peg reset sampler",
-        )
-
-        fake_env = _FakeEnv()
-        fake_agent = _FakeAgent()
-        sampled_task_state = {
-            "socket_pos": np.array([-0.10, 0.80, 0.47], dtype=np.float32),
-            "socket_quat": np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32),
-            "peg_pos": np.array([0.10, 0.82, 0.47], dtype=np.float32),
-            "peg_quat": np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32),
-        }
-        cfg = OmegaConf.create(
-            {
-                "agent": {},
-                "eval": {
-                    "ckpt_path": "checkpoints/vla_model_best.pt",
-                    "num_episodes": 1,
-                    "max_timesteps": 1,
-                    "device": "cpu",
-                    "task_name": "sim_air_insert_socket_peg",
-                    "camera_names": ["front"],
-                    "use_smoothing": False,
-                    "smooth_alpha": 0.3,
-                    "verbose_action": False,
-                    "headless": True,
-                },
-            }
-        )
-
-        with mock.patch.object(
-            eval_vla,
-            "load_checkpoint",
-            return_value=(fake_agent, None),
-        ), mock.patch.object(
-            eval_vla,
-            "make_sim_env",
-            return_value=fake_env,
-        ) as make_env, mock.patch.object(
-            eval_vla,
-            "sample_air_insert_socket_peg_state",
-            return_value=sampled_task_state,
-        ) as socket_peg_sampler, mock.patch.object(
-            eval_vla,
-            "sample_transfer_pose",
-            side_effect=AssertionError("sample_transfer_pose should not be used for sim_air_insert_socket_peg"),
-        ), mock.patch.object(
-            eval_vla,
-            "execute_policy_action",
-        ) as execute_policy_action, mock.patch.object(
-            eval_vla,
-            "tqdm",
-            side_effect=lambda iterable, **kwargs: iterable,
-        ):
-            eval_vla._run_eval(cfg)
-
-        make_env.assert_called_once_with("sim_air_insert_socket_peg", headless=True)
-        socket_peg_sampler.assert_called_once_with()
-        execute_policy_action.assert_called_once()
-        self.assertEqual(fake_env.reset_calls, [sampled_task_state])
-
 
 if __name__ == "__main__":
     unittest.main()

tests/test_eval_vla_rollout_artifacts.py

@@ -102,8 +102,10 @@ class EvalVLARolloutArtifactsTest(unittest.TestCase):
         self.assertIn('artifact_dir', eval_cfg)
         self.assertFalse(eval_cfg.save_summary_json)
         self.assertFalse(eval_cfg.save_trajectory_npz)
+        self.assertFalse(eval_cfg.save_trajectory_image)
         self.assertFalse(eval_cfg.record_video)
         self.assertIsNone(eval_cfg.artifact_dir)
+        self.assertIsNone(eval_cfg.trajectory_image_camera_name)
         self.assertIsNone(eval_cfg.video_camera_name)
         self.assertEqual(eval_cfg.video_fps, 30)
 
@@ -133,6 +135,8 @@ class EvalVLARolloutArtifactsTest(unittest.TestCase):
                         'artifact_dir': tmpdir,
                         'save_summary_json': True,
                         'save_trajectory_npz': True,
+                        'save_trajectory_image': True,
+                        'trajectory_image_camera_name': 'front',
                         'record_video': True,
                         'video_camera_name': 'front',
                         'video_fps': 12,
@@ -176,12 +180,14 @@ class EvalVLARolloutArtifactsTest(unittest.TestCase):
             trajectory_path = Path(artifacts['trajectory_npz'])
             summary_path = Path(artifacts['summary_json'])
             video_path = Path(artifacts['video_mp4'])
+            trajectory_image_path = Path(summary['episodes'][0]['artifact_paths']['trajectory_image'])
 
             self.assertEqual(Path(artifacts['output_dir']), Path(tmpdir))
             self.assertEqual(artifacts['video_camera_name'], 'front')
             self.assertTrue(trajectory_path.exists())
             self.assertTrue(summary_path.exists())
             self.assertTrue(video_path.exists())
+            self.assertTrue(trajectory_image_path.exists())
 
             rollout_npz = np.load(trajectory_path)
             np.testing.assert_array_equal(rollout_npz['episode_index'], np.array([0, 0]))
@@ -218,267 +224,120 @@ class EvalVLARolloutArtifactsTest(unittest.TestCase):
                 saved_summary = json.load(fh)
             self.assertEqual(saved_summary['artifacts']['trajectory_npz'], str(trajectory_path))
             self.assertEqual(saved_summary['artifacts']['video_mp4'], str(video_path))
+            self.assertEqual(
+                saved_summary['episodes'][0]['artifact_paths']['trajectory_image'],
+                str(trajectory_image_path),
+            )
             self.assertEqual(saved_summary['episode_rewards'], [3.0])
             self.assertAlmostEqual(summary['avg_reward'], 3.0)
             self.assertIn('avg_obs_read_time_ms', summary)
             self.assertIn('avg_env_step_time_ms', summary)
 
-    def test_run_eval_parallel_rejects_trajectory_and_video_exports(self):
+    def test_run_eval_exports_front_trajectory_images_without_video_dependency(self):
-        unsupported_flags = [
+        actions = [
-            "record_video",
+            np.arange(16, dtype=np.float32),
-            "save_trajectory",
+            np.arange(16, dtype=np.float32) + 10.0,
-            "save_trajectory_npz",
+            np.arange(16, dtype=np.float32) + 100.0,
+            np.arange(16, dtype=np.float32) + 110.0,
         ]
+        fake_agent = _FakeAgent(actions)
+        fake_env = _FakeEnv()
 
-        for flag_name in unsupported_flags:
+        with tempfile.TemporaryDirectory() as tmpdir:
-            with self.subTest(flag_name=flag_name):
+            cfg = OmegaConf.create(
-                cfg = OmegaConf.create(
+                {
+                    'agent': {},
+                    'eval': {
+                        'ckpt_path': 'checkpoints/vla_model_best.pt',
+                        'num_episodes': 2,
+                        'max_timesteps': 2,
+                        'device': 'cpu',
+                        'task_name': 'sim_transfer',
+                        'camera_names': ['top', 'front'],
+                        'use_smoothing': True,
+                        'smooth_alpha': 0.5,
+                        'verbose_action': False,
+                        'headless': True,
+                        'artifact_dir': tmpdir,
+                        'save_trajectory_image': True,
+                        'record_video': False,
+                    },
+                }
+            )
+
+            trajectory_image_calls = []
+
+            def fake_save_rollout_trajectory_image(
+                env,
+                output_path,
+                raw_actions,
+                camera_name,
+                *,
+                line_radius=0.004,
+                max_markers=1500,
+            ):
+                del env, line_radius, max_markers
+                trajectory_image_calls.append(
                     {
-                        "agent": {},
+                        'output_path': output_path,
-                        "eval": {
+                        'camera_name': camera_name,
-                            "ckpt_path": "checkpoints/vla_model_best.pt",
+                        'raw_actions': [np.array(action, copy=True) for action in raw_actions],
-                            "num_episodes": 2,
-                            "num_workers": 2,
-                            "max_timesteps": 1,
-                            "device": "cpu",
-                            "task_name": "sim_transfer",
-                            "camera_names": ["front"],
-                            "use_smoothing": False,
-                            "smooth_alpha": 0.3,
-                            "verbose_action": False,
-                            "headless": True,
-                            "save_artifacts": True,
-                            flag_name: True,
-                        },
                     }
                 )
-
+                if output_path is None:
-                with self.assertRaisesRegex(ValueError, flag_name):
+                    return None
-                    eval_vla._run_eval_parallel(cfg)
+                output_path = Path(output_path)
-
+                output_path.parent.mkdir(parents=True, exist_ok=True)
-    def test_run_eval_parallel_writes_merged_summary_timing_and_worker_dirs(self):
+                output_path.write_bytes(b'fake-png')
-        with tempfile.TemporaryDirectory() as tmpdir:
+                return str(output_path)
-            cfg = OmegaConf.create(
-                {
-                    "agent": {},
-                    "eval": {
-                        "ckpt_path": "checkpoints/vla_model_best.pt",
-                        "num_episodes": 3,
-                        "num_workers": 2,
-                        "max_timesteps": 1,
-                        "device": "cpu",
-                        "task_name": "sim_transfer",
-                        "camera_names": ["front"],
-                        "use_smoothing": False,
-                        "smooth_alpha": 0.3,
-                        "verbose_action": False,
-                        "headless": True,
-                        "artifact_dir": tmpdir,
-                        "save_summary_json": True,
-                        "save_timing": True,
-                    },
-                }
-            )
-
-            def fake_run_spawn_jobs(payloads, max_workers, worker_fn):
-                del max_workers, worker_fn
-                return [
-                    {
-                        "episodes": [
-                            {
-                                "episode_index": 2,
-                                "episode_reward": 3.0,
-                                "episode_max_reward": 3.0,
-                                "inference_fps": 30.0,
-                                "control_fps": 15.0,
-                            }
-                        ],
-                        "_merge_state": {
-                            "obs_read_time_ms": [3.0],
-                            "preprocess_time_ms": [1.0],
-                            "inference_time_ms": [2.0],
-                            "env_step_time_ms": [4.0],
-                            "total_time_ms": [5.0],
-                            "model_forward_flags": [True],
-                        },
-                    },
-                    {
-                        "episodes": [
-                            {
-                                "episode_index": 1,
-                                "episode_reward": 2.0,
-                                "episode_max_reward": 2.0,
-                                "inference_fps": 20.0,
-                                "control_fps": 10.0,
-                            },
-                            {
-                                "episode_index": 0,
-                                "episode_reward": 1.0,
-                                "episode_max_reward": 1.0,
-                                "inference_fps": 10.0,
-                                "control_fps": 5.0,
-                            },
-                        ],
-                        "_merge_state": {
-                            "obs_read_time_ms": [1.0, 2.0],
-                            "preprocess_time_ms": [1.0, 1.0],
-                            "inference_time_ms": [2.0, 2.0],
-                            "env_step_time_ms": [4.0, 4.0],
-                            "total_time_ms": [5.0, 5.0],
-                            "model_forward_flags": [False, True],
-                        },
-                    },
-                ]
 
             with mock.patch.object(
                 eval_vla,
-                "sample_transfer_pose",
+                'load_checkpoint',
-                side_effect=[
+                return_value=(fake_agent, None),
-                    np.array([0.1, 0.2, 0.3], dtype=np.float32),
-                    np.array([0.4, 0.5, 0.6], dtype=np.float32),
-                    np.array([0.7, 0.8, 0.9], dtype=np.float32),
-                ],
             ), mock.patch.object(
                 eval_vla,
-                "_run_spawn_jobs",
+                'make_sim_env',
-                side_effect=fake_run_spawn_jobs,
+                return_value=fake_env,
-            ):
+            ), mock.patch.object(
-                summary = eval_vla._run_eval_parallel(cfg)
+                eval_vla,
+                'sample_transfer_pose',
+                return_value=np.array([0.1, 0.2, 0.3], dtype=np.float32),
+            ), mock.patch.object(
+                eval_vla,
+                'tqdm',
+                side_effect=lambda iterable, **kwargs: iterable,
+            ), mock.patch.object(
+                eval_vla,
+                '_save_rollout_trajectory_image',
+                side_effect=fake_save_rollout_trajectory_image,
+            ) as save_trajectory_image_mock, mock.patch.object(
+                eval_vla,
+                '_open_video_writer',
+            ) as open_video_writer_mock:
+                summary = eval_vla._run_eval(cfg)
 
-            summary_path = Path(tmpdir) / "rollout_summary.json"
+            self.assertEqual(save_trajectory_image_mock.call_count, 2)
-            timing_path = Path(tmpdir) / "timing.json"
+            open_video_writer_mock.assert_not_called()
-            worker_00_dir = Path(tmpdir) / "workers" / "worker_00"
+            self.assertIsNone(summary['artifacts']['video_mp4'])
-            worker_01_dir = Path(tmpdir) / "workers" / "worker_01"
+            self.assertEqual(summary['artifacts']['trajectory_image_camera_name'], 'front')
-
+            self.assertEqual(
-            self.assertTrue(summary_path.exists())
+                [call['camera_name'] for call in trajectory_image_calls],
-            self.assertTrue(timing_path.exists())
+                ['front', 'front'],
-            self.assertTrue(worker_00_dir.is_dir())
-            self.assertTrue(worker_01_dir.is_dir())
-            self.assertEqual(summary["episode_rewards"], [1.0, 2.0, 3.0])
-
-            with summary_path.open("r", encoding="utf-8") as fh:
-                saved_summary = json.load(fh)
-            with timing_path.open("r", encoding="utf-8") as fh:
-                saved_timing = json.load(fh)
-
-            self.assertEqual(saved_summary["episode_rewards"], [1.0, 2.0, 3.0])
-            self.assertEqual(saved_summary["artifact_dir"], tmpdir)
-            self.assertEqual(saved_timing["count"], 3)
-            self.assertEqual(saved_timing["model_forward_count"], 2)
-
-    def test_run_eval_parallel_cuda_writes_merged_summary_timing_and_worker_dirs(self):
-        with tempfile.TemporaryDirectory() as tmpdir:
-            cfg = OmegaConf.create(
-                {
-                    "agent": {},
-                    "eval": {
-                        "ckpt_path": "checkpoints/vla_model_best.pt",
-                        "num_episodes": 3,
-                        "num_workers": 2,
-                        "cuda_devices": [0],
-                        "max_timesteps": 1,
-                        "device": "cuda",
-                        "task_name": "sim_transfer",
-                        "camera_names": ["front"],
-                        "use_smoothing": False,
-                        "smooth_alpha": 0.3,
-                        "verbose_action": False,
-                        "headless": True,
-                        "artifact_dir": tmpdir,
-                        "save_summary_json": True,
-                        "save_timing": True,
-                    },
-                }
             )
 
-            def fake_run_cuda_parallel_processes(server_payloads, worker_payloads):
+            first_episode_path = Path(summary['episodes'][0]['artifact_paths']['trajectory_image'])
-                self.assertEqual(len(server_payloads), 1)
+            second_episode_path = Path(summary['episodes'][1]['artifact_paths']['trajectory_image'])
-                self.assertEqual(server_payloads[0]["device_index"], 0)
+            self.assertTrue(first_episode_path.exists())
-                self.assertEqual([payload["server_index"] for payload in worker_payloads], [0, 0])
+            self.assertTrue(second_episode_path.exists())
-                return [
+            self.assertNotEqual(first_episode_path, second_episode_path)
-                    {
+            self.assertEqual(first_episode_path.parent, Path(tmpdir))
-                        "episodes": [
+            self.assertEqual(second_episode_path.parent, Path(tmpdir))
-                            {
-                                "episode_index": 2,
-                                "episode_reward": 3.0,
-                                "episode_max_reward": 3.0,
-                                "inference_fps": 30.0,
-                                "control_fps": 15.0,
-                            }
-                        ],
-                        "_merge_state": {
-                            "obs_read_time_ms": [3.0],
-                            "preprocess_time_ms": [1.0],
-                            "inference_time_ms": [2.0],
-                            "env_step_time_ms": [4.0],
-                            "total_time_ms": [5.0],
-                            "model_forward_flags": [True],
-                        },
-                    },
-                    {
-                        "episodes": [
-                            {
-                                "episode_index": 1,
-                                "episode_reward": 2.0,
-                                "episode_max_reward": 2.0,
-                                "inference_fps": 20.0,
-                                "control_fps": 10.0,
-                            },
-                            {
-                                "episode_index": 0,
-                                "episode_reward": 1.0,
-                                "episode_max_reward": 1.0,
-                                "inference_fps": 10.0,
-                                "control_fps": 5.0,
-                            },
-                        ],
-                        "_merge_state": {
-                            "obs_read_time_ms": [1.0, 2.0],
-                            "preprocess_time_ms": [1.0, 1.0],
-                            "inference_time_ms": [2.0, 2.0],
-                            "env_step_time_ms": [4.0, 4.0],
-                            "total_time_ms": [5.0, 5.0],
-                            "model_forward_flags": [False, True],
-                        },
-                    },
-                ]
 
-            with mock.patch.object(
+            np.testing.assert_array_equal(trajectory_image_calls[0]['raw_actions'][0], actions[0])
-                eval_vla,
+            np.testing.assert_array_equal(trajectory_image_calls[0]['raw_actions'][1], actions[1])
-                "sample_transfer_pose",
+            np.testing.assert_array_equal(trajectory_image_calls[1]['raw_actions'][0], actions[2])
-                side_effect=[
+            np.testing.assert_array_equal(trajectory_image_calls[1]['raw_actions'][1], actions[3])
-                    np.array([0.1, 0.2, 0.3], dtype=np.float32),
-                    np.array([0.4, 0.5, 0.6], dtype=np.float32),
-                    np.array([0.7, 0.8, 0.9], dtype=np.float32),
-                ],
-            ), mock.patch.object(
-                eval_vla,
-                "_run_cuda_parallel_processes",
-                side_effect=fake_run_cuda_parallel_processes,
-                create=True,
-            ):
-                summary = eval_vla._run_eval_parallel_cuda(cfg)
-
-            summary_path = Path(tmpdir) / "rollout_summary.json"
-            timing_path = Path(tmpdir) / "timing.json"
-            worker_00_dir = Path(tmpdir) / "workers" / "worker_00"
-            worker_01_dir = Path(tmpdir) / "workers" / "worker_01"
-
-            self.assertTrue(summary_path.exists())
-            self.assertTrue(timing_path.exists())
-            self.assertTrue(worker_00_dir.is_dir())
-            self.assertTrue(worker_01_dir.is_dir())
-            self.assertEqual(summary["episode_rewards"], [1.0, 2.0, 3.0])
-
-            with summary_path.open("r", encoding="utf-8") as fh:
-                saved_summary = json.load(fh)
-            with timing_path.open("r", encoding="utf-8") as fh:
-                saved_timing = json.load(fh)
-
-            self.assertEqual(saved_summary["episode_rewards"], [1.0, 2.0, 3.0])
-            self.assertEqual(saved_summary["artifact_dir"], tmpdir)
-            self.assertEqual(saved_timing["count"], 3)
-            self.assertEqual(saved_timing["model_forward_count"], 2)
 
 
 if __name__ == '__main__':

tests/test_imf_vla_agent.py

@@ -376,6 +376,49 @@ class _ForbiddenScheduler:
         raise AssertionError('IMF inference should not use DDIM scheduler step')
 
 
+class _StubFutureTokenPredictor(nn.Module):
+    def __init__(self, num_future_tokens=1):
+        super().__init__()
+        self.num_future_tokens = int(num_future_tokens)
+        self.calls = []
+
+    def forward(self, history_tokens):
+        self.calls.append(history_tokens.detach().clone())
+        summary = history_tokens.mean(dim=1, keepdim=True)
+        return summary.repeat(1, self.num_future_tokens, 1)
+
+
+class _RecordingDirectFutureDecoder(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.scale = nn.Parameter(torch.tensor(0.5))
+        self.calls = []
+
+    def forward(self, sample, r, t, cond=None):
+        record = {
+            'sample': sample.detach().clone(),
+            'r': r.detach().clone(),
+            't': t.detach().clone(),
+            'cond': None if cond is None else cond.detach().clone(),
+        }
+        self.calls.append(record)
+        cond_term = 0.0
+        if cond is not None:
+            cond_term = cond.mean(dim=1, keepdim=True)
+        return self.scale * sample + cond_term
+
+
+class _RecordingSigReg(nn.Module):
+    def __init__(self, value=0.5):
+        super().__init__()
+        self.value = float(value)
+        self.calls = []
+
+    def forward(self, embeddings):
+        self.calls.append(embeddings.detach().clone())
+        return embeddings.new_tensor(self.value)
+
+
 def _make_images(batch_size, obs_horizon, per_camera_fill):
     return {
         name: torch.full((batch_size, obs_horizon, 1, 2, 2), fill_value=value, dtype=torch.float32)
@@ -501,6 +544,311 @@ class IMFVLAAgentTest(unittest.TestCase):
         self.assertTrue(torch.allclose(head.calls[0]['t'], torch.ones(2)))
         self.assertTrue(torch.allclose(head.calls[0]['cond'], expected_cond))
 
+    def test_predict_action_appends_lewm_future_tokens_to_history_conditioning(self):
+        agent_cls, agent_module = _load_imf_agent_class()
+        head = _RecordingLinearIMFHead()
+        future_predictor = _StubFutureTokenPredictor(num_future_tokens=1)
+        agent = agent_cls(
+            vision_backbone=_StubVisionBackbone(),
+            state_encoder=nn.Identity(),
+            action_encoder=nn.Identity(),
+            head=head,
+            action_dim=2,
+            obs_dim=1,
+            pred_horizon=3,
+            obs_horizon=2,
+            diffusion_steps=10,
+            inference_steps=1,
+            num_cams=len(_CAMERA_NAMES),
+            camera_names=_CAMERA_NAMES,
+            num_action_steps=2,
+            head_type='transformer',
+            extra_condition_tokens=1,
+            lewm_history_horizon=3,
+            lewm_query_offsets=[8],
+            lewm_predictor=future_predictor,
+            lewm_pred_projector=nn.Identity(),
+            lewm_loss_weight=0.5,
+        )
+        agent.infer_scheduler = _ForbiddenScheduler()
+
+        images = _make_images(
+            batch_size=1,
+            obs_horizon=2,
+            per_camera_fill={'r_vis': 10.0, 'top': 20.0, 'front': 30.0},
+        )
+        qpos = torch.tensor([[[1.0], [2.0]]], dtype=torch.float32)
+        lewm_images = _make_images(
+            batch_size=1,
+            obs_horizon=3,
+            per_camera_fill={'r_vis': 10.0, 'top': 20.0, 'front': 30.0},
+        )
+        lewm_qpos = torch.tensor([[[0.5], [1.5], [2.5]]], dtype=torch.float32)
+        initial_noise = torch.tensor(
+            [[[1.0, -1.0], [0.0, 2.0], [3.0, -2.0]]],
+            dtype=torch.float32,
+        )
+
+        with mock.patch.object(agent_module.torch, 'randn', return_value=initial_noise):
+            _ = agent.predict_action(
+                images,
+                qpos,
+                lewm_images=lewm_images,
+                lewm_proprioception=lewm_qpos,
+            )
+
+        expected_history = torch.tensor(
+            [[[10.0, 20.0, 30.0, 1.0], [10.0, 20.0, 30.0, 2.0]]],
+            dtype=torch.float32,
+        )
+        expected_future = torch.tensor([[[10.0, 20.0, 30.0, 1.5]]], dtype=torch.float32)
+        expected_cond = torch.cat([expected_history, expected_future], dim=1)
+
+        self.assertEqual(agent.condition_sequence_length, 3)
+        self.assertEqual(agent.per_step_cond_dim, 4)
+        self.assertEqual(len(head.calls), 1)
+        self.assertTrue(torch.allclose(head.calls[0]['cond'], expected_cond))
+        self.assertEqual(len(future_predictor.calls), 1)
+
+    def test_compute_loss_tracks_action_and_lewm_loss_breakdown(self):
+        agent_cls, agent_module = _load_imf_agent_class()
+        head = _RecordingLinearIMFHead()
+        future_predictor = _StubFutureTokenPredictor(num_future_tokens=1)
+        sigreg = _RecordingSigReg(value=0.75)
+        agent = agent_cls(
+            vision_backbone=_StubVisionBackbone(),
+            state_encoder=nn.Identity(),
+            action_encoder=nn.Identity(),
+            head=head,
+            action_dim=2,
+            obs_dim=1,
+            pred_horizon=3,
+            obs_horizon=2,
+            diffusion_steps=10,
+            inference_steps=1,
+            num_cams=len(_CAMERA_NAMES),
+            camera_names=_CAMERA_NAMES,
+            num_action_steps=2,
+            head_type='transformer',
+            extra_condition_tokens=1,
+            lewm_history_horizon=3,
+            lewm_query_offsets=[8],
+            lewm_predictor=future_predictor,
+            lewm_pred_projector=nn.Identity(),
+            lewm_sigreg=sigreg,
+            lewm_sigreg_weight=0.09,
+            lewm_loss_weight=0.25,
+        )
+
+        images = _make_images(
+            batch_size=1,
+            obs_horizon=2,
+            per_camera_fill={'r_vis': 1.0, 'top': 2.0, 'front': 3.0},
+        )
+        qpos = torch.tensor([[[0.25], [0.75]]], dtype=torch.float32)
+        actions = torch.tensor(
+            [[[1.0, -1.0], [0.5, 0.25], [-0.5, 1.5]]],
+            dtype=torch.float32,
+        )
+        lewm_images = _make_images(
+            batch_size=1,
+            obs_horizon=3,
+            per_camera_fill={'r_vis': 1.0, 'top': 2.0, 'front': 3.0},
+        )
+        lewm_qpos = torch.tensor([[[0.1], [0.2], [0.3]]], dtype=torch.float32)
+        lewm_future_images = _make_images(
+            batch_size=1,
+            obs_horizon=1,
+            per_camera_fill={'r_vis': 1.0, 'top': 2.0, 'front': 3.0},
+        )
+        lewm_future_qpos = torch.tensor([[[0.4]]], dtype=torch.float32)
+        noise = torch.tensor(
+            [[[0.2, -0.4], [0.1, 0.3], [0.5, -0.2]]],
+            dtype=torch.float32,
+        )
+        t_sample = torch.tensor([0.8], dtype=torch.float32)
+        r_sample = torch.tensor([0.25], dtype=torch.float32)
+
+        with mock.patch.object(agent_module.torch, 'randn_like', return_value=noise), \
+             mock.patch.object(agent_module.torch, 'rand', side_effect=[t_sample, r_sample]):
+            loss = agent.compute_loss(
+                {
+                    'images': images,
+                    'qpos': qpos,
+                    'action': actions,
+                    'lewm_images': lewm_images,
+                    'lewm_qpos': lewm_qpos,
+                    'lewm_future_images': lewm_future_images,
+                    'lewm_future_qpos': lewm_future_qpos,
+                }
+            )
+
+        metrics = agent.get_last_loss_breakdown()
+        self.assertAlmostEqual(loss.item(), metrics['loss'], places=6)
+        self.assertIn('action_loss', metrics)
+        self.assertIn('lewm_pred_loss', metrics)
+        self.assertIn('lewm_sigreg_loss', metrics)
+        self.assertIn('lewm_loss', metrics)
+        self.assertAlmostEqual(metrics['lewm_sigreg_loss'], 0.75, places=6)
+        self.assertAlmostEqual(
+            metrics['lewm_loss'],
+            metrics['lewm_pred_loss'] + 0.09 * metrics['lewm_sigreg_loss'],
+            places=5,
+        )
+        self.assertAlmostEqual(
+            metrics['loss'],
+            metrics['action_loss'] + 0.25 * metrics['lewm_loss'],
+            places=5,
+        )
+        self.assertEqual(len(sigreg.calls), 1)
+        expected_lewm_history = torch.tensor(
+            [[[1.0, 2.0, 3.0, 0.1], [1.0, 2.0, 3.0, 0.2], [1.0, 2.0, 3.0, 0.3]]],
+            dtype=torch.float32,
+        )
+        torch.testing.assert_close(sigreg.calls[0], expected_lewm_history.transpose(0, 1))
+
+    def test_predict_action_with_dual_decoder_keeps_action_condition_history_only(self):
+        agent_cls, agent_module = _load_imf_agent_class()
+        head = _RecordingLinearIMFHead()
+        future_decoder = _RecordingDirectFutureDecoder()
+        agent = agent_cls(
+            vision_backbone=_StubVisionBackbone(),
+            state_encoder=nn.Identity(),
+            action_encoder=nn.Identity(),
+            head=head,
+            future_decoder=future_decoder,
+            action_dim=2,
+            obs_dim=1,
+            pred_horizon=3,
+            obs_horizon=2,
+            diffusion_steps=10,
+            inference_steps=1,
+            num_cams=len(_CAMERA_NAMES),
+            camera_names=_CAMERA_NAMES,
+            num_action_steps=2,
+            head_type='transformer',
+            lewm_history_horizon=3,
+            lewm_query_offsets=[8],
+            lewm_loss_weight=1.0,
+        )
+        agent.infer_scheduler = _ForbiddenScheduler()
+        with torch.no_grad():
+            agent.future_query_tokens.copy_(torch.tensor([[[0.1, 0.2, 0.3, 0.4]]], dtype=torch.float32))
+
+        images = _make_images(
+            batch_size=1,
+            obs_horizon=2,
+            per_camera_fill={'r_vis': 10.0, 'top': 20.0, 'front': 30.0},
+        )
+        qpos = torch.tensor([[[1.0], [2.0]]], dtype=torch.float32)
+        initial_noise = torch.tensor(
+            [[[1.0, -1.0], [0.0, 2.0], [3.0, -2.0]]],
+            dtype=torch.float32,
+        )
+
+        with mock.patch.object(agent_module.torch, 'randn', return_value=initial_noise):
+            _ = agent.predict_action(images, qpos)
+
+        expected_history = torch.tensor(
+            [[[10.0, 20.0, 30.0, 1.0], [10.0, 20.0, 30.0, 2.0]]],
+            dtype=torch.float32,
+        )
+        self.assertEqual(len(head.calls), 1)
+        self.assertTrue(torch.allclose(head.calls[0]['cond'], expected_history))
+        self.assertEqual(len(future_decoder.calls), 0)
+
+    def test_compute_loss_with_dual_decoder_tracks_lewm_loss_breakdown(self):
+        agent_cls, agent_module = _load_imf_agent_class()
+        head = _RecordingLinearIMFHead()
+        future_decoder = _RecordingDirectFutureDecoder()
+        sigreg = _RecordingSigReg(value=0.75)
+        agent = agent_cls(
+            vision_backbone=_StubVisionBackbone(),
+            state_encoder=nn.Identity(),
+            action_encoder=nn.Identity(),
+            head=head,
+            future_decoder=future_decoder,
+            action_dim=2,
+            obs_dim=1,
+            pred_horizon=3,
+            obs_horizon=2,
+            diffusion_steps=10,
+            inference_steps=1,
+            num_cams=len(_CAMERA_NAMES),
+            camera_names=_CAMERA_NAMES,
+            num_action_steps=2,
+            head_type='transformer',
+            lewm_history_horizon=3,
+            lewm_query_offsets=[8],
+            lewm_sigreg=sigreg,
+            lewm_sigreg_weight=0.09,
+            lewm_loss_weight=1.0,
+        )
+        with torch.no_grad():
+            agent.future_query_tokens.copy_(torch.tensor([[[0.2, 0.4, 0.6, 0.8]]], dtype=torch.float32))
+
+        images = _make_images(
+            batch_size=1,
+            obs_horizon=2,
+            per_camera_fill={'r_vis': 1.0, 'top': 2.0, 'front': 3.0},
+        )
+        qpos = torch.tensor([[[0.25], [0.75]]], dtype=torch.float32)
+        actions = torch.tensor(
+            [[[1.0, -1.0], [0.5, 0.25], [-0.5, 1.5]]],
+            dtype=torch.float32,
+        )
+        lewm_images = _make_images(
+            batch_size=1,
+            obs_horizon=3,
+            per_camera_fill={'r_vis': 1.0, 'top': 2.0, 'front': 3.0},
+        )
+        lewm_qpos = torch.tensor([[[0.1], [0.2], [0.3]]], dtype=torch.float32)
+        lewm_future_images = _make_images(
+            batch_size=1,
+            obs_horizon=1,
+            per_camera_fill={'r_vis': 1.0, 'top': 2.0, 'front': 3.0},
+        )
+        lewm_future_qpos = torch.tensor([[[0.4]]], dtype=torch.float32)
+        noise = torch.tensor(
+            [[[0.2, -0.4], [0.1, 0.3], [0.5, -0.2]]],
+            dtype=torch.float32,
+        )
+        t_sample = torch.tensor([0.8], dtype=torch.float32)
+        r_sample = torch.tensor([0.25], dtype=torch.float32)
+
+        with mock.patch.object(agent_module.torch, 'randn_like', return_value=noise), \
+             mock.patch.object(agent_module.torch, 'rand', side_effect=[t_sample, r_sample]):
+            loss = agent.compute_loss(
+                {
+                    'images': images,
+                    'qpos': qpos,
+                    'action': actions,
+                    'lewm_images': lewm_images,
+                    'lewm_qpos': lewm_qpos,
+                    'lewm_future_images': lewm_future_images,
+                    'lewm_future_qpos': lewm_future_qpos,
+                }
+            )
+
+        metrics = agent.get_last_loss_breakdown()
+        self.assertAlmostEqual(loss.item(), metrics['loss'], places=6)
+        self.assertEqual(len(head.calls), 2)
+        self.assertEqual(head.calls[0]['cond'].shape, (1, 2, 4))
+        self.assertEqual(len(future_decoder.calls), 1)
+        self.assertEqual(future_decoder.calls[0]['cond'].shape, (1, 3, 4))
+        self.assertAlmostEqual(
+            metrics['loss'],
+            metrics['action_loss'] + metrics['lewm_loss'],
+            places=5,
+        )
+        self.assertAlmostEqual(
+            metrics['lewm_loss'],
+            metrics['lewm_pred_loss'] + 0.09 * metrics['lewm_sigreg_loss'],
+            places=5,
+        )
+        self.assertGreater(metrics['lewm_pred_loss'], 0.0)
+        self.assertAlmostEqual(metrics['lewm_sigreg_loss'], 0.75, places=6)
+
     def test_select_action_only_regenerates_when_action_queue_is_empty(self):
         agent, _head, _agent_module = self._make_agent(pred_horizon=4, obs_horizon=2, num_action_steps=2)
         observation = {
@@ -851,6 +1199,80 @@ class IMFVLAAgentTest(unittest.TestCase):
         self.assertEqual(agent.vision_encoder.output_dim, 96)
         self.assertEqual(agent.vision_encoder.eval_image_resize_shape, (256, 256))
 
+    def test_hydra_config_instantiates_lewm_resnet_query_imf_attnres_with_future_tokens(self):
+        cfg = _compose_cfg(
+            overrides=[
+                'agent=lewm_resnet_query_imf_attnres',
+                'agent.head.n_layer=1',
+                'agent.head.n_emb=16',
+                'agent.lewm_query_offsets=[8]',
+            ]
+        )
+
+        self.assertEqual(cfg.agent._target_, 'roboimi.vla.agent_imf.IMFVLAAgent')
+        self.assertEqual(
+            cfg.agent.vision_backbone._target_,
+            'roboimi.vla.models.backbones.lewm_resnet_query_fusion.LeWMMultiViewResNetBackbone',
+        )
+        self.assertEqual(
+            cfg.agent.state_encoder._target_,
+            'roboimi.vla.modules.encoders.LeWMStateEncoder',
+        )
+        self.assertEqual(cfg.agent.head.cond_dim, 288)
+        self.assertEqual(cfg.agent.cond_projector.output_dim, 288)
+        self.assertEqual(cfg.agent.extra_condition_tokens, 1)
+        self.assertEqual(
+            cfg.agent.lewm_sigreg._target_,
+            'roboimi.vla.models.backbones.lewm_resnet_query_fusion.SIGReg',
+        )
+        self.assertAlmostEqual(cfg.agent.lewm_sigreg_weight, 0.09)
+
+        with _stub_optional_modules(include_imf_head=True):
+            agent = instantiate(cfg.agent)
+
+        self.assertEqual(agent.per_step_cond_dim, 288)
+        self.assertEqual(agent.condition_sequence_length, agent.obs_horizon + 1)
+        self.assertEqual(agent.noise_pred_net.constructor_kwargs['cond_dim'], 288)
+        self.assertEqual(
+            agent.noise_pred_net.constructor_kwargs['n_obs_steps'],
+            agent.condition_sequence_length,
+        )
+        self.assertIsNotNone(agent.lewm_sigreg)
+
+    def test_hydra_config_instantiates_lewm_resnet_dual_decoder_imf_attnres(self):
+        cfg = _compose_cfg(
+            overrides=[
+                'agent=lewm_resnet_dual_decoder_imf_attnres',
+                'agent.head.n_layer=1',
+                'agent.head.n_emb=16',
+                'agent.future_decoder.n_layer=1',
+                'agent.future_decoder.n_emb=16',
+                'agent.lewm_query_offsets=[8]',
+            ]
+        )
+
+        self.assertEqual(cfg.agent._target_, 'roboimi.vla.agent_imf.IMFVLAAgent')
+        self.assertEqual(cfg.agent.extra_condition_tokens, 0)
+        self.assertEqual(
+            cfg.agent.future_decoder._target_,
+            'roboimi.vla.models.heads.imf_transformer1d.IMFTransformer1D',
+        )
+        self.assertEqual(cfg.agent.head.cond_dim, 288)
+        self.assertEqual(cfg.agent.future_decoder.cond_dim, 288)
+
+        with _stub_optional_modules(include_imf_head=True):
+            agent = instantiate(cfg.agent)
+
+        self.assertEqual(agent.per_step_cond_dim, 288)
+        self.assertEqual(agent.condition_sequence_length, agent.obs_horizon)
+        self.assertEqual(agent.noise_pred_net.constructor_kwargs['n_obs_steps'], agent.obs_horizon)
+        self.assertEqual(agent.future_decoder.constructor_kwargs['cond_dim'], 288)
+        self.assertEqual(
+            agent.future_decoder.constructor_kwargs['n_obs_steps'],
+            agent.lewm_history_horizon,
+        )
+        self.assertEqual(agent.future_query_tokens.shape, (1, 1, 288))
+
 
     def test_hydra_config_instantiates_resnet_imf_attnres_multitoken_with_sequence_length_three_times_obs_horizon(self):
         cfg = _compose_cfg(

tests/test_robot_asset_paths.py

@@ -4,7 +4,7 @@ import unittest
 from pathlib import Path
 from unittest import mock
 
-from roboimi.assets.robots import diana_med
+from roboimi.assets.robots.diana_med import BiDianaMed
 
 
 class _FakeKDL:
@@ -24,7 +24,6 @@ class RobotAssetPathResolutionTest(unittest.TestCase):
         _FakeKDL.reset_calls = []
 
     def test_bidianamed_resolves_robot_asset_paths_independent_of_cwd(self):
-        BiDianaMed = diana_med.BiDianaMed
         repo_root = Path(__file__).resolve().parents[1]
         expected_xml = repo_root / 'roboimi/assets/models/manipulators/DianaMed/bi_diana_transfer_ee.xml'
         expected_urdf = repo_root / 'roboimi/assets/models/manipulators/DianaMed/DualDianaMed.urdf'
@@ -59,47 +58,6 @@ class RobotAssetPathResolutionTest(unittest.TestCase):
         self.assertEqual({Path(path) for path in _FakeKDL.init_calls}, {expected_urdf})
         self.assertTrue(all(Path(path).is_absolute() for path in _FakeKDL.init_calls))
 
-    def test_bidianamed_socket_peg_resolves_robot_asset_paths_independent_of_cwd(self):
-        BiDianaMedSocketPeg = getattr(diana_med, 'BiDianaMedSocketPeg', None)
-        self.assertIsNotNone(
-            BiDianaMedSocketPeg,
-            'Expected roboimi.assets.robots.diana_med.BiDianaMedSocketPeg',
-        )
-
-        repo_root = Path(__file__).resolve().parents[1]
-        expected_xml = repo_root / 'roboimi/assets/models/manipulators/DianaMed/bi_diana_socket_peg_ee.xml'
-        expected_urdf = repo_root / 'roboimi/assets/models/manipulators/DianaMed/DualDianaMed.urdf'
-        xml_calls = []
-
-        def fake_from_xml_path(*, filename, assets=None):
-            xml_calls.append((filename, assets))
-            return object()
-
-        with tempfile.TemporaryDirectory() as tempdir:
-            previous_cwd = os.getcwd()
-            try:
-                os.chdir(tempdir)
-                with mock.patch(
-                    'roboimi.assets.robots.arm_base.mujoco.MjModel.from_xml_path',
-                    side_effect=fake_from_xml_path,
-                ), mock.patch(
-                    'roboimi.assets.robots.arm_base.mujoco.MjData',
-                    return_value=object(),
-                ), mock.patch(
-                    'roboimi.assets.robots.arm_base.KDL_utils',
-                    _FakeKDL,
-                ):
-                    BiDianaMedSocketPeg()
-            finally:
-                os.chdir(previous_cwd)
-
-        self.assertEqual(len(xml_calls), 1)
-        self.assertEqual(Path(xml_calls[0][0]), expected_xml)
-        self.assertTrue(Path(xml_calls[0][0]).is_absolute())
-        self.assertGreaterEqual(len(_FakeKDL.init_calls), 2)
-        self.assertEqual({Path(path) for path in _FakeKDL.init_calls}, {expected_urdf})
-        self.assertTrue(all(Path(path).is_absolute() for path in _FakeKDL.init_calls))
-
 
 if __name__ == '__main__':
     unittest.main()

tests/test_simple_robot_dataset_episode_filter.py

@@ -1,42 +0,0 @@
-import tempfile
-import unittest
-from pathlib import Path
-
-import h5py
-import numpy as np
-
-
-def _write_episode(path: Path, length: int = 2):
-    with h5py.File(path, 'w') as f:
-        f.create_dataset('action', data=np.zeros((length, 16), dtype=np.float32))
-        obs = f.create_group('observations')
-        obs.create_dataset('qpos', data=np.zeros((length, 16), dtype=np.float32))
-        images = obs.create_group('images')
-        for cam_name in ('l_vis', 'r_vis', 'front'):
-            images.create_dataset(cam_name, data=np.zeros((length, 4, 4, 3), dtype=np.uint8))
-
-
-class SimpleRobotDatasetEpisodeFilterTest(unittest.TestCase):
-    def test_filters_by_original_episode_indices_and_exposes_available_indices(self):
-        from roboimi.vla.data.simpe_robot_dataset import SimpleRobotDataset
-
-        with tempfile.TemporaryDirectory() as tmpdir:
-            root = Path(tmpdir)
-            _write_episode(root / 'episode_0.hdf5')
-            _write_episode(root / 'episode_2.hdf5')
-            _write_episode(root / 'episode_10.hdf5')
-
-            dataset = SimpleRobotDataset(
-                root,
-                camera_names=['l_vis', 'r_vis', 'front'],
-                image_resize_shape=None,
-                episode_indices=[10, 2],
-            )
-
-            self.assertEqual(dataset.available_episode_indices, [2, 10])
-            self.assertEqual(set(dataset.episodes.keys()), {2, 10})
-            self.assertEqual(len(dataset), 4)
-
-
-if __name__ == '__main__':
-    unittest.main()

tests/test_simple_robot_dataset_image_loading.py

@@ -12,18 +12,21 @@ from roboimi.vla.data.simpe_robot_dataset import SimpleRobotDataset
 
 
 class SimpleRobotDatasetImageLoadingTest(unittest.TestCase):
-    def _write_episode(self, dataset_dir: Path) -> None:
+    def _write_episode(self, dataset_dir: Path, episode_idx: int = 0, *, base_value: float = 0.0) -> None:
-        episode_path = dataset_dir / "episode_0.hdf5"
+        episode_path = dataset_dir / f"episode_{episode_idx}.hdf5"
         with h5py.File(episode_path, "w") as root:
-            root.create_dataset("action", data=np.arange(8, dtype=np.float32).reshape(4, 2))
+            root.create_dataset(
+                "action",
+                data=(np.arange(8, dtype=np.float32).reshape(4, 2) + base_value),
+            )
             root.create_dataset(
                 "observations/qpos",
-                data=np.arange(16, dtype=np.float32).reshape(4, 4),
+                data=(np.arange(16, dtype=np.float32).reshape(4, 4) + base_value),
             )
             root.create_dataset("task", data=np.array([b"sim_transfer"]))
             root.create_dataset(
                 "observations/images/front",
-                data=np.arange(4 * 8 * 8 * 3, dtype=np.uint8).reshape(4, 8, 8, 3),
+                data=((np.arange(4 * 8 * 8 * 3, dtype=np.uint8) + int(base_value)) % 255).reshape(4, 8, 8, 3),
             )
 
     def test_getitem_only_resizes_observation_horizon_images(self):
@@ -79,3 +82,46 @@ class SimpleRobotDatasetImageLoadingTest(unittest.TestCase):
 
         fake_cv2.resize.assert_not_called()
         self.assertEqual(tuple(sample["observation.front"].shape), (2, 3, 8, 8))
+
+    def test_getitem_can_emit_lewm_history_and_future_observations(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            dataset_dir = Path(tmpdir)
+            self._write_episode(dataset_dir)
+            dataset = SimpleRobotDataset(
+                dataset_dir,
+                obs_horizon=2,
+                pred_horizon=3,
+                camera_names=["front"],
+                image_resize_shape=None,
+                lewm_history_horizon=3,
+                lewm_query_offsets=[1, 2],
+            )
+
+            sample = dataset[1]
+
+        self.assertEqual(tuple(sample["lewm.observation.state"].shape), (3, 4))
+        self.assertEqual(tuple(sample["lewm.observation.front"].shape), (3, 3, 8, 8))
+        self.assertEqual(tuple(sample["lewm.future.state"].shape), (2, 4))
+        self.assertEqual(tuple(sample["lewm.future.front"].shape), (2, 3, 8, 8))
+
+    def test_dataset_can_limit_loading_to_specific_episode_indices(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            dataset_dir = Path(tmpdir)
+            self._write_episode(dataset_dir, episode_idx=0, base_value=0.0)
+            self._write_episode(dataset_dir, episode_idx=1, base_value=100.0)
+
+            dataset = SimpleRobotDataset(
+                dataset_dir,
+                obs_horizon=2,
+                pred_horizon=3,
+                camera_names=["front"],
+                image_resize_shape=None,
+                episode_indices=[1],
+            )
+
+            sample = dataset[0]
+
+        self.assertEqual(len(dataset.hdf5_files), 1)
+        self.assertEqual(dataset.available_episode_indices, [1])
+        self.assertEqual(len(dataset), 4)
+        self.assertTrue(np.allclose(sample["observation.state"][0].numpy(), np.array([100.0, 101.0, 102.0, 103.0])))

tests/test_smolvla_imf_agent.py

@@ -1,311 +0,0 @@
-import contextlib
-import importlib
-import importlib.machinery
-import sys
-import types
-import unittest
-from pathlib import Path
-from unittest import mock
-
-import torch
-from hydra import compose, initialize_config_dir
-from hydra.core.global_hydra import GlobalHydra
-from hydra.utils import instantiate
-from omegaconf import OmegaConf
-from torch import nn
-
-
-_REPO_ROOT = Path(__file__).resolve().parents[1]
-_CONFIG_DIR = str((_REPO_ROOT / 'roboimi/vla/conf').resolve())
-_CAMERA_NAMES = ('r_vis', 'top', 'front')
-_MISSING = object()
-
-
-class _RecordingHead(nn.Module):
-    def __init__(self):
-        super().__init__()
-        self.scale = nn.Parameter(torch.tensor(0.5))
-        self.calls = []
-
-    @staticmethod
-    def _broadcast(value, reference):
-        while value.ndim < reference.ndim:
-            value = value.unsqueeze(-1)
-        return value
-
-    def forward(self, sample, r, t, cond=None):
-        self.calls.append({
-            'sample': sample.detach().clone(),
-            'r': r.detach().clone(),
-            't': t.detach().clone(),
-            'cond': None if cond is None else cond.detach().clone(),
-        })
-        cond_term = 0.0 if cond is None else cond.mean(dim=(1, 2), keepdim=True)
-        return self.scale * sample + self._broadcast(r, sample) + 2.0 * self._broadcast(t, sample) + cond_term
-
-
-class _TaskAwareConditionEncoder(nn.Module):
-    output_dim = 4
-    condition_sequence_length = 3
-    tokens_per_step = 3
-    joint_output_dim = 4
-    camera_names = _CAMERA_NAMES
-    num_cameras = 3
-
-    def __init__(self, **kwargs):
-        super().__init__()
-        self.constructor_kwargs = dict(kwargs)
-        self.bias = nn.Parameter(torch.tensor(0.0))
-        self.calls = []
-
-    def forward(self, images, state, task=None):
-        self.calls.append({'task': task, 'state': state.detach().clone(), 'image_keys': tuple(images.keys())})
-        batch_size = state.shape[0]
-        state_last = state[:, -1, 0]
-        if isinstance(task, str):
-            task_lengths = torch.full((batch_size,), float(len(task)), dtype=state.dtype, device=state.device)
-        else:
-            task_lengths = torch.tensor([float(len(item)) for item in task], dtype=state.dtype, device=state.device)
-        image_marker = images['r_vis'][:, -1].mean(dim=(1, 2, 3))
-        token0 = torch.stack([state_last, task_lengths, image_marker, torch.ones_like(state_last)], dim=-1)
-        token1 = token0 + 1.0
-        token2 = token0 + 2.0
-        return torch.stack([token0, token1, token2], dim=1) + self.bias
-
-
-class _BF16TaskAwareConditionEncoder(_TaskAwareConditionEncoder):
-    def forward(self, images, state, task=None):
-        return super().forward(images, state, task=task).to(dtype=torch.bfloat16)
-
-
-class _StubIMFHead(nn.Module):
-    def __init__(self, input_dim, output_dim, horizon, n_obs_steps, cond_dim, **kwargs):
-        super().__init__()
-        self.constructor_kwargs = {
-            'input_dim': input_dim,
-            'output_dim': output_dim,
-            'horizon': horizon,
-            'n_obs_steps': n_obs_steps,
-            'cond_dim': cond_dim,
-            **kwargs,
-        }
-        self.proj = nn.Linear(input_dim, output_dim)
-        self.cond_obs_emb = nn.Linear(cond_dim, max(cond_dim, 1))
-
-    def forward(self, sample, r, t, cond=None):
-        return torch.zeros_like(sample)
-
-    def get_optim_groups(self, weight_decay):
-        return [
-            {'params': [self.proj.weight], 'weight_decay': weight_decay},
-            {'params': [self.proj.bias, self.cond_obs_emb.weight, self.cond_obs_emb.bias], 'weight_decay': 0.0},
-        ]
-
-
-@contextlib.contextmanager
-def _stub_optional_modules(include_head=False, include_condition_encoder=False):
-    previous = {}
-
-    def inject(name, module):
-        if name not in previous:
-            previous[name] = sys.modules.get(name, _MISSING)
-        sys.modules[name] = module
-
-    diffusers_module = types.ModuleType('diffusers')
-    schedulers_module = types.ModuleType('diffusers.schedulers')
-    ddpm_module = types.ModuleType('diffusers.schedulers.scheduling_ddpm')
-    ddim_module = types.ModuleType('diffusers.schedulers.scheduling_ddim')
-
-    class _FakeScheduler:
-        def __init__(self, num_train_timesteps=100, **kwargs):
-            self.config = types.SimpleNamespace(num_train_timesteps=num_train_timesteps)
-
-    ddpm_module.DDPMScheduler = _FakeScheduler
-    ddim_module.DDIMScheduler = _FakeScheduler
-    diffusers_module.DDPMScheduler = _FakeScheduler
-    diffusers_module.DDIMScheduler = _FakeScheduler
-    diffusers_module.schedulers = schedulers_module
-
-    try:
-        inject('diffusers', diffusers_module)
-        inject('diffusers.schedulers', schedulers_module)
-        inject('diffusers.schedulers.scheduling_ddpm', ddpm_module)
-        inject('diffusers.schedulers.scheduling_ddim', ddim_module)
-        if include_head:
-            import roboimi.vla.models.heads as heads_package
-            head_module = types.ModuleType('roboimi.vla.models.heads.imf_transformer1d')
-            head_module.IMFTransformer1D = _StubIMFHead
-            inject('roboimi.vla.models.heads.imf_transformer1d', head_module)
-            setattr(heads_package, 'imf_transformer1d', head_module)
-        if include_condition_encoder:
-            module = types.ModuleType('tests.fake_smolvla_condition_encoder')
-            module.TaskAwareConditionEncoder = _TaskAwareConditionEncoder
-            module.BF16TaskAwareConditionEncoder = _BF16TaskAwareConditionEncoder
-            inject('tests.fake_smolvla_condition_encoder', module)
-        yield
-    finally:
-        for name, old in reversed(list(previous.items())):
-            if old is _MISSING:
-                sys.modules.pop(name, None)
-            else:
-                sys.modules[name] = old
-
-
-def _compose_cfg(overrides=None):
-    if not OmegaConf.has_resolver('len'):
-        OmegaConf.register_new_resolver('len', lambda x: len(x))
-    GlobalHydra.instance().clear()
-    with initialize_config_dir(version_base=None, config_dir=_CONFIG_DIR):
-        return compose(config_name='config', overrides=list(overrides or []))
-
-
-class SmolVLAIMFAgentTest(unittest.TestCase):
-    def test_compute_loss_and_predict_action_pass_variable_task_to_condition_encoder(self):
-        from roboimi.vla.agent_smolvla_conditioned import SmolVLAIMFAttnResAgent
-
-        condition_encoder = _BF16TaskAwareConditionEncoder()
-        head = _RecordingHead()
-        agent = SmolVLAIMFAttnResAgent(
-            condition_encoder=condition_encoder,
-            action_encoder=nn.Identity(),
-            head=head,
-            action_dim=2,
-            obs_dim=1,
-            pred_horizon=3,
-            obs_horizon=2,
-            diffusion_steps=10,
-            inference_steps=1,
-            num_cams=3,
-            camera_names=_CAMERA_NAMES,
-            num_action_steps=2,
-            head_type='transformer',
-        )
-        images = {
-            'r_vis': torch.full((2, 2, 1, 2, 2), 1.0),
-            'top': torch.full((2, 2, 1, 2, 2), 2.0),
-            'front': torch.full((2, 2, 1, 2, 2), 3.0),
-        }
-        qpos = torch.tensor([[[1.0], [2.0]], [[3.0], [4.0]]])
-        actions = torch.zeros(2, 3, 2)
-        tasks = ['short', 'longer task']
-        loss = agent.compute_loss({'images': images, 'qpos': qpos, 'action': actions, 'task': tasks})
-        self.assertTrue(torch.isfinite(loss))
-        self.assertEqual(condition_encoder.calls[-1]['task'], tasks)
-        self.assertEqual(head.calls[-1]['cond'].shape, (2, 3, 4))
-        self.assertTrue(torch.allclose(head.calls[-1]['cond'][:, 0, 1], torch.tensor([5.0, 11.0])))
-
-        with mock.patch('roboimi.vla.agent_imf.torch.randn', return_value=torch.zeros(2, 3, 2)):
-            pred = agent.predict_action(images, qpos, task=tasks)
-        self.assertEqual(pred.shape, (2, 3, 2))
-        self.assertEqual(condition_encoder.calls[-1]['task'], tasks)
-
-    def test_condition_tokens_are_cast_to_action_head_dtype(self):
-        from roboimi.vla.agent_smolvla_conditioned import SmolVLAIMFAttnResAgent
-
-        condition_encoder = _TaskAwareConditionEncoder()
-        head = _RecordingHead()
-        agent = SmolVLAIMFAttnResAgent(
-            condition_encoder=condition_encoder,
-            action_encoder=nn.Identity(),
-            head=head,
-            action_dim=2,
-            obs_dim=1,
-            pred_horizon=3,
-            obs_horizon=2,
-            diffusion_steps=10,
-            inference_steps=1,
-            num_cams=3,
-            camera_names=_CAMERA_NAMES,
-            num_action_steps=2,
-            head_type='transformer',
-        )
-        images = {
-            'r_vis': torch.full((1, 2, 1, 2, 2), 1.0),
-            'top': torch.full((1, 2, 1, 2, 2), 2.0),
-            'front': torch.full((1, 2, 1, 2, 2), 3.0),
-        }
-        qpos = torch.tensor([[[1.0], [2.0]]], dtype=torch.float32)
-
-        cond = agent._build_cond(images, qpos, task=['pick'])
-
-        self.assertEqual(cond.dtype, head.scale.dtype)
-
-    def test_unknown_dataset_task_uses_configured_task_description(self):
-        from roboimi.vla.agent_smolvla_conditioned import SmolVLAIMFAttnResAgent
-
-        condition_encoder = _TaskAwareConditionEncoder()
-        head = _RecordingHead()
-        agent = SmolVLAIMFAttnResAgent(
-            condition_encoder=condition_encoder,
-            action_encoder=nn.Identity(),
-            head=head,
-            action_dim=2,
-            obs_dim=1,
-            pred_horizon=3,
-            obs_horizon=2,
-            diffusion_steps=10,
-            inference_steps=1,
-            num_cams=3,
-            camera_names=_CAMERA_NAMES,
-            num_action_steps=2,
-            head_type='transformer',
-            task_description='insert the peg into the socket',
-        )
-        images = {
-            'r_vis': torch.full((2, 2, 1, 2, 2), 1.0),
-            'top': torch.full((2, 2, 1, 2, 2), 2.0),
-            'front': torch.full((2, 2, 1, 2, 2), 3.0),
-        }
-        qpos = torch.tensor([[[1.0], [2.0]], [[3.0], [4.0]]], dtype=torch.float32)
-
-        agent._build_cond(images, qpos, task=['unknown', ''])
-
-        self.assertEqual(
-            condition_encoder.calls[-1]['task'],
-            ['insert the peg into the socket', 'insert the peg into the socket'],
-        )
-
-    def test_hydra_config_instantiates_smolvla_imf_attnres_with_condition_encoder_contract(self):
-        cfg = _compose_cfg(overrides=[
-            'agent=smolvla_imf_attnres',
-            'agent.condition_encoder._target_=tests.fake_smolvla_condition_encoder.TaskAwareConditionEncoder',
-            'agent.condition_dim=4',
-            'agent.condition_sequence_length=3',
-            'agent.head.n_layer=1',
-            'agent.head.n_emb=16',
-        ])
-        self.assertEqual(cfg.agent._target_, 'roboimi.vla.agent_smolvla_conditioned.SmolVLAIMFAttnResAgent')
-        self.assertEqual(cfg.agent.head.cond_dim, cfg.agent.condition_dim)
-        self.assertEqual(cfg.agent.head.n_obs_steps, cfg.agent.condition_sequence_length)
-
-        with _stub_optional_modules(include_head=True, include_condition_encoder=True):
-            agent = instantiate(cfg.agent)
-
-        self.assertEqual(agent.per_step_cond_dim, 4)
-        self.assertEqual(agent.condition_sequence_length, 3)
-        self.assertIsInstance(agent.noise_pred_net, _StubIMFHead)
-        self.assertEqual(agent.noise_pred_net.constructor_kwargs['cond_dim'], 4)
-        self.assertEqual(agent.noise_pred_net.constructor_kwargs['n_obs_steps'], 3)
-
-    def test_hydra_config_exposes_smolvla_pretrained_vlm_defaults(self):
-        cfg = _compose_cfg(overrides=[
-            'agent=smolvla_imf_attnres',
-        ])
-
-        self.assertEqual(cfg.agent.condition_encoder.model_name, 'HuggingFaceTB/SmolVLM2-500M-Video-Instruct')
-        self.assertTrue(cfg.agent.condition_encoder.load_vlm_weights)
-        self.assertEqual(cfg.agent.condition_encoder.num_vlm_layers, 16)
-        self.assertTrue(cfg.agent.condition_encoder.freeze_vlm)
-        self.assertTrue(cfg.agent.condition_encoder.freeze_vision_encoder)
-        self.assertTrue(cfg.agent.condition_encoder.run_text_model)
-        self.assertEqual(cfg.agent.condition_encoder.max_state_dim, 32)
-        self.assertIsNone(cfg.agent.condition_encoder.dataset_image_resize_shape)
-        self.assertIsNone(cfg.agent.condition_encoder.eval_image_resize_shape)
-        self.assertEqual(cfg.agent.condition_dim, 960)
-        self.assertEqual(cfg.agent.condition_sequence_length, 241)
-        self.assertEqual(cfg.agent.head.cond_dim, 960)
-        self.assertEqual(cfg.agent.head.n_obs_steps, 241)
-
-
-if __name__ == '__main__':
-    unittest.main()

tests/test_smolvla_native_agent.py

@@ -1,445 +0,0 @@
-import contextlib
-import sys
-import types
-import unittest
-from pathlib import Path
-
-import torch
-from hydra import compose, initialize_config_dir
-from hydra.core.global_hydra import GlobalHydra
-from omegaconf import OmegaConf
-from torch import nn
-
-
-_REPO_ROOT = Path(__file__).resolve().parents[1]
-_CONFIG_DIR = str((_REPO_ROOT / 'roboimi/vla/conf').resolve())
-_CAMERA_NAMES = ('r_vis', 'top', 'front')
-_MISSING = object()
-
-
-def _stats():
-    return {
-        'qpos_min': [0.0, 10.0],
-        'qpos_max': [10.0, 20.0],
-        'action_min': [-10.0, 10.0],
-        'action_max': [10.0, 30.0],
-    }
-
-
-class FakeTokenizer:
-    def __init__(self):
-        self.calls = []
-
-    def __call__(self, texts, **kwargs):
-        self.calls.append({'texts': list(texts), 'kwargs': dict(kwargs)})
-        batch = len(texts)
-        if kwargs.get('padding') == 'max_length' and kwargs.get('max_length') is not None:
-            max_len = int(kwargs['max_length'])
-        else:
-            max_len = max(len(text) for text in texts) if texts else 0
-        input_ids = torch.zeros(batch, max_len, dtype=torch.long)
-        attention_mask = torch.zeros(batch, max_len, dtype=torch.long)
-        for i, text in enumerate(texts):
-            encoded = [ord(ch) % 127 for ch in text]
-            if kwargs.get('truncation') and len(encoded) > max_len:
-                encoded = encoded[:max_len]
-            if encoded:
-                input_ids[i, : len(encoded)] = torch.tensor(encoded, dtype=torch.long)
-                attention_mask[i, : len(encoded)] = 1
-        return {'input_ids': input_ids, 'attention_mask': attention_mask}
-
-
-class FakeNativeModel(nn.Module):
-    def __init__(self, action_dim=2, chunk_size=3):
-        super().__init__()
-        self.anchor = nn.Parameter(torch.tensor(0.0))
-        self.action_dim = action_dim
-        self.chunk_size = chunk_size
-        self.forward_calls = []
-        self.sample_calls = []
-        self.sample_return = torch.tensor([[[-1.0, 1.0], [0.0, 0.0], [1.0, -1.0]]])
-
-    def forward(self, **kwargs):
-        self.forward_calls.append({k: _clone(v) for k, v in kwargs.items()})
-        actions = kwargs['actions']
-        loss = (actions**2).sum(dim=-1)
-        if 'action_is_pad' in kwargs and kwargs['action_is_pad'] is not None:
-            mask = (~kwargs['action_is_pad']).to(loss.dtype)
-            loss = (loss * mask).sum() / mask.sum().clamp_min(1.0)
-        else:
-            loss = loss.mean()
-        return {'loss': loss + self.anchor}
-
-    def sample_actions(self, **kwargs):
-        self.sample_calls.append({k: _clone(v) for k, v in kwargs.items()})
-        return self.sample_return.to(device=kwargs['state'].device, dtype=kwargs['state'].dtype)
-
-
-def _clone(value):
-    if torch.is_tensor(value):
-        return value.detach().clone()
-    if isinstance(value, dict):
-        return {k: _clone(v) for k, v in value.items()}
-    if isinstance(value, list):
-        return list(value)
-    return value
-
-
-def _make_agent(**kwargs):
-    from roboimi.vla.agent_smolvla_native import SmolVLANativeAgent
-
-    params = dict(
-        model=FakeNativeModel(),
-        tokenizer=FakeTokenizer(),
-        action_dim=2,
-        obs_dim=2,
-        chunk_size=3,
-        n_action_steps=2,
-        obs_horizon=2,
-        camera_names=_CAMERA_NAMES,
-        num_cams=3,
-        dataset_stats=_stats(),
-        normalization_type='min_max',
-        task_description='fallback task',
-        model_config={'resize_imgs_with_padding': None},
-    )
-    params.update(kwargs)
-    return SmolVLANativeAgent(**params)
-
-
-def _batch(task=None):
-    images = {
-        'front': torch.full((2, 2, 1, 2, 2), 3.0),
-        'r_vis': torch.full((2, 2, 1, 2, 2), 1.0),
-        'top': torch.full((2, 2, 1, 2, 2), 2.0),
-    }
-    batch = {
-        'images': images,
-        'qpos': torch.tensor([[[0.0, 10.0], [10.0, 20.0]], [[5.0, 15.0], [10.0, 10.0]]]),
-        'action': torch.tensor([[[-10.0, 10.0], [0.0, 20.0], [10.0, 30.0]], [[0.0, 20.0], [10.0, 10.0], [-10.0, 30.0]]]),
-        'action_is_pad': torch.tensor([[False, False, True], [False, True, True]]),
-    }
-    if task is not None:
-        batch['task'] = task
-    return batch
-
-
-@contextlib.contextmanager
-def _stub_native_modules():
-    old_pkg = sys.modules.get('roboimi.vla.models.smolvla', _MISSING)
-    old_conf = sys.modules.get('roboimi.vla.models.smolvla.configuration', _MISSING)
-    try:
-        pkg = types.ModuleType('roboimi.vla.models.smolvla')
-        conf = types.ModuleType('roboimi.vla.models.smolvla.configuration')
-
-        class NativeSmolVLAConfig:
-            def __init__(self, **kwargs):
-                self.kwargs = kwargs
-
-        conf.NativeSmolVLAConfig = NativeSmolVLAConfig
-        sys.modules['roboimi.vla.models.smolvla'] = pkg
-        sys.modules['roboimi.vla.models.smolvla.configuration'] = conf
-        yield
-    finally:
-        for name, old in [
-            ('roboimi.vla.models.smolvla.configuration', old_conf),
-            ('roboimi.vla.models.smolvla', old_pkg),
-        ]:
-            if old is _MISSING:
-                sys.modules.pop(name, None)
-            else:
-                sys.modules[name] = old
-
-
-def _compose_cfg(overrides=None):
-    if not OmegaConf.has_resolver('len'):
-        OmegaConf.register_new_resolver('len', lambda x: len(x))
-    GlobalHydra.instance().clear()
-    with initialize_config_dir(version_base=None, config_dir=_CONFIG_DIR):
-        return compose(config_name='config', overrides=list(overrides or []))
-
-
-
-
-class StrictSignatureNativeModel(nn.Module):
-    def __init__(self, action_dim=2, max_action_dim=4, chunk_size=3):
-        super().__init__()
-        self.anchor = nn.Parameter(torch.tensor(0.0))
-        self.action_dim = action_dim
-        self.max_action_dim = max_action_dim
-        self.chunk_size = chunk_size
-        self.forward_calls = []
-        self.sample_calls = []
-
-    def forward(self, images, img_masks, lang_tokens, lang_masks, state, actions):
-        self.forward_calls.append({
-            'images': _clone(images),
-            'img_masks': _clone(img_masks),
-            'lang_tokens': _clone(lang_tokens),
-            'lang_masks': _clone(lang_masks),
-            'state': _clone(state),
-            'actions': _clone(actions),
-        })
-        # Per-element losses in the native core's max_action_dim space.
-        return actions.pow(2) + self.anchor
-
-    def sample_actions(self, images, img_masks, lang_tokens, lang_masks, state):
-        self.sample_calls.append({
-            'images': _clone(images),
-            'img_masks': _clone(img_masks),
-            'lang_tokens': _clone(lang_tokens),
-            'lang_masks': _clone(lang_masks),
-            'state': _clone(state),
-        })
-        batch_size = state.shape[0]
-        out = torch.zeros(batch_size, self.chunk_size, self.max_action_dim, device=state.device, dtype=state.dtype)
-        out[..., : self.action_dim] = torch.tensor(
-            [[[-1.0, 1.0], [0.0, 0.0], [1.0, -1.0]]],
-            device=state.device,
-            dtype=state.dtype,
-        ).expand(batch_size, -1, -1)
-        out[..., self.action_dim :] = 123.0
-        return out
-
-
-class SmolVLANativeAgentTest(unittest.TestCase):
-    def test_compute_loss_orders_normalizes_tokenizes_and_passes_action_pad(self):
-        agent = _make_agent()
-        batch = _batch(task=['pick', 'place'])
-
-        loss = agent.compute_loss(batch)
-
-        self.assertTrue(torch.isfinite(loss))
-        call = agent.model.forward_calls[-1]
-        self.assertIsInstance(call['images'], list)
-        self.assertEqual(len(call['images']), len(_CAMERA_NAMES))
-        self.assertTrue(torch.allclose(call['images'][0], torch.full((2, 1, 2, 2), 1.0)))
-        self.assertTrue(torch.allclose(call['state'], torch.tensor([[1.0, 1.0], [1.0, -1.0]])))
-        self.assertTrue(torch.allclose(call['actions'][0], torch.tensor([[-1.0, -1.0], [0.0, 0.0], [1.0, 1.0]])))
-        self.assertNotIn('action_is_pad', call)
-        self.assertEqual(agent.tokenizer.calls[-1]['texts'], ['pick\n', 'place\n'])
-        self.assertIn('lang_tokens', call)
-        self.assertIn('lang_masks', call)
-        self.assertEqual(call['lang_tokens'].dtype, torch.long)
-        self.assertEqual(call['lang_masks'].dtype, torch.bool)
-
-    def test_task_fallback_for_missing_empty_and_unknown_but_valid_task_wins(self):
-        agent = _make_agent(task_description='default instruction')
-
-        agent.compute_loss(_batch())
-        self.assertEqual(agent.tokenizer.calls[-1]['texts'], ['default instruction\n', 'default instruction\n'])
-
-        agent.compute_loss(_batch(task=[None, '']))
-        self.assertEqual(agent.tokenizer.calls[-1]['texts'], ['default instruction\n', 'default instruction\n'])
-
-        agent.compute_loss(_batch(task=['unknown', 'valid task']))
-        self.assertEqual(agent.tokenizer.calls[-1]['texts'], ['default instruction\n', 'valid task\n'])
-
-    def test_missing_camera_raises_value_error(self):
-        agent = _make_agent()
-        batch = _batch(task=['pick', 'place'])
-        del batch['images']['top']
-
-        with self.assertRaisesRegex(ValueError, 'missing.*top'):
-            agent.compute_loss(batch)
-
-    def test_predict_action_chunk_samples_and_denormalizes_actions(self):
-        agent = _make_agent()
-        batch = _batch(task=['pick', 'place'])
-        batch.pop('action')
-        batch.pop('action_is_pad')
-        agent.model.sample_return = torch.tensor([[[-1.0, 1.0], [0.0, 0.0], [1.0, -1.0]], [[0.5, -0.5], [-0.5, 0.5], [0.0, 1.0]]])
-
-        actions = agent.predict_action_chunk(batch)
-
-        self.assertTrue(torch.allclose(actions[0], torch.tensor([[-10.0, 30.0], [0.0, 20.0], [10.0, 10.0]])))
-        self.assertEqual(actions.shape, (2, 3, 2))
-        self.assertEqual(len(agent.model.sample_calls), 1)
-        self.assertTrue(torch.allclose(agent.model.sample_calls[-1]['state'][0], torch.tensor([1.0, 1.0])))
-
-    def test_select_action_queues_actions_and_defaults_chunk_start_zero(self):
-        agent = _make_agent(n_action_steps=2, action_chunk_start=0)
-        agent.model.sample_return = torch.tensor([[[-1.0, -1.0], [0.0, 0.0], [1.0, 1.0]]])
-        obs = {
-            'images': {name: torch.full((1, 2, 2), float(i + 1)) for i, name in enumerate(_CAMERA_NAMES)},
-            'qpos': torch.tensor([0.0, 10.0]),
-            'task': 'rollout task',
-        }
-
-        action0 = agent.select_action(obs)
-        action1 = agent.select_action(obs)
-
-        self.assertEqual(len(agent.model.sample_calls), 1)
-        self.assertTrue(torch.allclose(action0, torch.tensor([-10.0, 10.0])))
-        self.assertTrue(torch.allclose(action1, torch.tensor([0.0, 20.0])))
-        self.assertEqual(agent.action_chunk_start, 0)
-
-    def test_get_normalization_stats_returns_stats(self):
-        stats = _make_agent().get_normalization_stats()
-        self.assertEqual(stats['normalization_type'], 'min_max')
-        self.assertEqual(stats['qpos_min'], [0.0, 10.0])
-
-
-    def test_agent_adapts_roboimi_batch_to_native_vla_signature_and_reduces_loss(self):
-        from roboimi.vla.agent_smolvla_native import SmolVLANativeAgent
-
-        model = StrictSignatureNativeModel(action_dim=2, max_action_dim=4, chunk_size=3)
-        agent = SmolVLANativeAgent(
-            model=model,
-            tokenizer=FakeTokenizer(),
-            action_dim=2,
-            obs_dim=2,
-            chunk_size=3,
-            n_action_steps=2,
-            obs_horizon=2,
-            camera_names=_CAMERA_NAMES,
-            num_cams=3,
-            dataset_stats=_stats(),
-            normalization_type='min_max',
-            model_config={'max_state_dim': 4, 'max_action_dim': 4, 'resize_imgs_with_padding': None},
-        )
-
-        batch = _batch(task=['pick', 'place'])
-        batch['images'] = {
-            'front': torch.full((2, 2, 1, 2, 2), 0.6),
-            'r_vis': torch.full((2, 2, 1, 2, 2), 0.2),
-            'top': torch.full((2, 2, 1, 2, 2), 0.4),
-        }
-
-        loss = agent.compute_loss(batch)
-
-        self.assertEqual(loss.ndim, 0)
-        self.assertTrue(torch.isfinite(loss))
-        call = model.forward_calls[-1]
-        self.assertIsInstance(call['images'], list)
-        self.assertEqual(len(call['images']), len(_CAMERA_NAMES))
-        self.assertEqual(tuple(call['images'][0].shape), (2, 1, 2, 2))
-        self.assertTrue(torch.allclose(call['images'][0], torch.full((2, 1, 2, 2), -0.6)))
-        self.assertTrue(torch.allclose(call['images'][1], torch.full((2, 1, 2, 2), -0.2)))
-        self.assertTrue(torch.allclose(call['images'][2], torch.full((2, 1, 2, 2), 0.2)))
-        self.assertEqual(len(call['img_masks']), len(_CAMERA_NAMES))
-        self.assertTrue(torch.equal(call['img_masks'][0], torch.ones(2, dtype=torch.bool)))
-        self.assertTrue(torch.equal(call['lang_tokens'], model.forward_calls[-1]['lang_tokens']))
-        self.assertEqual(call['lang_tokens'].dtype, torch.long)
-        self.assertEqual(call['lang_masks'].dtype, torch.bool)
-        self.assertEqual(tuple(call['state'].shape), (2, 4))
-        self.assertTrue(torch.allclose(call['state'][0], torch.tensor([1.0, 1.0, 0.0, 0.0])))
-        self.assertEqual(tuple(call['actions'].shape), (2, 3, 4))
-        self.assertTrue(torch.allclose(call['actions'][0, :, :2], torch.tensor([[-1.0, -1.0], [0.0, 0.0], [1.0, 1.0]])))
-        self.assertTrue(torch.allclose(call['actions'][..., 2:], torch.zeros(2, 3, 2)))
-        # Valid entries: sample0 steps 0,1 and sample1 step 0, only first action_dim losses are reduced.
-        expected = (2.0 + 0.0 + 0.0) / (3 * 2)
-        self.assertTrue(torch.allclose(loss, torch.tensor(expected)))
-
-    def test_predict_action_chunk_accepts_native_max_action_dim_and_crops_before_denorm(self):
-        from roboimi.vla.agent_smolvla_native import SmolVLANativeAgent
-
-        model = StrictSignatureNativeModel(action_dim=2, max_action_dim=4, chunk_size=3)
-        agent = SmolVLANativeAgent(
-            model=model,
-            tokenizer=FakeTokenizer(),
-            action_dim=2,
-            obs_dim=2,
-            chunk_size=3,
-            n_action_steps=2,
-            obs_horizon=2,
-            camera_names=_CAMERA_NAMES,
-            num_cams=3,
-            dataset_stats=_stats(),
-            normalization_type='min_max',
-            model_config={'max_state_dim': 4, 'max_action_dim': 4, 'resize_imgs_with_padding': None},
-        )
-        batch = _batch(task=['pick', 'place'])
-        batch.pop('action')
-        batch.pop('action_is_pad')
-
-        actions = agent.predict_action_chunk(batch)
-
-        self.assertEqual(actions.shape, (2, 3, 2))
-        self.assertTrue(torch.allclose(actions[0], torch.tensor([[-10.0, 30.0], [0.0, 20.0], [10.0, 10.0]])))
-        call = model.sample_calls[-1]
-        self.assertEqual(tuple(call['state'].shape), (2, 4))
-        self.assertEqual(len(call['images']), len(_CAMERA_NAMES))
-
-    def test_build_model_filters_and_maps_non_core_config_fields(self):
-        from roboimi.vla.agent_smolvla_native import SmolVLANativeAgent
-        from roboimi.vla.models.smolvla.configuration import NativeSmolVLAConfig
-
-        captured = {}
-
-        class BuildOnlyAgent(SmolVLANativeAgent):
-            def _build_tokenizer(self, tokenizer_name):
-                return FakeTokenizer()
-
-            def _build_model(self):
-                cfg_kwargs = self._native_config_kwargs()
-                captured.update(cfg_kwargs)
-                return FakeNativeModel(action_dim=2, chunk_size=3)
-
-        agent = BuildOnlyAgent(
-            action_dim=2,
-            obs_dim=2,
-            chunk_size=3,
-            n_action_steps=2,
-            obs_horizon=2,
-            camera_names=_CAMERA_NAMES,
-            num_cams=3,
-            dataset_stats=_stats(),
-            normalization_type='min_max',
-            model_config={
-                'state_dim': 2,
-                'action_dim': 2,
-                'max_state_dim': 4,
-                'max_action_dim': 4,
-                'tokenizer_name': 'dummy-tokenizer',
-                'freeze_vlm': True,
-                'image_resize_shape': [512, 320],
-                'num_cameras': 3,
-                'load_vlm_weights': False,
-            },
-        )
-        self.assertIsNotNone(agent.model)
-        config = NativeSmolVLAConfig(**captured)
-        self.assertEqual(config.resize_imgs_with_padding, (512, 320))
-        self.assertEqual(config.max_state_dim, 4)
-        self.assertEqual(config.max_action_dim, 4)
-        self.assertNotIn('state_dim', captured)
-        self.assertNotIn('action_dim', captured)
-        self.assertNotIn('tokenizer_name', captured)
-        self.assertNotIn('freeze_vlm', captured)
-        self.assertNotIn('image_resize_shape', captured)
-        self.assertNotIn('num_cameras', captured)
-
-    def test_hydra_config_target_and_key_fields(self):
-        with _stub_native_modules():
-            cfg = _compose_cfg(overrides=['agent=smolvla_native'])
-        self.assertEqual(cfg.agent._target_, 'roboimi.vla.agent_smolvla_native.SmolVLANativeAgent')
-        self.assertEqual(cfg.agent.action_dim, 16)
-        self.assertEqual(cfg.agent.obs_dim, 16)
-        self.assertEqual(cfg.agent.normalization_type, 'gaussian')
-        self.assertEqual(list(cfg.agent.camera_names), list(cfg.data.camera_names))
-        self.assertEqual(cfg.agent.num_cams, len(cfg.data.camera_names))
-        self.assertEqual(cfg.agent.chunk_size, 32)
-        self.assertEqual(cfg.agent.n_action_steps, 16)
-        self.assertEqual(cfg.agent.model_config.max_state_dim, 32)
-        self.assertEqual(cfg.agent.model_config.max_action_dim, 32)
-        self.assertEqual(cfg.agent.model_config.chunk_size, 32)
-        self.assertEqual(cfg.agent.model_config.n_action_steps, 16)
-        self.assertIsNone(cfg.agent.dataset_image_resize_shape)
-        self.assertIsNone(cfg.agent.eval_image_resize_shape)
-        self.assertEqual(list(cfg.agent.model_config.resize_imgs_with_padding), [512, 512])
-        self.assertNotIn('state_dim', cfg.agent.model_config)
-        self.assertNotIn('action_dim', cfg.agent.model_config)
-        self.assertNotIn('tokenizer_name', cfg.agent.model_config)
-        self.assertEqual(cfg.agent.model_config.vlm_model_name, 'HuggingFaceTB/SmolVLM2-500M-Video-Instruct')
-
-    def test_tokenize_tasks_preserves_existing_single_newline(self):
-        agent = _make_agent(model_config={'resize_imgs_with_padding': None, 'pad_language_to': 'max_length', 'tokenizer_max_length': 12})
-
-        agent._tokenize_tasks(['already newline\n', 'needs newline'], batch_size=2, device=torch.device('cpu'))
-
-        self.assertEqual(agent.tokenizer.calls[-1]['texts'], ['already newline\n', 'needs newline\n'])
-
-
-if __name__ == '__main__':
-    unittest.main()

tests/test_smolvla_native_modeling.py

@@ -1,121 +0,0 @@
-import unittest
-from types import SimpleNamespace
-
-import torch
-from torch import nn
-
-from roboimi.vla.models.smolvla import NativeSmolVLAConfig
-from roboimi.vla.models.smolvla.modeling import (
-    VLAFlowMatching,
-    make_att_2d_masks,
-    pad_vector,
-    resize_with_pad,
-)
-
-
-class FakeTokenizer:
-    fake_image_token_id = 32000
-    global_image_token_id = 32001
-
-
-class FakeProcessor:
-    tokenizer = FakeTokenizer()
-
-
-class FakeVLMWithExpert(nn.Module):
-    def __init__(self, vlm_hidden_size=8, expert_hidden_size=6, image_tokens=3, vocab_size=64):
-        super().__init__()
-        self.config = SimpleNamespace(text_config=SimpleNamespace(hidden_size=vlm_hidden_size))
-        self.expert_hidden_size = expert_hidden_size
-        self.image_tokens = image_tokens
-        self.processor = FakeProcessor()
-        self.vlm = SimpleNamespace(device=torch.device("cpu"))
-        self.image_proj = nn.Linear(3, vlm_hidden_size)
-        self.token_emb = nn.Embedding(vocab_size, vlm_hidden_size)
-        self.suffix_proj = nn.Linear(expert_hidden_size, expert_hidden_size)
-
-    def embed_image(self, image):
-        # Deterministic lightweight image embedding: pool pixels, then repeat.
-        pooled = image.mean(dim=(-1, -2)).to(dtype=torch.float32)
-        return self.image_proj(pooled).unsqueeze(1).expand(-1, self.image_tokens, -1)
-
-    def embed_language_tokens(self, tokens):
-        return self.token_emb(tokens)
-
-    def forward(self, attention_mask, position_ids, past_key_values, inputs_embeds, use_cache, fill_kv_cache):
-        prefix_embs, suffix_embs = inputs_embeds
-        prefix_out = prefix_embs if prefix_embs is not None else None
-        suffix_out = self.suffix_proj(suffix_embs) if suffix_embs is not None else None
-        cache = ("fake-cache",) if fill_kv_cache else past_key_values
-        return (prefix_out, suffix_out), cache
-
-
-class NativeSmolVLAModelingTest(unittest.TestCase):
-    def test_config_rejects_action_steps_greater_than_chunk_size(self):
-        with self.assertRaisesRegex(ValueError, "n_action_steps"):
-            NativeSmolVLAConfig(chunk_size=2, n_action_steps=3)
-
-    def test_pad_vector_pads_returns_original_for_equal_and_rejects_truncation(self):
-        vector = torch.tensor([[1.0, 2.0, 3.0]])
-        padded = pad_vector(vector, 5)
-        self.assertEqual(tuple(padded.shape), (1, 5))
-        torch.testing.assert_close(padded, torch.tensor([[1.0, 2.0, 3.0, 0.0, 0.0]]))
-
-        same = pad_vector(vector, 3)
-        self.assertIs(same, vector)
-
-        with self.assertRaisesRegex(ValueError, "target dimension"):
-            pad_vector(vector, 2)
-
-    def test_resize_with_pad_returns_requested_spatial_size(self):
-        img = torch.arange(2 * 3 * 4 * 8, dtype=torch.float32).reshape(2, 3, 4, 8)
-        resized = resize_with_pad(img, width=10, height=10, pad_value=-1)
-        self.assertEqual(tuple(resized.shape), (2, 3, 10, 10))
-
-    def test_make_att_2d_masks_implements_prefix_lm_semantics(self):
-        pad_masks = torch.tensor([[True, True, True, True, False]])
-        # First two tokens are bidirectional prefix, later valid tokens are causal.
-        att_masks = torch.tensor([[False, False, True, True, True]])
-        mask = make_att_2d_masks(pad_masks, att_masks)
-        expected = torch.tensor(
-            [[
-                [True, True, False, False, False],
-                [True, True, False, False, False],
-                [True, True, True, False, False],
-                [True, True, True, True, False],
-                [False, False, False, False, False],
-            ]]
-        )
-        torch.testing.assert_close(mask, expected)
-
-    def test_vla_flow_matching_forward_and_sample_actions_with_fake_vlm(self):
-        torch.manual_seed(0)
-        config = NativeSmolVLAConfig(
-            chunk_size=4,
-            n_action_steps=4,
-            max_state_dim=5,
-            max_action_dim=3,
-            num_steps=2,
-            prefix_length=8,
-            add_image_special_tokens=False,
-        )
-        fake_vlm = FakeVLMWithExpert(vlm_hidden_size=8, expert_hidden_size=6)
-        model = VLAFlowMatching(config, vlm_with_expert=fake_vlm)
-
-        bsize = 2
-        images = [torch.randn(bsize, 3, 6, 6)]
-        img_masks = [torch.tensor([True, False])]
-        lang_tokens = torch.tensor([[1, 2, 3], [4, 5, 0]])
-        lang_masks = torch.tensor([[True, True, True], [True, True, False]])
-        state = torch.randn(bsize, config.max_state_dim)
-        actions = torch.randn(bsize, config.chunk_size, config.max_action_dim)
-
-        losses = model(images, img_masks, lang_tokens, lang_masks, state, actions)
-        self.assertEqual(tuple(losses.shape), (bsize, config.chunk_size, config.max_action_dim))
-
-        sampled = model.sample_actions(images, img_masks, lang_tokens, lang_masks, state)
-        self.assertEqual(tuple(sampled.shape), (bsize, config.chunk_size, config.max_action_dim))
-
-
-if __name__ == "__main__":
-    unittest.main()

tests/test_smolvla_prefix_encoder.py

@@ -1,328 +0,0 @@
-import types
-import unittest
-from unittest import mock
-
-import torch
-from torch import nn
-
-
-class _FakeVisionOutput:
-    def __init__(self, last_hidden_state):
-        self.last_hidden_state = last_hidden_state
-
-
-class _FakeVisionModel(nn.Module):
-    def __init__(self, hidden_size=4):
-        super().__init__()
-        self.dtype = torch.float32
-        self.scale = nn.Parameter(torch.tensor(1.0))
-        self.calls = []
-        self.hidden_size = hidden_size
-
-    def forward(self, pixel_values=None, patch_attention_mask=None):
-        self.calls.append({
-            'pixel_values': pixel_values.detach().clone(),
-            'patch_attention_mask': patch_attention_mask,
-        })
-        pooled = pixel_values.mean(dim=(2, 3)) * self.scale
-        tokens = torch.stack([pooled, pooled + 1.0], dim=1)
-        return _FakeVisionOutput(tokens)
-
-
-class _FakeConnector(nn.Module):
-    def __init__(self, in_dim=3, out_dim=4):
-        super().__init__()
-        self.proj = nn.Linear(in_dim, out_dim, bias=False)
-        with torch.no_grad():
-            self.proj.weight.copy_(
-                torch.tensor(
-                    [
-                        [1.0, 0.0, 0.0],
-                        [0.0, 1.0, 0.0],
-                        [0.0, 0.0, 1.0],
-                        [1.0, 1.0, 1.0],
-                    ]
-                )
-            )
-
-    def forward(self, x):
-        return self.proj(x)
-
-
-class _FakeTextModel(nn.Module):
-    def __init__(self, vocab_size=32, hidden_size=4, num_layers=6):
-        super().__init__()
-        self.embed = nn.Embedding(vocab_size, hidden_size)
-        self.layers = nn.ModuleList([nn.Linear(hidden_size, hidden_size) for _ in range(num_layers)])
-        self.norm = nn.Identity()
-        self.forward_calls = []
-        with torch.no_grad():
-            for idx in range(vocab_size):
-                self.embed.weight[idx].fill_(float(idx))
-
-    def get_input_embeddings(self):
-        return self.embed
-
-    def forward(
-        self,
-        input_ids=None,
-        attention_mask=None,
-        position_ids=None,
-        past_key_values=None,
-        inputs_embeds=None,
-        use_cache=None,
-        return_dict=True,
-        cache_position=None,
-        **kwargs,
-    ):
-        del input_ids, past_key_values, use_cache, cache_position, kwargs
-        self.forward_calls.append({
-            'attention_mask': None if attention_mask is None else attention_mask.detach().clone(),
-            'position_ids': None if position_ids is None else position_ids.detach().clone(),
-            'inputs_embeds': inputs_embeds.detach().clone(),
-        })
-        hidden = inputs_embeds
-        for layer in self.layers:
-            hidden = layer(hidden)
-        hidden = self.norm(hidden)
-        if return_dict:
-            return types.SimpleNamespace(last_hidden_state=hidden)
-        return (hidden,)
-
-
-class _FakeVLM(nn.Module):
-    def __init__(self):
-        super().__init__()
-        text_config = types.SimpleNamespace(hidden_size=4, head_dim=2, num_attention_heads=2, num_key_value_heads=1)
-        self.config = types.SimpleNamespace(text_config=text_config)
-        self.model = types.SimpleNamespace(
-            vision_model=_FakeVisionModel(hidden_size=4),
-            connector=_FakeConnector(in_dim=3, out_dim=4),
-            text_model=_FakeTextModel(hidden_size=4, num_layers=6),
-        )
-
-
-class _FakeTokenizer:
-    fake_image_token_id = 29
-    global_image_token_id = 30
-
-    def __init__(self):
-        self.padding_side = 'left'
-        self.calls = []
-
-    def __call__(self, text, *, padding, max_length, return_tensors, truncation):
-        self.calls.append({
-            'text': list(text),
-            'padding': padding,
-            'max_length': max_length,
-            'return_tensors': return_tensors,
-            'truncation': truncation,
-            'padding_side': self.padding_side,
-        })
-        batch = len(text)
-        ids = torch.zeros(batch, max_length, dtype=torch.long)
-        mask = torch.zeros(batch, max_length, dtype=torch.bool)
-        for row, item in enumerate(text):
-            del item
-            ids[row, :3] = torch.tensor([1, 2, 3])
-            mask[row, :3] = True
-        return {'input_ids': ids, 'attention_mask': mask}
-
-
-class SmolVLAPrefixEncoderTest(unittest.TestCase):
-    def test_loads_pretrained_vlm_with_local_files_only_crops_layers_and_freezes_vlm(self):
-        from roboimi.vla.models.backbones.smolvla_prefix_encoder import SmolVLAPrefixEncoder
-
-        fake_vlm = _FakeVLM()
-        fake_tokenizer = _FakeTokenizer()
-        with mock.patch(
-            'roboimi.vla.models.backbones.smolvla_prefix_encoder.AutoModelForImageTextToText.from_pretrained',
-            return_value=fake_vlm,
-        ) as model_loader, mock.patch(
-            'roboimi.vla.models.backbones.smolvla_prefix_encoder.AutoTokenizer.from_pretrained',
-            return_value=fake_tokenizer,
-        ) as tokenizer_loader:
-            encoder = SmolVLAPrefixEncoder(
-                model_name='HuggingFaceTB/SmolVLM2-500M-Video-Instruct',
-                local_files_only=True,
-                num_vlm_layers=2,
-                freeze_vlm=True,
-                max_state_dim=5,
-                tokenizer_max_length=4,
-                camera_names=('r_vis', 'top'),
-                resize_imgs_with_padding=None,
-            )
-
-        model_loader.assert_called_once()
-        self.assertEqual(model_loader.call_args.kwargs['local_files_only'], True)
-        self.assertEqual(model_loader.call_args.kwargs['torch_dtype'], 'bfloat16')
-        tokenizer_loader.assert_called_once_with(
-            'HuggingFaceTB/SmolVLM2-500M-Video-Instruct',
-            local_files_only=True,
-        )
-        self.assertEqual(len(encoder.vlm.model.text_model.layers), 2)
-        self.assertTrue(all(not p.requires_grad for p in encoder.vlm.parameters()))
-        self.assertFalse(encoder.vlm.training)
-        encoder.train()
-        self.assertFalse(encoder.vlm.training)
-
-    def test_accepts_train_eval_resize_compatibility_fields(self):
-        from roboimi.vla.models.backbones.smolvla_prefix_encoder import SmolVLAPrefixEncoder
-
-        encoder = SmolVLAPrefixEncoder(
-            vlm=_FakeVLM(),
-            tokenizer=_FakeTokenizer(),
-            num_vlm_layers=3,
-            camera_names=('r_vis', 'top'),
-            resize_imgs_with_padding=None,
-            dataset_image_resize_shape=None,
-            eval_image_resize_shape=(640, 480),
-        )
-
-        self.assertIsNone(encoder.dataset_image_resize_shape)
-        self.assertEqual(encoder.eval_image_resize_shape, (640, 480))
-
-    def test_embed_prefix_uses_variable_tasks_camera_order_last_state_and_smolvla_masks(self):
-        from roboimi.vla.models.backbones.smolvla_prefix_encoder import SmolVLAPrefixEncoder
-
-        fake_vlm = _FakeVLM()
-        fake_tokenizer = _FakeTokenizer()
-        encoder = SmolVLAPrefixEncoder(
-            vlm=fake_vlm,
-            tokenizer=fake_tokenizer,
-            num_vlm_layers=3,
-            freeze_vlm=True,
-            train_state_proj=True,
-            max_state_dim=5,
-            tokenizer_max_length=4,
-            camera_names=('r_vis', 'top'),
-            resize_imgs_with_padding=None,
-        )
-        with torch.no_grad():
-            encoder.state_proj.weight.zero_()
-            encoder.state_proj.bias.zero_()
-            encoder.state_proj.weight[:, :4] = torch.eye(4)
-
-        images = {
-            'top': torch.full((2, 2, 3, 2, 2), 0.75),
-            'r_vis': torch.full((2, 2, 3, 2, 2), 0.25),
-        }
-        state = torch.tensor(
-            [
-                [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]],
-                [[7.0, 8.0, 9.0], [10.0, 11.0, 12.0]],
-            ]
-        )
-        tasks = ['pick red cube', 'open drawer']
-        out = encoder.embed_prefix(images=images, state=state, task=tasks)
-
-        # 2 cameras * 2 image tokens + 4 language tokens + 1 state token
-        self.assertEqual(out.shape, (2, 9, 4))
-        self.assertEqual(encoder.output_dim, 4)
-        self.assertEqual(encoder.tokens_per_step, 9)
-        self.assertEqual(encoder.condition_sequence_length, 9)
-
-        self.assertEqual(fake_tokenizer.calls[-1]['text'], ['pick red cube\n', 'open drawer\n'])
-        self.assertEqual(fake_tokenizer.calls[-1]['padding'], 'max_length')
-        self.assertEqual(fake_tokenizer.calls[-1]['max_length'], 4)
-        self.assertEqual(fake_tokenizer.calls[-1]['padding_side'], 'right')
-
-        # Camera order is r_vis then top, and pixels are mapped [0,1] -> [-1,1].
-        first_camera_pixels = fake_vlm.model.vision_model.calls[0]['pixel_values']
-        second_camera_pixels = fake_vlm.model.vision_model.calls[1]['pixel_values']
-        self.assertTrue(torch.allclose(first_camera_pixels, torch.full((2, 3, 2, 2), -0.5)))
-        self.assertTrue(torch.allclose(second_camera_pixels, torch.full((2, 3, 2, 2), 0.5)))
-
-        # Last token is padded last state projected from [4,5,6,0,0] and [10,11,12,0,0].
-        self.assertTrue(torch.allclose(out[0, -1], torch.tensor([4.0, 5.0, 6.0, 0.0])))
-        self.assertTrue(torch.allclose(out[1, -1], torch.tensor([10.0, 11.0, 12.0, 0.0])))
-
-        pad_mask, att_mask = encoder.last_prefix_pad_mask, encoder.last_prefix_att_mask
-        self.assertEqual(pad_mask.shape, (2, 9))
-        self.assertEqual(att_mask.shape, (2, 9))
-        self.assertTrue(torch.all(att_mask[:, :8] == 0))
-        self.assertTrue(torch.all(att_mask[:, -1] == 1))
-
-    def test_forward_can_run_cropped_frozen_text_model_over_prefix_tokens(self):
-        from roboimi.vla.models.backbones.smolvla_prefix_encoder import SmolVLAPrefixEncoder
-
-        fake_vlm = _FakeVLM()
-        fake_tokenizer = _FakeTokenizer()
-        encoder = SmolVLAPrefixEncoder(
-            vlm=fake_vlm,
-            tokenizer=fake_tokenizer,
-            num_vlm_layers=2,
-            freeze_vlm=True,
-            max_state_dim=5,
-            tokenizer_max_length=4,
-            camera_names=('r_vis',),
-            resize_imgs_with_padding=None,
-            run_text_model=True,
-        )
-        images = {
-            'r_vis': torch.full((2, 1, 3, 2, 2), 0.25),
-        }
-        state = torch.tensor([[[1.0, 2.0, 3.0]], [[4.0, 5.0, 6.0]]])
-
-        out = encoder(images, state=state, task=['pick', 'place'])
-
-        # 1 camera * 2 image tokens + 4 language tokens + 1 state token.
-        self.assertEqual(out.shape, (2, 7, 4))
-        self.assertEqual(len(fake_vlm.model.text_model.layers), 2)
-        self.assertEqual(len(fake_vlm.model.text_model.forward_calls), 1)
-        text_call = fake_vlm.model.text_model.forward_calls[-1]
-        self.assertEqual(tuple(text_call['attention_mask'].shape), (2, 1, 7, 7))
-        self.assertEqual(tuple(text_call['position_ids'].shape), (2, 7))
-        self.assertTrue(torch.all(out != 0))
-
-    def test_frozen_vlm_still_backpropagates_through_text_model_to_state_projection(self):
-        from roboimi.vla.models.backbones.smolvla_prefix_encoder import SmolVLAPrefixEncoder
-
-        fake_vlm = _FakeVLM()
-        encoder = SmolVLAPrefixEncoder(
-            vlm=fake_vlm,
-            tokenizer=_FakeTokenizer(),
-            num_vlm_layers=3,
-            freeze_vlm=True,
-            train_state_proj=True,
-            max_state_dim=5,
-            tokenizer_max_length=4,
-            camera_names=('r_vis',),
-            resize_imgs_with_padding=None,
-            run_text_model=True,
-        )
-        images = {
-            'r_vis': torch.full((2, 1, 3, 2, 2), 0.25),
-        }
-        state = torch.tensor([[[1.0, 2.0, 3.0]], [[4.0, 5.0, 6.0]]])
-
-        out = encoder(images, state=state, task=['pick', 'place'])
-        out[:, -1].sum().backward()
-
-        self.assertIsNotNone(encoder.state_proj.weight.grad)
-        self.assertGreater(float(encoder.state_proj.weight.grad.abs().sum()), 0.0)
-        self.assertTrue(all(param.grad is None for param in fake_vlm.parameters()))
-
-    def test_embed_prefix_rejects_missing_camera_and_task_batch_mismatch(self):
-        from roboimi.vla.models.backbones.smolvla_prefix_encoder import SmolVLAPrefixEncoder
-
-        encoder = SmolVLAPrefixEncoder(
-            vlm=_FakeVLM(),
-            tokenizer=_FakeTokenizer(),
-            num_vlm_layers=3,
-            camera_names=('r_vis', 'top'),
-            resize_imgs_with_padding=None,
-            max_state_dim=5,
-            tokenizer_max_length=4,
-        )
-        images = {'r_vis': torch.rand(2, 1, 3, 2, 2)}
-        state = torch.rand(2, 1, 3)
-        with self.assertRaisesRegex(ValueError, 'missing.*top'):
-            encoder.embed_prefix(images=images, state=state, task=['a', 'b'])
-        images['top'] = torch.rand(2, 1, 3, 2, 2)
-        with self.assertRaisesRegex(ValueError, 'task batch'):
-            encoder.embed_prefix(images=images, state=state, task=['only one'])
-
-
-if __name__ == '__main__':
-    unittest.main()

tests/test_train_vla_rollout_validation.py

@@ -14,8 +14,6 @@ from roboimi.demos.vla_scripts import eval_vla, train_vla
 
 
 class _FakeDataset:
-    available_episode_indices = [0, 1]
-
     def __len__(self):
         return 4
 
@@ -31,13 +29,6 @@ class _FakeLoader:
         return iter(self._batches)
 
 
-class _FakeValDataset(_FakeDataset):
-    available_episode_indices = [1]
-
-    def __len__(self):
-        return 2
-
-
 class _FakeOptimizer:
     def __init__(self, lr=1e-3):
         self.param_groups = [{'lr': lr}]
@@ -100,16 +91,6 @@ class _FakeAgent(nn.Module):
         return {}
 
 
-class _CapturingAgent(_FakeAgent):
-    def __init__(self):
-        super().__init__()
-        self.compute_loss_inputs = []
-
-    def compute_loss(self, agent_input):
-        self.compute_loss_inputs.append(agent_input)
-        return (self.weight - torch.tensor(0.5)).pow(2)
-
-
 class _SequentialLossAgent(nn.Module):
     def __init__(self, losses):
         super().__init__()
@@ -169,94 +150,6 @@ class _FakeEvalEnv:
 
 
 class TrainVLARolloutValidationTest(unittest.TestCase):
-    def test_run_training_passes_variable_batch_task_to_agent_input(self):
-        cfg = OmegaConf.create(
-            {
-                'train': {
-                    'device': 'cpu',
-                    'batch_size': 2,
-                    'num_workers': 0,
-                    'val_split': 0.0,
-                    'seed': 0,
-                    'lr': 1e-3,
-                    'max_steps': 1,
-                    'log_freq': 100,
-                    'save_freq': 1000,
-                    'warmup_steps': 1,
-                    'scheduler_type': 'constant',
-                    'min_lr': 0.0,
-                    'grad_clip': 1.0,
-                    'weight_decay': 0.0,
-                    'pretrained_ckpt': None,
-                    'resume_ckpt': None,
-                    'use_swanlab': False,
-                    'rollout_val_freq_epochs': 0,
-                    'rollout_validate_on_checkpoint': False,
-                    'rollout_num_episodes': 1,
-                },
-                'data': {
-                    'camera_names': ['front'],
-                    'dataset_dir': 'unused',
-                },
-                'agent': {
-                    '_target_': 'fake.agent',
-                    'normalization_type': 'min_max',
-                },
-                'eval': {
-                    'ckpt_path': 'unused.pt',
-                    'num_episodes': 1,
-                    'max_timesteps': 1,
-                    'device': 'cpu',
-                    'task_name': 'sim_transfer',
-                    'camera_names': ['front'],
-                    'use_smoothing': False,
-                    'smooth_alpha': 0.3,
-                    'verbose_action': False,
-                    'headless': True,
-                },
-                'experiment': {},
-            }
-        )
-        agent = _CapturingAgent()
-        batch_task = ['pick the red cube', 'insert the peg into the socket']
-
-        def fake_instantiate(config_node, **_kwargs):
-            if config_node is cfg.data:
-                return _FakeDataset()
-            if config_node is cfg.agent:
-                return agent
-            raise AssertionError(f'unexpected instantiate config: {config_node!r}')
-
-        def fake_dataloader(_dataset, *, shuffle, **_kwargs):
-            del shuffle, _kwargs
-            return _FakeLoader(
-                {
-                    'observation.front': torch.zeros(2, 2, 3, 4, 4),
-                    'observation.state': torch.zeros(2, 2, 4),
-                    'action': torch.zeros(2, 8, 2),
-                    'action_is_pad': torch.zeros(2, 8, dtype=torch.bool),
-                    'task': list(batch_task),
-                },
-                length=1,
-            )
-
-        with tempfile.TemporaryDirectory() as tempdir:
-            previous_cwd = os.getcwd()
-            try:
-                os.chdir(tempdir)
-                with mock.patch.object(train_vla, 'instantiate', side_effect=fake_instantiate), \
-                     mock.patch.object(train_vla, 'DataLoader', side_effect=fake_dataloader), \
-                     mock.patch.object(train_vla, 'build_training_optimizer', return_value=_FakeOptimizer(cfg.train.lr)), \
-                     mock.patch.object(train_vla, 'get_lr_schedule_with_warmup', return_value=_FakeScheduler()), \
-                     mock.patch.object(train_vla, 'tqdm', side_effect=lambda iterable, **kwargs: _FakeProgressBar(iterable)), \
-                     mock.patch.object(train_vla.torch, 'save', return_value=None):
-                    train_vla._run_training(cfg)
-            finally:
-                os.chdir(previous_cwd)
-
-        self.assertEqual(len(agent.compute_loss_inputs), 1)
-        self.assertEqual(agent.compute_loss_inputs[0]['task'], batch_task)
-
     def test_default_train_config_uses_full_dataset_and_epoch_rollout_validation(self):
         cfg = OmegaConf.load(Path('roboimi/vla/conf/config.yaml'))
 
@@ -269,40 +162,6 @@ class TrainVLARolloutValidationTest(unittest.TestCase):
         self.assertIsNone(cfg.train.rollout_num_workers)
         self.assertIsNone(cfg.train.rollout_cuda_devices)
 
-    def test_explicit_val_episode_indices_builds_held_out_dataset(self):
-        cfg = OmegaConf.create(
-            {
-                'train': {
-                    'val_episode_indices': [1],
-                    'val_split': 0.0,
-                    'seed': 42,
-                },
-                'data': {},
-            }
-        )
-        instantiate_calls = []
-
-        def fake_instantiate(config_node, **kwargs):
-            del config_node
-            instantiate_calls.append(dict(kwargs))
-            if kwargs.get('episode_indices') == [1]:
-                return _FakeValDataset()
-            return _FakeDataset()
-
-        with mock.patch.object(train_vla, 'instantiate', side_effect=fake_instantiate):
-            dataset, train_dataset, val_dataset, explicit = train_vla.build_train_val_datasets(
-                cfg,
-                dataset_image_resize_shape=None,
-            )
-
-        self.assertIsInstance(dataset, _FakeDataset)
-        self.assertIsInstance(train_dataset, _FakeDataset)
-        self.assertIsInstance(val_dataset, _FakeValDataset)
-        self.assertEqual(explicit, [1])
-        self.assertEqual(instantiate_calls[1]['episode_indices'], [0])
-        self.assertEqual(instantiate_calls[2]['episode_indices'], [1])
-
-
     def test_run_training_rollout_validation_propagates_gpu_parallel_settings(self):
         cfg = OmegaConf.create(
             {
@@ -377,7 +236,13 @@ class TrainVLARolloutValidationTest(unittest.TestCase):
             previous_cwd = os.getcwd()
             try:
                 os.chdir(tempdir)
-                with mock.patch.object(train_vla, 'instantiate', side_effect=fake_instantiate),                      mock.patch.object(train_vla, 'DataLoader', side_effect=fake_dataloader),                      mock.patch.object(train_vla, 'build_training_optimizer', return_value=_FakeOptimizer(cfg.train.lr)),                      mock.patch.object(train_vla, 'get_lr_schedule_with_warmup', return_value=_FakeScheduler()),                      mock.patch.object(train_vla, 'tqdm', side_effect=lambda iterable, **kwargs: _FakeProgressBar(iterable)),                      mock.patch.object(train_vla.torch, 'save', return_value=None),                      mock.patch.object(eval_vla, '_run_eval', rollout_mock, create=True):
+                with mock.patch.object(train_vla, 'instantiate', side_effect=fake_instantiate), \
+                     mock.patch.object(train_vla, 'DataLoader', side_effect=fake_dataloader), \
+                     mock.patch.object(train_vla, 'build_training_optimizer', return_value=_FakeOptimizer(cfg.train.lr)), \
+                     mock.patch.object(train_vla, 'get_lr_schedule_with_warmup', return_value=_FakeScheduler()), \
+                     mock.patch.object(train_vla, 'tqdm', side_effect=lambda iterable, **kwargs: _FakeProgressBar(iterable)), \
+                     mock.patch.object(train_vla.torch, 'save', return_value=None), \
+                     mock.patch.object(eval_vla, '_run_eval', rollout_mock, create=True):
                     train_vla._run_training(cfg)
             finally:
                 os.chdir(previous_cwd)
@@ -394,23 +259,6 @@ class TrainVLARolloutValidationTest(unittest.TestCase):
         self.assertTrue(rollout_cfg.eval.save_summary_json)
         self.assertTrue(rollout_cfg.eval.save_trajectory_image)
 
-    def test_resolve_dataset_image_resize_shape_prefers_agent_top_level_override(self):
-        cfg = OmegaConf.create(
-            {
-                'agent': {
-                    'dataset_image_resize_shape': None,
-                    'vision_backbone': {
-                        'dataset_image_resize_shape': [256, 256],
-                    },
-                },
-                'data': {
-                    'image_resize_shape': [224, 224],
-                },
-            }
-        )
-
-        self.assertIsNone(train_vla._resolve_dataset_image_resize_shape(cfg))
-
     def test_training_passes_backbone_image_resize_override_to_dataset_instantiation(self):
         cfg = OmegaConf.create(
             {
@@ -497,93 +345,6 @@ class TrainVLARolloutValidationTest(unittest.TestCase):
         self.assertIn('image_resize_shape', captured_dataset_kwargs)
         self.assertIsNone(captured_dataset_kwargs['image_resize_shape'])
 
-    def test_training_passes_condition_encoder_image_resize_override_to_dataset_instantiation(self):
-        cfg = OmegaConf.create(
-            {
-                'agent': {
-                    'condition_encoder': {
-                        'dataset_image_resize_shape': None,
-                    },
-                    'normalization_type': 'min_max',
-                },
-                'data': {
-                    'dataset_dir': 'unused',
-                    'camera_names': ['front'],
-                    'image_resize_shape': [224, 224],
-                },
-                'train': {
-                    'batch_size': 2,
-                    'lr': 1e-4,
-                    'max_steps': 0,
-                    'device': 'cpu',
-                    'disable_cudnn': False,
-                    'num_workers': 0,
-                    'val_split': 0.0,
-                    'seed': 42,
-                    'log_freq': 1,
-                    'save_freq': 10,
-                    'use_swanlab': False,
-                    'rollout_val_freq_epochs': 0,
-                    'rollout_validate_on_checkpoint': False,
-                    'rollout_num_episodes': 1,
-                    'warmup_steps': 1,
-                    'scheduler_type': 'constant',
-                    'min_lr': 1e-6,
-                    'weight_decay': 1e-5,
-                    'grad_clip': 1.0,
-                    'pretrained_ckpt': None,
-                },
-                'eval': {
-                    'ckpt_path': 'unused.pt',
-                    'num_episodes': 1,
-                    'headless': True,
-                    'device': 'cpu',
-                    'verbose_action': False,
-                },
-                'experiment': {},
-            }
-        )
-        captured_dataset_kwargs = {}
-
-        def fake_instantiate(config_node, **kwargs):
-            if config_node is cfg.data:
-                captured_dataset_kwargs.update(kwargs)
-                return _FakeDataset()
-            if config_node is cfg.agent:
-                return _FakeAgent()
-            raise AssertionError(f'unexpected instantiate config: {config_node!r}')
-
-        def fake_dataloader(_dataset, *, shuffle, **_kwargs):
-            del shuffle, _kwargs
-            return _FakeLoader(
-                {
-                    'observation.front': torch.zeros(1, 3, 2, 2),
-                    'observation.state': torch.zeros(1, 4),
-                    'action': torch.zeros(1, 2),
-                    'action_is_pad': torch.zeros(1, 1, dtype=torch.bool),
-                },
-                length=1,
-            )
-
-        with tempfile.TemporaryDirectory() as tempdir:
-            previous_cwd = os.getcwd()
-            try:
-                os.chdir(tempdir)
-                with mock.patch.object(train_vla, 'instantiate', side_effect=fake_instantiate), \
-                     mock.patch.object(train_vla, 'DataLoader', side_effect=fake_dataloader), \
-                     mock.patch.object(train_vla, 'build_training_optimizer', return_value=_FakeOptimizer(cfg.train.lr)), \
-                     mock.patch.object(train_vla, 'get_lr_schedule_with_warmup', return_value=_FakeScheduler()), \
-                     mock.patch.object(train_vla, 'tqdm', side_effect=lambda iterable, **kwargs: _FakeProgressBar(iterable)), \
-                     mock.patch.object(train_vla, '_init_swanlab', return_value=None), \
-                     mock.patch.object(train_vla, '_finish_swanlab', return_value=None), \
-                     mock.patch.object(train_vla.torch, 'save', return_value=None):
-                    train_vla._run_training(cfg)
-            finally:
-                os.chdir(previous_cwd)
-
-        self.assertIn('image_resize_shape', captured_dataset_kwargs)
-        self.assertIsNone(captured_dataset_kwargs['image_resize_shape'])
-
     def test_eval_main_delegates_to_plain_run_eval_helper(self):
         cfg = OmegaConf.create(
             {

tests/test_train_vla_swanlab_logging.py

@@ -41,6 +41,19 @@ class FakeDataset:
         return 4
 
 
+class SplitAwareFakeDataset(FakeDataset):
+    def __init__(self, episode_indices=None):
+        self.episode_indices = None if episode_indices is None else list(episode_indices)
+        if self.episode_indices is None:
+            self.episodes = {0: [0], 1: [1], 2: [2]}
+        else:
+            self.episodes = {idx: [idx] for idx in self.episode_indices}
+
+    @property
+    def available_episode_indices(self):
+        return sorted(self.episodes.keys())
+
+
 class FakeLoader:
     def __init__(self, batch):
         self.batch = batch
@@ -114,6 +127,26 @@ class FakeAgent(nn.Module):
         return {}
 
 
+class RecordingAgent(FakeAgent):
+    def __init__(self):
+        super().__init__()
+        self.seen_inputs = []
+
+    def compute_loss(self, agent_input):
+        self.seen_inputs.append(agent_input)
+        return super().compute_loss(agent_input)
+
+    def predict_action_chunk(self, agent_input):
+        self.seen_inputs.append({'predict_action_chunk': agent_input})
+        return torch.ones_like(agent_input['action'])
+
+
+class ShapeMixedFakeAgent(FakeAgent):
+    def __init__(self):
+        super().__init__()
+        self.bias = nn.Parameter(torch.zeros(2))
+
+
 class FakeSwanLab:
     def __init__(self, init_error=None, log_errors=None, finish_error=None, image_errors=None):
         self.init_error = init_error
@@ -339,6 +372,8 @@ class TrainVLASwanLabLoggingTest(unittest.TestCase):
                 batch_size=2,
                 num_workers=0,
                 val_split=0.25,
+                val_episode_indices=None,
+                action_mse_val_freq_epochs=0,
                 seed=0,
                 lr=1e-3,
                 max_steps=2,
@@ -388,6 +423,18 @@ class TrainVLASwanLabLoggingTest(unittest.TestCase):
             'action_is_pad': torch.zeros(1, 1, dtype=torch.bool),
         }
 
+    def _make_lewm_batch(self):
+        batch = self._make_batch()
+        batch.update(
+            {
+                'lewm.observation.front': torch.ones(1, 3, 2, 2),
+                'lewm.observation.state': torch.ones(1, 4),
+                'lewm.future.front': torch.full((1, 3, 2, 2), 2.0),
+                'lewm.future.state': torch.full((1, 4), 2.0),
+            }
+        )
+        return batch
+
     def _loader_factory(self):
         train_batch = self._make_batch()
         val_batch = self._make_batch()
@@ -397,6 +444,15 @@ class TrainVLASwanLabLoggingTest(unittest.TestCase):
 
         return factory
 
+    def _lewm_loader_factory(self):
+        train_batch = self._make_lewm_batch()
+        val_batch = self._make_lewm_batch()
+
+        def factory(_dataset, *, shuffle, **_kwargs):
+            return FakeLoader(train_batch if shuffle else val_batch)
+
+        return factory
+
     def test_run_training_logs_metrics_and_checkpoint_paths_to_swanlab(self):
         module = self._load_train_vla_module()
         run_training = self._get_run_training(module)
@@ -442,6 +498,8 @@ class TrainVLASwanLabLoggingTest(unittest.TestCase):
                         'batch_size': 2,
                         'num_workers': 0,
                         'val_split': 0.25,
+                        'val_episode_indices': None,
+                        'action_mse_val_freq_epochs': 0,
                         'seed': 0,
                         'lr': 1e-3,
                         'max_steps': 2,
@@ -487,6 +545,95 @@ class TrainVLASwanLabLoggingTest(unittest.TestCase):
         self.assertTrue(final_payload['final/best_checkpoint_path'].endswith('checkpoints/vla_model_best.pt'))
         self.assertEqual(fake_swanlab.finish_calls, 1)
 
+    def test_run_training_passes_lewm_history_and_future_batches_into_agent_input(self):
+        module = self._load_train_vla_module()
+        run_training = self._get_run_training(module)
+        cfg = self._make_cfg(use_swanlab=False)
+        cfg.train.max_steps = 1
+        cfg.train.save_freq = 100
+        agent = RecordingAgent()
+
+        def fake_instantiate(config_node, **_kwargs):
+            if config_node is cfg.data:
+                return FakeDataset()
+            if config_node is cfg.agent:
+                return agent
+            raise AssertionError(f'unexpected instantiate config: {config_node!r}')
+
+        with tempfile.TemporaryDirectory() as tempdir:
+            previous_cwd = os.getcwd()
+            try:
+                os.chdir(tempdir)
+                with mock.patch.object(module, 'instantiate', side_effect=fake_instantiate), \
+                     mock.patch.object(module, 'DataLoader', side_effect=self._lewm_loader_factory()), \
+                     mock.patch.object(module, 'get_lr_schedule_with_warmup', return_value=FakeScheduler()), \
+                     mock.patch.object(module, 'tqdm', side_effect=lambda iterable, **kwargs: FakeProgressBar(iterable)), \
+                     mock.patch.object(module.torch, 'save', return_value=None):
+                    run_training(cfg)
+            finally:
+                os.chdir(previous_cwd)
+
+        self.assertGreaterEqual(len(agent.seen_inputs), 1)
+        first_input = agent.seen_inputs[0]
+        self.assertIn('lewm_images', first_input)
+        self.assertIn('lewm_qpos', first_input)
+        self.assertIn('lewm_future_images', first_input)
+        self.assertIn('lewm_future_qpos', first_input)
+        self.assertIn('front', first_input['lewm_images'])
+        self.assertIn('front', first_input['lewm_future_images'])
+
+    def test_run_training_logs_epoch_action_mse_for_held_out_val_episode(self):
+        module = self._load_train_vla_module()
+        run_training = self._get_run_training(module)
+        cfg = self._make_cfg()
+        cfg.train.max_steps = 1
+        cfg.train.save_freq = 100
+        cfg.train.val_split = 0.0
+        cfg.train.val_episode_indices = [2]
+        cfg.train.action_mse_val_freq_epochs = 1
+        agent = RecordingAgent()
+        fake_swanlab = FakeSwanLab()
+        real_import_module = importlib.import_module
+
+        def fake_instantiate(config_node, **kwargs):
+            if config_node is cfg.data:
+                return SplitAwareFakeDataset(kwargs.get('episode_indices'))
+            if config_node is cfg.agent:
+                return agent
+            raise AssertionError(f'unexpected instantiate config: {config_node!r}')
+
+        def fake_loader_factory(dataset, *, shuffle, **_kwargs):
+            action_value = 0.0 if shuffle else 2.0
+            batch = {
+                'observation.front': torch.zeros(1, 3, 2, 2),
+                'observation.state': torch.zeros(1, 4),
+                'action': torch.full((1, 1, 2), action_value),
+                'action_is_pad': torch.zeros(1, 1, dtype=torch.bool),
+            }
+            return FakeLoader(batch)
+
+        def fake_import_module(name, package=None):
+            if name == 'swanlab':
+                return fake_swanlab
+            return real_import_module(name, package)
+
+        with tempfile.TemporaryDirectory() as tempdir:
+            previous_cwd = os.getcwd()
+            try:
+                os.chdir(tempdir)
+                with mock.patch.object(module, 'instantiate', side_effect=fake_instantiate), \
+                     mock.patch.object(module, 'DataLoader', side_effect=fake_loader_factory), \
+                     mock.patch.object(module, 'get_lr_schedule_with_warmup', return_value=FakeScheduler()), \
+                     mock.patch.object(module, 'tqdm', side_effect=lambda iterable, **kwargs: FakeProgressBar(iterable)), \
+                     mock.patch.object(module.torch, 'save', return_value=None), \
+                     mock.patch.object(module.importlib, 'import_module', side_effect=fake_import_module):
+                    run_training(cfg)
+            finally:
+                os.chdir(previous_cwd)
+
+        logged_keys = set().union(*(payload.keys() for payload, _ in fake_swanlab.log_calls))
+        self.assertIn('val/action_mse', logged_keys)
+
     def test_run_training_skips_swanlab_when_disabled(self):
         module = self._load_train_vla_module()
         run_training = self._get_run_training(module)
@@ -668,6 +815,52 @@ class TrainVLASwanLabLoggingTest(unittest.TestCase):
         self.assertTrue(final_payload['final/best_checkpoint_path'].endswith('checkpoints/vla_model_best.pt'))
         self.assertFalse(any(path.endswith('checkpoints/vla_model_best.pt') for path in saved_paths))
 
+    def test_run_training_pretrained_ckpt_loads_matching_keys_even_if_some_shapes_mismatch(self):
+        module = self._load_train_vla_module()
+        run_training = self._get_run_training(module)
+        cfg = self._make_cfg(use_swanlab=False)
+        cfg.train.max_steps = 0
+        cfg.train.save_freq = 100
+        cfg.train.pretrained_ckpt = 'pretrained.pt'
+        agent = ShapeMixedFakeAgent()
+
+        def fake_instantiate(config_node, **_kwargs):
+            if config_node is cfg.data:
+                return FakeDataset()
+            if config_node is cfg.agent:
+                return agent
+            raise AssertionError(f'unexpected instantiate config: {config_node!r}')
+
+        def fake_torch_load(path, map_location=None):
+            del map_location
+            if Path(path).name != 'pretrained.pt':
+                raise AssertionError(f'unexpected load path: {path}')
+            return {
+                'model_state_dict': {
+                    'weight': torch.tensor(3.0),
+                    'bias': torch.tensor([1.0, 2.0, 3.0]),
+                },
+                'step': 123,
+                'loss': 0.5,
+            }
+
+        with tempfile.TemporaryDirectory() as tempdir:
+            previous_cwd = os.getcwd()
+            try:
+                os.chdir(tempdir)
+                Path('pretrained.pt').write_bytes(b'pretend')
+                with mock.patch.object(module, 'instantiate', side_effect=fake_instantiate), \
+                     mock.patch.object(module, 'DataLoader', side_effect=self._loader_factory()), \
+                     mock.patch.object(module, 'get_lr_schedule_with_warmup', return_value=FakeScheduler()), \
+                     mock.patch.object(module, 'tqdm', side_effect=lambda iterable, **kwargs: FakeProgressBar(iterable)), \
+                     mock.patch.object(module.torch, 'save', return_value=None), \
+                     mock.patch.object(module.torch, 'load', side_effect=fake_torch_load):
+                    run_training(cfg)
+            finally:
+                os.chdir(previous_cwd)
+
+        self.assertAlmostEqual(agent.weight.item(), 3.0, places=6)
+
     def test_run_training_resume_ignores_best_checkpoint_without_rollout_metric(self):
         module = self._load_train_vla_module()
         run_training = self._get_run_training(module)

tests/test_train_vla_transformer_optimizer.py

@@ -39,17 +39,6 @@ class FakeLoader:
         return iter(())
 
 
-class FakeTqdm:
-    def __init__(self, iterable, **_kwargs):
-        self.iterable = iterable
-
-    def __iter__(self):
-        return iter(self.iterable)
-
-    def set_postfix(self, *_args, **_kwargs):
-        return None
-
-
 class FakeScheduler:
     def state_dict(self):
         return {}
@@ -57,18 +46,13 @@ class FakeScheduler:
     def load_state_dict(self, state_dict):
         return None
 
-    def step(self):
-        return None
-
 
 class RecordingAdamW:
     created = []
 
-    def __init__(self, params, lr, weight_decay, betas=(0.9, 0.999), eps=1e-8):
+    def __init__(self, params, lr, weight_decay):
         self.lr = lr
         self.weight_decay = weight_decay
-        self.betas = betas
-        self.eps = eps
         self.param_groups = self._normalize_param_groups(params, lr, weight_decay)
         RecordingAdamW.created.append(self)
 
@@ -95,12 +79,6 @@ class RecordingAdamW:
     def load_state_dict(self, state_dict):
         return None
 
-    def zero_grad(self):
-        return None
-
-    def step(self):
-        return None
-
 
 class RecordingTransformerHead(nn.Module):
     def __init__(self):
@@ -346,7 +324,7 @@ class TrainVLATransformerOptimizerTest(unittest.TestCase):
                      mock.patch.object(module, 'get_lr_schedule_with_warmup', return_value=FakeScheduler()), \
                      mock.patch.object(module, 'AdamW', RecordingAdamW), \
                      mock.patch.object(module.torch, 'save', return_value=None), \
-                     mock.patch.object(module, 'tqdm', side_effect=lambda iterable, **kwargs: FakeTqdm(iterable, **kwargs)):
+                     mock.patch.object(module, 'tqdm', side_effect=lambda iterable, **kwargs: iterable):
                     module.main(cfg)
             finally:
                 os.chdir(previous_cwd)
@@ -426,7 +404,7 @@ class TrainVLATransformerOptimizerTest(unittest.TestCase):
                      mock.patch.object(module, 'get_lr_schedule_with_warmup', return_value=FakeScheduler()), \
                      mock.patch.object(module, 'AdamW', RecordingAdamW), \
                      mock.patch.object(module.torch, 'save', return_value=None), \
-                     mock.patch.object(module, 'tqdm', side_effect=lambda iterable, **kwargs: FakeTqdm(iterable, **kwargs)):
+                     mock.patch.object(module, 'tqdm', side_effect=lambda iterable, **kwargs: iterable):
                     module.main(cfg)
             finally:
                 os.chdir(previous_cwd)
@@ -444,159 +422,3 @@ class TrainVLATransformerOptimizerTest(unittest.TestCase):
 
 if __name__ == '__main__':
     unittest.main()
-
-class TrainVLASmolVLAOptimizerTest(unittest.TestCase):
-    def test_build_training_optimizer_excludes_frozen_vlm_parameters_and_keeps_state_proj_and_head(self):
-        module = TrainVLATransformerOptimizerTest()._load_train_vla_module()
-
-        class _Head(nn.Module):
-            def __init__(self):
-                super().__init__()
-                self.proj = nn.Linear(2, 2)
-
-            def get_optim_groups(self, weight_decay):
-                return [{'params': list(self.parameters()), 'weight_decay': weight_decay}]
-
-        class _ConditionEncoder(nn.Module):
-            def __init__(self):
-                super().__init__()
-                self.vlm = nn.Linear(2, 2)
-                for param in self.vlm.parameters():
-                    param.requires_grad = False
-                self.state_proj = nn.Linear(2, 2)
-
-        class _Agent(nn.Module):
-            def __init__(self):
-                super().__init__()
-                self.noise_pred_net = _Head()
-                self.condition_encoder = _ConditionEncoder()
-
-        agent = _Agent()
-        with mock.patch.object(module, 'AdamW', RecordingAdamW):
-            optimizer = module.build_training_optimizer(agent, lr=1e-4, weight_decay=0.01)
-
-        names_by_param_id = {id(param): name for name, param in agent.named_parameters()}
-        optimizer_names = {
-            names_by_param_id[id(param)]
-            for group in optimizer.param_groups
-            for param in group['params']
-        }
-        self.assertIn('condition_encoder.state_proj.weight', optimizer_names)
-        self.assertIn('condition_encoder.state_proj.bias', optimizer_names)
-        self.assertIn('noise_pred_net.proj.weight', optimizer_names)
-        self.assertNotIn('condition_encoder.vlm.weight', optimizer_names)
-        self.assertNotIn('condition_encoder.vlm.bias', optimizer_names)
-
-    def test_smolvla_native_training_preset_overrides_optimizer_scheduler_and_grad_clip(self):
-        module = TrainVLATransformerOptimizerTest()._load_train_vla_module()
-
-        class _NativeAgent(nn.Module):
-            def __init__(self):
-                super().__init__()
-                self.model = nn.Linear(2, 2)
-
-            def to(self, device):
-                return self
-
-            def get_normalization_stats(self):
-                return {}
-
-        agent = _NativeAgent()
-        cfg = TrainVLATransformerOptimizerTest()._make_cfg()
-        cfg.agent = AttrDict(_target_='roboimi.vla.agent_smolvla_native.SmolVLANativeAgent')
-        cfg.train.lr = 9e-4
-        cfg.train.max_steps = 1
-        cfg.train.weight_decay = 0.123
-        cfg.train.grad_clip = 1.0
-        cfg.train.warmup_steps = 7
-        cfg.train.scheduler_type = 'constant'
-        cfg.train.min_lr = 1e-7
-
-        scheduler_calls = []
-        clip_calls = []
-
-        def fake_instantiate(config_node, **_kwargs):
-            if config_node is cfg.data:
-                return FakeDataset()
-            if config_node is cfg.agent:
-                return agent
-            raise AssertionError(f'unexpected instantiate config: {config_node!r}')
-
-        def fake_scheduler(*args, **kwargs):
-            scheduler_calls.append(kwargs)
-            return FakeScheduler()
-
-        def fake_clip(parameters, max_norm):
-            clip_calls.append(float(max_norm))
-            return torch.tensor(0.0)
-
-        class OneBatchLoader:
-            def __len__(self):
-                return 1
-
-            def __iter__(self):
-                batch = {
-                    'observation.front': torch.zeros(1, 1, 1, 2, 2),
-                    'observation.state': torch.zeros(1, 1, 2),
-                    'action': torch.zeros(1, 1, 2),
-                }
-                return iter([batch])
-
-        def fake_compute_loss(_batch):
-            return agent.model.weight.sum() * 0.0 + torch.tensor(1.0, requires_grad=True)
-
-        agent.compute_loss = fake_compute_loss
-
-        with tempfile.TemporaryDirectory() as tempdir:
-            previous_cwd = os.getcwd()
-            try:
-                os.chdir(tempdir)
-                with mock.patch.object(module, 'instantiate', side_effect=fake_instantiate), \
-                     mock.patch.object(module, 'DataLoader', side_effect=lambda *args, **kwargs: OneBatchLoader()), \
-                     mock.patch.object(module, 'get_lr_schedule_with_warmup', side_effect=fake_scheduler), \
-                     mock.patch.object(module, 'AdamW', RecordingAdamW), \
-                     mock.patch.object(module.torch.nn.utils, 'clip_grad_norm_', side_effect=fake_clip), \
-                     mock.patch.object(module.torch, 'save', return_value=None), \
-                     mock.patch.object(module, 'tqdm', side_effect=lambda iterable, **kwargs: FakeTqdm(iterable, **kwargs)):
-                    module.main(cfg)
-            finally:
-                os.chdir(previous_cwd)
-
-        optimizer = RecordingAdamW.created[-1]
-        self.assertEqual(optimizer.lr, 1e-4)
-        self.assertEqual(optimizer.weight_decay, 1e-10)
-        self.assertEqual(optimizer.betas, (0.9, 0.95))
-        self.assertEqual(optimizer.eps, 1e-8)
-        self.assertEqual(scheduler_calls[-1], {
-            'warmup_steps': 1000,
-            'max_steps': 1,
-            'scheduler_type': 'cosine',
-            'min_lr': 2.5e-6,
-        })
-        self.assertEqual(clip_calls, [10.0])
-
-    def test_cosine_scheduler_spans_requested_training_steps_then_clamps(self):
-        module = TrainVLATransformerOptimizerTest()._load_train_vla_module()
-        param = nn.Parameter(torch.tensor(1.0))
-        optimizer = torch.optim.SGD([param], lr=1e-4)
-        base_lr = optimizer.param_groups[0]['lr']
-
-        scheduler = module.get_lr_schedule_with_warmup(
-            optimizer,
-            warmup_steps=1000,
-            max_steps=150000,
-            scheduler_type='cosine',
-            min_lr=2.5e-6,
-        )
-
-        lr_lambda = scheduler.lr_lambdas[0]
-        observed = {}
-        for step in (0, 1, 1000, 30000, 40000, 150000, 160000):
-            observed[step] = base_lr * lr_lambda(step)
-
-        self.assertGreater(observed[1], observed[0])
-        self.assertLess(observed[1000], 1e-4)
-        self.assertGreater(observed[30000], observed[40000])
-        self.assertGreater(observed[40000], observed[150000])
-        self.assertAlmostEqual(observed[150000], 2.5e-6, places=12)
-        self.assertAlmostEqual(observed[160000], 2.5e-6, places=12)