Compare commits
22 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
 Logic
|
ebac9860fe | ||
|
Logic
|
2ac926d427 | ||
|
Logic
|
d94eb8f70b | ||
|
Logic
|
acbd7c605a | ||
|
Logic
|
73f5b6e3d9 | ||
|
Logic
|
bc4caf355b | ||
|
Logic
|
b1116e489f | ||
|
Logic
|
a2c018acce | ||
|
Logic
|
4890f54b13 | ||
|
Logic
|
5c5cb299e9 | ||
|
Logic
|
4c3646a3d5 | ||
|
Logic
|
4936cf2635 | ||
|
Logic
|
d245d64def | ||
|
Logic
|
8145c9eb62 | ||
|
Logic
|
a837a982f7 | ||
|
Logic
|
f1ede7690f | ||
|
Logic
|
06ac6c6d18 | ||
|
Logic
|
fce6839daa | ||
|
Logic
|
3eb1a83940 | ||
|
Logic
|
636290d36a | ||
|
Logic
|
4ea75966ee | ||
|
Logic
|
27f4a07632 |
@@ -1,471 +0,0 @@
|
||||
# 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`
|
||||
@@ -0,0 +1,311 @@
|
||||
# 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.
|
||||
@@ -0,0 +1,184 @@
|
||||
# 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.
|
||||
@@ -0,0 +1,316 @@
|
||||
# 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
|
||||
@@ -0,0 +1,86 @@
|
||||
# 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.
|
||||
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,200 @@
|
||||
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)
|
||||
@@ -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.25" euler="0.0 9.4 0.0 " fovy="15" resolution="1920 1200"/>
|
||||
<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"/>
|
||||
</body>
|
||||
</body>
|
||||
<body name="l_finger_left" pos="0 0.01 0.0444">
|
||||
|
||||
@@ -0,0 +1,6 @@
|
||||
<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>
|
||||
@@ -0,0 +1,19 @@
|
||||
<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>
|
||||
@@ -7,7 +7,6 @@
|
||||
<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>
|
||||
|
||||
@@ -91,3 +91,39 @@ 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])
|
||||
|
||||
|
||||
@@ -0,0 +1,184 @@
|
||||
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,
|
||||
},
|
||||
]
|
||||
@@ -2,9 +2,11 @@ import time
|
||||
import os
|
||||
import numpy as np
|
||||
from roboimi.envs.double_pos_ctrl_env import make_sim_env
|
||||
from diana_policy import TestPickAndTransferPolicy
|
||||
from roboimi.demos.diana_air_insert_policy import TestAirInsertPolicy
|
||||
from roboimi.demos.diana_policy import TestPickAndTransferPolicy
|
||||
import cv2
|
||||
from roboimi.utils.act_ex_utils import sample_transfer_pose
|
||||
from roboimi.utils.act_ex_utils import sample_air_insert_socket_peg_state, sample_transfer_pose
|
||||
from roboimi.utils.constants import SIM_TASK_CONFIGS
|
||||
from roboimi.utils.streaming_episode_writer import StreamingEpisodeWriter
|
||||
|
||||
import pathlib
|
||||
@@ -12,16 +14,34 @@ HOME_PATH = str(pathlib.Path(__file__).parent.resolve())
|
||||
DATASET_DIR = HOME_PATH + '/dataset'
|
||||
|
||||
|
||||
def main():
|
||||
task_name = 'sim_transfer'
|
||||
dataset_dir = DATASET_DIR + '/sim_transfer' #SIM_TASK_CONFIGS[task_name]['dataset_dir']
|
||||
num_episodes = 100 #SIM_TASK_CONFIGS[task_name]['num_episodes']
|
||||
def sample_task_state(task_name):
|
||||
if task_name == 'sim_transfer':
|
||||
return sample_transfer_pose()
|
||||
if task_name == 'sim_air_insert_socket_peg':
|
||||
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 = 700 #SIM_TASK_CONFIGS[task_name]['episode_len']
|
||||
camera_names = ['angle','r_vis', 'top', 'front'] #SIM_TASK_CONFIGS[task_name]['camera_names']
|
||||
episode_len = task_cfg['episode_len']
|
||||
camera_names = task_cfg['camera_names']
|
||||
image_size = (256, 256)
|
||||
if task_name == 'sim_transfer':
|
||||
if task_name in {'sim_transfer', 'sim_air_insert_socket_peg'}:
|
||||
print(task_name)
|
||||
else:
|
||||
raise NotImplementedError
|
||||
@@ -29,7 +49,7 @@ def main():
|
||||
success = []
|
||||
|
||||
env = make_sim_env(task_name)
|
||||
policy = TestPickAndTransferPolicy(inject_noise)
|
||||
policy = make_policy(task_name, inject_noise=inject_noise)
|
||||
|
||||
# 等待osmesa完全启动后再开始收集数据
|
||||
print("等待osmesa线程启动...")
|
||||
@@ -41,8 +61,8 @@ def main():
|
||||
max_reward = float('-inf')
|
||||
print(f'\n{episode_idx=}')
|
||||
print('Rollout out EE space scripted policy')
|
||||
box_pos = sample_transfer_pose()
|
||||
env.reset(box_pos)
|
||||
task_state = sample_task_state(task_name)
|
||||
env.reset(task_state)
|
||||
episode_writer = StreamingEpisodeWriter(
|
||||
dataset_path=os.path.join(dataset_dir, f'episode_{episode_idx}.hdf5'),
|
||||
max_timesteps=episode_len,
|
||||
@@ -50,7 +70,7 @@ def main():
|
||||
image_size=image_size,
|
||||
)
|
||||
for step in range(episode_len):
|
||||
raw_action = policy.predict(box_pos,step)
|
||||
raw_action = policy.predict(task_state, step)
|
||||
env.step(raw_action)
|
||||
env.render()
|
||||
sum_reward += env.rew
|
||||
|
||||
@@ -30,7 +30,10 @@ 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 sample_transfer_pose
|
||||
from roboimi.utils.act_ex_utils import (
|
||||
sample_air_insert_socket_peg_state,
|
||||
sample_transfer_pose,
|
||||
)
|
||||
from roboimi.vla.eval_utils import execute_policy_action
|
||||
|
||||
sys.path.append(os.getcwd())
|
||||
@@ -110,6 +113,7 @@ def prepare_observation(
|
||||
obs: Dict,
|
||||
camera_names: list,
|
||||
image_resize_shape: Optional[tuple[int, int]] = (224, 224),
|
||||
task_description: Optional[str] = None,
|
||||
) -> Dict:
|
||||
"""
|
||||
将环境观测转换为 agent 格式。
|
||||
@@ -136,7 +140,34 @@ def prepare_observation(
|
||||
# 转换 qpos: numpy -> tensor
|
||||
qpos = torch.from_numpy(obs['qpos']).float()
|
||||
|
||||
return {'qpos': qpos, 'images': images}
|
||||
observation = {'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]:
|
||||
@@ -156,6 +187,7 @@ 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(),
|
||||
}
|
||||
|
||||
@@ -171,6 +203,8 @@ 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(
|
||||
@@ -199,7 +233,10 @@ def _prepare_local_policy_batch(
|
||||
).unsqueeze(0)
|
||||
for camera_name in ordered_camera_names
|
||||
}
|
||||
return {'qpos': batch_qpos, 'images': batch_images}
|
||||
batch = {'qpos': batch_qpos, 'images': batch_images}
|
||||
if queues.get('task'):
|
||||
batch['task'] = [list(queues['task'])[-1]]
|
||||
return batch
|
||||
|
||||
|
||||
def _enqueue_predicted_actions(
|
||||
@@ -207,13 +244,14 @@ 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
|
||||
start = int(obs_horizon) - 1 if action_chunk_start is None else int(action_chunk_start)
|
||||
end = start + int(num_action_steps)
|
||||
executable_actions = predicted_actions[:, start:end]
|
||||
for action_index in range(executable_actions.shape[1]):
|
||||
@@ -223,23 +261,29 @@ def _enqueue_predicted_actions(
|
||||
|
||||
|
||||
def _serialize_policy_batch(batch: Dict[str, torch.Tensor]) -> dict[str, Any]:
|
||||
return {
|
||||
serialized = {
|
||||
'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]:
|
||||
return {
|
||||
deserialized = {
|
||||
'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:
|
||||
@@ -275,6 +319,7 @@ 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)
|
||||
@@ -284,6 +329,7 @@ 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
|
||||
@@ -330,6 +376,7 @@ 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']:
|
||||
@@ -964,6 +1011,14 @@ 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)
|
||||
@@ -1004,6 +1059,8 @@ 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'],
|
||||
@@ -1030,17 +1087,19 @@ def _run_eval_episode_plans(
|
||||
episode_max_rewards = []
|
||||
for plan in episode_plans:
|
||||
episode_idx = int(plan['episode_index'])
|
||||
box_pos = plan.get('box_pos')
|
||||
if box_pos is None:
|
||||
box_pos = sample_transfer_pose()
|
||||
box_pos = np.asarray(box_pos, dtype=np.float32)
|
||||
task_state = plan.get('box_pos')
|
||||
if task_state is None:
|
||||
task_state = _sample_task_reset_state(str(eval_cfg.task_name))
|
||||
elif isinstance(task_state, np.ndarray):
|
||||
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(box_pos)
|
||||
env.reset(task_state)
|
||||
|
||||
# 为新回合重置 rollout policy 状态
|
||||
policy_runner.reset()
|
||||
@@ -1076,7 +1135,12 @@ def _run_eval_episode_plans(
|
||||
video_recorder.write(video_frame)
|
||||
|
||||
# 准备给 agent 的观测
|
||||
observation = prepare_observation(obs, camera_names)
|
||||
observation = prepare_observation(
|
||||
obs,
|
||||
camera_names,
|
||||
image_resize_shape=image_resize_shape,
|
||||
task_description=task_description,
|
||||
)
|
||||
end_preprocess = time.perf_counter()
|
||||
|
||||
# 选择动作（本地 agent 或远端 inference server）
|
||||
@@ -1349,6 +1413,7 @@ 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,
|
||||
@@ -1357,6 +1422,7 @@ 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(
|
||||
|
||||
@@ -71,6 +71,19 @@ 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))
|
||||
@@ -118,127 +131,6 @@ 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 路径。
|
||||
@@ -275,7 +167,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: 最小学习率（用于余弦衰减）
|
||||
|
||||
@@ -288,16 +180,24 @@ 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):
|
||||
# 预热阶段：从 0 线性增加到 1
|
||||
# LeRobot CosineDecayWithWarmupSchedulerConfig 的线性预热：
|
||||
# 从一个很小的非零 LR 开始，避免首步完全为 0。
|
||||
if step < warmup_steps:
|
||||
return float(step) / float(max(1, warmup_steps))
|
||||
if step <= 0:
|
||||
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':
|
||||
# 从 1 到 min_lr_ratio 的余弦退火
|
||||
progress = float(step - warmup_steps) / float(max(1, max_steps - warmup_steps))
|
||||
cosine_decay = 0.5 * (1.0 + math.cos(math.pi * progress))
|
||||
return max(min_lr_ratio, cosine_decay)
|
||||
# 与 LeRobot SmolVLA/PI0 的 CosineDecayWithWarmupSchedulerConfig 对齐：
|
||||
# 1) 余弦衰减步数是固定的 num_decay_steps（SmolVLA 默认 30k），
|
||||
# 2) 超过 decay_steps 后 clamp 到 decay_lr，不能继续 cos() 进入下一周期；
|
||||
# 否则 150k 训练会显示成“正弦波”式反复升降。
|
||||
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
|
||||
@@ -305,7 +205,119 @@ def get_lr_schedule_with_warmup(optimizer, warmup_steps, max_steps, scheduler_ty
|
||||
return LambdaLR(optimizer, lr_lambda)
|
||||
|
||||
|
||||
def build_training_optimizer(agent, lr, weight_decay):
|
||||
def _is_smolvla_native_agent_config(agent_cfg) -> bool:
|
||||
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)
|
||||
@@ -313,7 +325,7 @@ def build_training_optimizer(agent, lr, weight_decay):
|
||||
use_head_groups = callable(get_optim_groups)
|
||||
|
||||
if not use_head_groups:
|
||||
return AdamW(trainable_params, lr=lr, weight_decay=weight_decay)
|
||||
return AdamW(trainable_params, lr=lr, weight_decay=weight_decay, betas=betas, eps=eps)
|
||||
|
||||
head_groups = []
|
||||
grouped_param_ids = set()
|
||||
@@ -355,33 +367,7 @@ def build_training_optimizer(agent, lr, weight_decay):
|
||||
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)
|
||||
|
||||
|
||||
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),
|
||||
}
|
||||
return AdamW(optim_groups, lr=lr, weight_decay=weight_decay, betas=betas, eps=eps)
|
||||
|
||||
|
||||
def _init_swanlab(cfg):
|
||||
@@ -428,7 +414,12 @@ def _init_swanlab(cfg):
|
||||
init_kwargs['experiment_name'] = run_name
|
||||
|
||||
try:
|
||||
swanlab.init(**init_kwargs)
|
||||
run = 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}"
|
||||
@@ -437,6 +428,24 @@ 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
|
||||
@@ -515,7 +524,13 @@ 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"
|
||||
@@ -527,10 +542,7 @@ def _run_training(cfg: DictConfig):
|
||||
# =========================================================================
|
||||
log.info("📦 加载数据集...")
|
||||
try:
|
||||
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_image_resize_shape = _resolve_dataset_image_resize_shape(cfg)
|
||||
dataset, train_dataset, val_dataset, explicit_val_episode_indices = (
|
||||
build_train_val_datasets(cfg, dataset_image_resize_shape)
|
||||
)
|
||||
@@ -656,23 +668,18 @@ def _run_training(cfg: DictConfig):
|
||||
try:
|
||||
checkpoint = torch.load(ckpt_path, map_location=cfg.train.device)
|
||||
|
||||
load_info = load_state_dict_ignoring_shape_mismatches(
|
||||
agent,
|
||||
# 只加载模型权重（不加载 optimizer、scheduler）
|
||||
missing_keys, unexpected_keys = agent.load_state_dict(
|
||||
checkpoint['model_state_dict'],
|
||||
strict=False # 允许部分加载（结构不完全匹配时）
|
||||
)
|
||||
|
||||
log.info(f"✅ [Finetune] 模型权重加载成功")
|
||||
|
||||
if load_info['missing_keys']:
|
||||
log.warning(
|
||||
f"⚠️ [Finetune] checkpoint 中存在本地模型没有的键 ({len(load_info['missing_keys'])} 个): "
|
||||
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]}..."
|
||||
)
|
||||
if missing_keys:
|
||||
log.warning(f"⚠️ [Finetune] 缺少的键 ({len(missing_keys)} 个): {missing_keys[:5]}...")
|
||||
if unexpected_keys:
|
||||
log.warning(f"⚠️ [Finetune] 多余的键 ({len(unexpected_keys)} 个): {unexpected_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}）")
|
||||
@@ -687,21 +694,35 @@ def _run_training(cfg: DictConfig):
|
||||
# =========================================================================
|
||||
# 4. 设置优化器与学习率调度器
|
||||
# =========================================================================
|
||||
weight_decay = float(cfg.train.get('weight_decay', 1e-5))
|
||||
grad_clip = float(cfg.train.get('grad_clip', 1.0))
|
||||
|
||||
optimizer = build_training_optimizer(agent, lr=cfg.train.lr, weight_decay=weight_decay)
|
||||
log.info(f"🔧 优化器: AdamW (学习率={cfg.train.lr}, weight_decay={weight_decay})")
|
||||
training_recipe = resolve_training_recipe(cfg)
|
||||
weight_decay = training_recipe['weight_decay']
|
||||
grad_clip = training_recipe['grad_clip']
|
||||
train_lr = training_recipe['lr']
|
||||
optimizer = build_training_optimizer(
|
||||
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 = int(cfg.train.get('warmup_steps', 500))
|
||||
scheduler_type = cfg.train.get('scheduler_type', 'cosine')
|
||||
min_lr = float(cfg.train.get('min_lr', 1e-6))
|
||||
warmup_steps = training_recipe['warmup_steps']
|
||||
scheduler_type = training_recipe['scheduler_type']
|
||||
min_lr = training_recipe['min_lr']
|
||||
|
||||
scheduler = get_lr_schedule_with_warmup(
|
||||
optimizer,
|
||||
warmup_steps=warmup_steps,
|
||||
max_steps=cfg.train.max_steps,
|
||||
max_steps=training_recipe['scheduler_steps'],
|
||||
scheduler_type=scheduler_type,
|
||||
min_lr=min_lr
|
||||
)
|
||||
@@ -795,6 +816,26 @@ 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({
|
||||
@@ -887,7 +928,6 @@ 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
|
||||
@@ -966,15 +1006,6 @@ 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)
|
||||
|
||||
# =====================================================================
|
||||
# 检查点保存与验证
|
||||
@@ -1050,6 +1081,31 @@ 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
|
||||
@@ -1057,33 +1113,6 @@ 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"
|
||||
|
||||
@@ -0,0 +1,157 @@
|
||||
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())
|
||||
@@ -52,7 +52,6 @@ class DualDianaMed(MujocoEnv):
|
||||
self.r_vis = None
|
||||
self.l_vis = None
|
||||
self.top = None
|
||||
self.angle = None
|
||||
self.front = None
|
||||
self.obs = None
|
||||
|
||||
@@ -92,7 +91,6 @@ 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)
|
||||
@@ -105,6 +103,7 @@ class DualDianaMed(MujocoEnv):
|
||||
|
||||
super().step(action)
|
||||
self.base_time = time.time() - ctrl_cur_time
|
||||
self.obs = self._get_obs()
|
||||
|
||||
|
||||
def preStep(self, action):
|
||||
@@ -168,10 +167,9 @@ 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
|
||||
|
||||
@@ -180,10 +178,9 @@ 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
|
||||
|
||||
@@ -202,14 +199,12 @@ class DualDianaMed(MujocoEnv):
|
||||
|
||||
@property
|
||||
def cam_view(self):
|
||||
if self.cam == 'top':
|
||||
return self.top
|
||||
elif self.cam == 'angle':
|
||||
return self.angle
|
||||
elif self.cam == 'r_vis':
|
||||
if 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:
|
||||
@@ -230,8 +225,6 @@ 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]
|
||||
|
||||
|
||||
@@ -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,6 +137,18 @@ 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(
|
||||
@@ -144,7 +156,7 @@ def make_sim_env(task_name, headless=False):
|
||||
is_render=not headless,
|
||||
control_freq=30,
|
||||
is_interpolate=True,
|
||||
cam_view='angle'
|
||||
cam_view='top'
|
||||
)
|
||||
return env
|
||||
else:
|
||||
@@ -170,4 +182,3 @@ if __name__ == "__main__":
|
||||
env.step(action)
|
||||
if env.is_render:
|
||||
env.render()
|
||||
|
||||
|
||||
@@ -1,267 +0,0 @@
|
||||
#!/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())
|
||||
@@ -1,5 +1,6 @@
|
||||
import numpy as np
|
||||
|
||||
|
||||
def sample_insertion_pose():
|
||||
# Peg
|
||||
x_range = [0.1, 0.2]
|
||||
@@ -36,3 +37,22 @@ 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,
|
||||
}
|
||||
|
||||
@@ -20,7 +20,14 @@ SIM_TASK_CONFIGS = {
|
||||
'dataset_dir': DATASET_DIR + '/sim_transfer',
|
||||
'num_episodes': 20,
|
||||
'episode_len': 700,
|
||||
'camera_names': ['top','r_vis','front'],
|
||||
'camera_names': ['r_vis', 'top', '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'
|
||||
},
|
||||
|
||||
@@ -52,13 +59,3 @@ 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
|
||||
|
||||
+4
-13
@@ -28,7 +28,6 @@ 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__()
|
||||
# 保存参数
|
||||
@@ -40,9 +39,6 @@ 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
|
||||
@@ -75,14 +71,11 @@ 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)
|
||||
@@ -94,11 +87,8 @@ class VLAAgent(nn.Module):
|
||||
else:
|
||||
per_token_vision_dim = int(single_cam_feat_dim) * int(num_cams)
|
||||
|
||||
self.history_condition_sequence_length = self.obs_horizon * self.condition_tokens_per_step
|
||||
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
|
||||
self.condition_sequence_length = self.obs_horizon * self.condition_tokens_per_step
|
||||
self.raw_per_step_cond_dim = per_token_vision_dim + obs_dim
|
||||
if cond_projector is None:
|
||||
self.cond_projector = None
|
||||
self.per_step_cond_dim = self.raw_per_step_cond_dim
|
||||
@@ -149,6 +139,7 @@ class VLAAgent(nn.Module):
|
||||
global_cond_dim=self.global_cond_dim
|
||||
)
|
||||
|
||||
self.state_encoder = state_encoder
|
||||
self.action_encoder = action_encoder
|
||||
|
||||
# 初始化队列（用于在线推理）
|
||||
@@ -229,7 +220,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.history_condition_sequence_length
|
||||
expected_length = self.condition_sequence_length
|
||||
if cond.shape[1] != expected_length:
|
||||
raise RuntimeError(
|
||||
f"条件序列长度不匹配: got {cond.shape[1]}, expected {expected_length}"
|
||||
|
||||
+10
-416
@@ -1,12 +1,9 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from contextlib import nullcontext
|
||||
from collections import deque
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, Mapping, Optional, Sequence
|
||||
from typing import Dict, Optional
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from roboimi.vla.agent import VLAAgent
|
||||
@@ -18,87 +15,14 @@ except ImportError: # pragma: no cover
|
||||
|
||||
|
||||
class IMFVLAAgent(VLAAgent):
|
||||
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,
|
||||
):
|
||||
def __init__(self, *args, inference_steps: int = 1, **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:
|
||||
@@ -195,251 +119,14 @@ 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, 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),
|
||||
)
|
||||
cond = self._build_cond(images, states)
|
||||
|
||||
x = actions
|
||||
e = torch.randn_like(x)
|
||||
@@ -459,109 +146,16 @@ 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]
|
||||
action_loss = (loss * mask).sum() / valid_count.clamp_min(1.0)
|
||||
loss = (loss * mask).sum() / valid_count.clamp_min(1.0)
|
||||
else:
|
||||
action_loss = loss.mean()
|
||||
|
||||
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
|
||||
loss = loss.mean()
|
||||
return loss
|
||||
|
||||
@torch.no_grad()
|
||||
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,
|
||||
):
|
||||
def predict_action(self, images, proprioception):
|
||||
batch_size = proprioception.shape[0]
|
||||
if self.lewm_predictor is not None:
|
||||
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),
|
||||
)
|
||||
proprioception = self.normalization.normalize_qpos(proprioception)
|
||||
cond = self._build_cond(images, 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)
|
||||
|
||||
@@ -0,0 +1,277 @@
|
||||
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()
|
||||
@@ -0,0 +1,398 @@
|
||||
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()
|
||||
@@ -1,74 +0,0 @@
|
||||
# @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
|
||||
@@ -1,77 +0,0 @@
|
||||
# @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
|
||||
@@ -0,0 +1,46 @@
|
||||
# @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
|
||||
@@ -0,0 +1,35 @@
|
||||
# @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]
|
||||
@@ -1,7 +0,0 @@
|
||||
_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
|
||||
@@ -0,0 +1,19 @@
|
||||
_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
|
||||
@@ -18,8 +18,6 @@ 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 # 随机种子（用于数据划分）
|
||||
|
||||
# 日志和检查点
|
||||
|
||||
@@ -11,8 +11,6 @@ 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}
|
||||
|
||||
# ====================
|
||||
# 相机配置
|
||||
|
||||
@@ -9,6 +9,7 @@ server_startup_timeout_s: 300.0 # parent 等待 inference server 就绪的超时
|
||||
max_timesteps: 700 # 每回合最大时间步
|
||||
device: ${train.device} # 与训练保持一致
|
||||
task_name: "sim_transfer" # 环境任务名称
|
||||
task_description: null # 可选语言指令；环境 obs 没有 task 时注入给语言条件策略
|
||||
|
||||
# ====================
|
||||
# 策略执行参数
|
||||
|
||||
@@ -1,5 +0,0 @@
|
||||
_target_: roboimi.vla.modules.encoders.LeWMStateEncoder
|
||||
|
||||
input_dim: ${agent.obs_dim}
|
||||
hidden_dim: 256
|
||||
output_dim: 64
|
||||
@@ -4,6 +4,7 @@ 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):
|
||||
@@ -24,8 +25,6 @@ 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,
|
||||
):
|
||||
"""
|
||||
@@ -36,6 +35,7 @@ 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,13 +45,6 @@ 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
|
||||
@@ -72,19 +65,13 @@ 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:
|
||||
stem = hdf5_path.stem
|
||||
if stem.startswith("episode_"):
|
||||
try:
|
||||
idx = int(stem.split("_")[-1])
|
||||
except ValueError:
|
||||
continue
|
||||
if idx in requested:
|
||||
filtered.append(hdf5_path)
|
||||
match = re.search(r'episode_(\d+)$', hdf5_path.stem)
|
||||
if match and int(match.group(1)) in requested:
|
||||
filtered.append(hdf5_path)
|
||||
self.hdf5_files = filtered
|
||||
if not self.hdf5_files:
|
||||
raise FileNotFoundError(
|
||||
@@ -98,12 +85,9 @@ class SimpleRobotDataset(Dataset):
|
||||
with h5py.File(hdf5_path, 'r') as f:
|
||||
T = f['action'].shape[0]
|
||||
dataset_episode_idx = ep_idx
|
||||
stem = hdf5_path.stem
|
||||
if stem.startswith("episode_"):
|
||||
try:
|
||||
dataset_episode_idx = int(stem.split("_")[-1])
|
||||
except ValueError:
|
||||
pass
|
||||
match = re.search(r'episode_(\d+)$', hdf5_path.stem)
|
||||
if match:
|
||||
dataset_episode_idx = int(match.group(1))
|
||||
start_idx = len(self.frame_meta)
|
||||
for t in range(T):
|
||||
self.frame_meta.append({
|
||||
@@ -258,60 +242,6 @@ 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
|
||||
|
||||
@@ -1,14 +1,5 @@
|
||||
# Backbone models
|
||||
__all__ = [
|
||||
"LEWMViTBackbone",
|
||||
"LeWMMultiViewResNetBackbone",
|
||||
"QueryTokenPredictor",
|
||||
"LeWMProjectorMLP",
|
||||
"SIGReg",
|
||||
"ResNetBackbone",
|
||||
"ResNetDiffusionBackbone",
|
||||
"SigLIP2DiffusionBackbone",
|
||||
]
|
||||
__all__ = ["LEWMViTBackbone", "ResNetBackbone", "ResNetDiffusionBackbone", "SigLIP2DiffusionBackbone"]
|
||||
|
||||
|
||||
def __getattr__(name):
|
||||
@@ -18,19 +9,6 @@ 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
|
||||
|
||||
@@ -1,409 +0,0 @@
|
||||
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()
|
||||
@@ -0,0 +1,410 @@
|
||||
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
|
||||
@@ -0,0 +1,13 @@
|
||||
"""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",
|
||||
]
|
||||
@@ -0,0 +1,75 @@
|
||||
"""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
|
||||
@@ -0,0 +1,421 @@
|
||||
"""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)
|
||||
@@ -0,0 +1,571 @@
|
||||
# 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
|
||||
@@ -16,23 +16,3 @@ 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)
|
||||
|
||||
@@ -0,0 +1,510 @@
|
||||
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()
|
||||
@@ -129,6 +129,55 @@ 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)
|
||||
|
||||
@@ -168,6 +217,58 @@ 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(
|
||||
@@ -186,6 +287,25 @@ 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(
|
||||
@@ -204,6 +324,7 @@ 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),
|
||||
@@ -231,8 +352,46 @@ 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([20.0], dtype=np.float32))
|
||||
np.testing.assert_array_equal(second_action.numpy(), np.array([30.0], dtype=np.float32))
|
||||
np.testing.assert_array_equal(first_action.numpy(), np.array([10.0], dtype=np.float32))
|
||||
np.testing.assert_array_equal(second_action.numpy(), np.array([20.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 = [
|
||||
|
||||
+204
-30
@@ -15,6 +15,7 @@ class _FakeAgent:
|
||||
def __init__(self):
|
||||
self.reset_calls = 0
|
||||
self.last_observation = None
|
||||
self.observation_shapes = []
|
||||
|
||||
def eval(self):
|
||||
return self
|
||||
@@ -27,6 +28,7 @@ class _FakeAgent:
|
||||
|
||||
def select_action(self, observation):
|
||||
self.last_observation = observation
|
||||
self.observation_shapes.append(tuple(observation["images"]["front"].shape))
|
||||
return torch.zeros(16)
|
||||
|
||||
|
||||
@@ -36,8 +38,8 @@ class _FakeEnv:
|
||||
self.render_calls = 0
|
||||
self.reset_calls = []
|
||||
|
||||
def reset(self, box_pos):
|
||||
self.reset_calls.append(np.array(box_pos))
|
||||
def reset(self, task_state):
|
||||
self.reset_calls.append(task_state)
|
||||
|
||||
def _get_image_obs(self):
|
||||
self.image_obs_calls += 1
|
||||
@@ -74,7 +76,7 @@ class _FakeRenderer:
|
||||
self._env = env
|
||||
self._frames = [
|
||||
np.full((4, 4, 3), fill_value=index, dtype=np.uint8)
|
||||
for index in range(5)
|
||||
for index in range(8)
|
||||
]
|
||||
self._index = 0
|
||||
|
||||
@@ -108,6 +110,23 @@ 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):
|
||||
@@ -125,26 +144,8 @@ class EvalVLAHeadlessTest(unittest.TestCase):
|
||||
|
||||
self.assertIn("headless", eval_cfg)
|
||||
self.assertFalse(eval_cfg.headless)
|
||||
|
||||
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)
|
||||
self.assertIn("task_description", eval_cfg)
|
||||
self.assertIsNone(eval_cfg.task_description)
|
||||
|
||||
def test_make_sim_env_accepts_headless_and_disables_render(self):
|
||||
fake_env = object()
|
||||
@@ -164,7 +165,7 @@ class EvalVLAHeadlessTest(unittest.TestCase):
|
||||
is_render=False,
|
||||
control_freq=30,
|
||||
is_interpolate=True,
|
||||
cam_view="angle",
|
||||
cam_view="top",
|
||||
)
|
||||
|
||||
def test_headless_sync_camera_capture_populates_images_without_gui_calls(self):
|
||||
@@ -173,11 +174,10 @@ class EvalVLAHeadlessTest(unittest.TestCase):
|
||||
env.mj_data = object()
|
||||
env.exit_flag = False
|
||||
env.is_render = False
|
||||
env.cam = 'angle'
|
||||
env.cam = 'top'
|
||||
env.r_vis = None
|
||||
env.l_vis = None
|
||||
env.top = None
|
||||
env.angle = None
|
||||
env.front = None
|
||||
env._offscreen_renderer = None
|
||||
|
||||
@@ -196,7 +196,6 @@ 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):
|
||||
@@ -216,11 +215,10 @@ class EvalVLAHeadlessTest(unittest.TestCase):
|
||||
env.mj_data = object()
|
||||
env.exit_flag = False
|
||||
env.is_render = False
|
||||
env.cam = "angle"
|
||||
env.cam = "top"
|
||||
env.r_vis = None
|
||||
env.l_vis = None
|
||||
env.top = None
|
||||
env.angle = None
|
||||
env.front = None
|
||||
|
||||
with mock.patch(
|
||||
@@ -237,9 +235,33 @@ 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()
|
||||
@@ -290,6 +312,74 @@ 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],
|
||||
@@ -347,6 +437,27 @@ 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(
|
||||
{
|
||||
@@ -513,6 +624,69 @@ 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()
|
||||
|
||||
@@ -102,10 +102,8 @@ 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)
|
||||
|
||||
@@ -135,8 +133,6 @@ 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,
|
||||
@@ -180,14 +176,12 @@ 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]))
|
||||
@@ -224,120 +218,267 @@ 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_exports_front_trajectory_images_without_video_dependency(self):
|
||||
actions = [
|
||||
np.arange(16, dtype=np.float32),
|
||||
np.arange(16, dtype=np.float32) + 10.0,
|
||||
np.arange(16, dtype=np.float32) + 100.0,
|
||||
np.arange(16, dtype=np.float32) + 110.0,
|
||||
def test_run_eval_parallel_rejects_trajectory_and_video_exports(self):
|
||||
unsupported_flags = [
|
||||
"record_video",
|
||||
"save_trajectory",
|
||||
"save_trajectory_npz",
|
||||
]
|
||||
fake_agent = _FakeAgent(actions)
|
||||
fake_env = _FakeEnv()
|
||||
|
||||
for flag_name in unsupported_flags:
|
||||
with self.subTest(flag_name=flag_name):
|
||||
cfg = OmegaConf.create(
|
||||
{
|
||||
"agent": {},
|
||||
"eval": {
|
||||
"ckpt_path": "checkpoints/vla_model_best.pt",
|
||||
"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,
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
with self.assertRaisesRegex(ValueError, flag_name):
|
||||
eval_vla._run_eval_parallel(cfg)
|
||||
|
||||
def test_run_eval_parallel_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': 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,
|
||||
"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,
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
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(
|
||||
def fake_run_spawn_jobs(payloads, max_workers, worker_fn):
|
||||
del max_workers, worker_fn
|
||||
return [
|
||||
{
|
||||
'output_path': output_path,
|
||||
'camera_name': camera_name,
|
||||
'raw_actions': [np.array(action, copy=True) for action in raw_actions],
|
||||
}
|
||||
)
|
||||
if output_path is None:
|
||||
return None
|
||||
output_path = Path(output_path)
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
output_path.write_bytes(b'fake-png')
|
||||
return str(output_path)
|
||||
"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,
|
||||
'load_checkpoint',
|
||||
return_value=(fake_agent, None),
|
||||
"sample_transfer_pose",
|
||||
side_effect=[
|
||||
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,
|
||||
'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], 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)
|
||||
"_run_spawn_jobs",
|
||||
side_effect=fake_run_spawn_jobs,
|
||||
):
|
||||
summary = eval_vla._run_eval_parallel(cfg)
|
||||
|
||||
self.assertEqual(save_trajectory_image_mock.call_count, 2)
|
||||
open_video_writer_mock.assert_not_called()
|
||||
self.assertIsNone(summary['artifacts']['video_mp4'])
|
||||
self.assertEqual(summary['artifacts']['trajectory_image_camera_name'], 'front')
|
||||
self.assertEqual(
|
||||
[call['camera_name'] for call in trajectory_image_calls],
|
||||
['front', 'front'],
|
||||
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)
|
||||
|
||||
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,
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
first_episode_path = Path(summary['episodes'][0]['artifact_paths']['trajectory_image'])
|
||||
second_episode_path = Path(summary['episodes'][1]['artifact_paths']['trajectory_image'])
|
||||
self.assertTrue(first_episode_path.exists())
|
||||
self.assertTrue(second_episode_path.exists())
|
||||
self.assertNotEqual(first_episode_path, second_episode_path)
|
||||
self.assertEqual(first_episode_path.parent, Path(tmpdir))
|
||||
self.assertEqual(second_episode_path.parent, Path(tmpdir))
|
||||
def fake_run_cuda_parallel_processes(server_payloads, worker_payloads):
|
||||
self.assertEqual(len(server_payloads), 1)
|
||||
self.assertEqual(server_payloads[0]["device_index"], 0)
|
||||
self.assertEqual([payload["server_index"] for payload in worker_payloads], [0, 0])
|
||||
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],
|
||||
},
|
||||
},
|
||||
]
|
||||
|
||||
np.testing.assert_array_equal(trajectory_image_calls[0]['raw_actions'][0], actions[0])
|
||||
np.testing.assert_array_equal(trajectory_image_calls[0]['raw_actions'][1], actions[1])
|
||||
np.testing.assert_array_equal(trajectory_image_calls[1]['raw_actions'][0], actions[2])
|
||||
np.testing.assert_array_equal(trajectory_image_calls[1]['raw_actions'][1], actions[3])
|
||||
with mock.patch.object(
|
||||
eval_vla,
|
||||
"sample_transfer_pose",
|
||||
side_effect=[
|
||||
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__':
|
||||
|
||||
@@ -376,49 +376,6 @@ 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)
|
||||
@@ -544,311 +501,6 @@ 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 = {
|
||||
@@ -1199,80 +851,6 @@ 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(
|
||||
|
||||
@@ -4,7 +4,7 @@ import unittest
|
||||
from pathlib import Path
|
||||
from unittest import mock
|
||||
|
||||
from roboimi.assets.robots.diana_med import BiDianaMed
|
||||
from roboimi.assets.robots import diana_med
|
||||
|
||||
|
||||
class _FakeKDL:
|
||||
@@ -24,6 +24,7 @@ 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'
|
||||
@@ -58,6 +59,47 @@ 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()
|
||||
|
||||
@@ -0,0 +1,42 @@
|
||||
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()
|
||||
@@ -12,21 +12,18 @@ from roboimi.vla.data.simpe_robot_dataset import SimpleRobotDataset
|
||||
|
||||
|
||||
class SimpleRobotDatasetImageLoadingTest(unittest.TestCase):
|
||||
def _write_episode(self, dataset_dir: Path, episode_idx: int = 0, *, base_value: float = 0.0) -> None:
|
||||
episode_path = dataset_dir / f"episode_{episode_idx}.hdf5"
|
||||
def _write_episode(self, dataset_dir: Path) -> None:
|
||||
episode_path = dataset_dir / "episode_0.hdf5"
|
||||
with h5py.File(episode_path, "w") as root:
|
||||
root.create_dataset(
|
||||
"action",
|
||||
data=(np.arange(8, dtype=np.float32).reshape(4, 2) + base_value),
|
||||
)
|
||||
root.create_dataset("action", data=np.arange(8, dtype=np.float32).reshape(4, 2))
|
||||
root.create_dataset(
|
||||
"observations/qpos",
|
||||
data=(np.arange(16, dtype=np.float32).reshape(4, 4) + base_value),
|
||||
data=np.arange(16, dtype=np.float32).reshape(4, 4),
|
||||
)
|
||||
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) + int(base_value)) % 255).reshape(4, 8, 8, 3),
|
||||
data=np.arange(4 * 8 * 8 * 3, dtype=np.uint8).reshape(4, 8, 8, 3),
|
||||
)
|
||||
|
||||
def test_getitem_only_resizes_observation_horizon_images(self):
|
||||
@@ -82,46 +79,3 @@ 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])))
|
||||
|
||||
@@ -0,0 +1,311 @@
|
||||
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()
|
||||
@@ -0,0 +1,445 @@
|
||||
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()
|
||||
@@ -0,0 +1,121 @@
|
||||
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()
|
||||
@@ -0,0 +1,328 @@
|
||||
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()
|
||||
@@ -14,6 +14,8 @@ from roboimi.demos.vla_scripts import eval_vla, train_vla
|
||||
|
||||
|
||||
class _FakeDataset:
|
||||
available_episode_indices = [0, 1]
|
||||
|
||||
def __len__(self):
|
||||
return 4
|
||||
|
||||
@@ -29,6 +31,13 @@ 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}]
|
||||
@@ -91,6 +100,16 @@ 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__()
|
||||
@@ -150,6 +169,94 @@ 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'))
|
||||
|
||||
@@ -162,6 +269,40 @@ 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(
|
||||
{
|
||||
@@ -236,13 +377,7 @@ 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)
|
||||
@@ -259,6 +394,23 @@ 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(
|
||||
{
|
||||
@@ -345,6 +497,93 @@ 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(
|
||||
{
|
||||
|
||||
@@ -41,19 +41,6 @@ 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
|
||||
@@ -127,26 +114,6 @@ 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
|
||||
@@ -372,8 +339,6 @@ 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,
|
||||
@@ -423,18 +388,6 @@ 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()
|
||||
@@ -444,15 +397,6 @@ 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)
|
||||
@@ -498,8 +442,6 @@ 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,
|
||||
@@ -545,95 +487,6 @@ 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)
|
||||
@@ -815,52 +668,6 @@ 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)
|
||||
|
||||
@@ -39,6 +39,17 @@ 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 {}
|
||||
@@ -46,13 +57,18 @@ 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):
|
||||
def __init__(self, params, lr, weight_decay, betas=(0.9, 0.999), eps=1e-8):
|
||||
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)
|
||||
|
||||
@@ -79,6 +95,12 @@ 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):
|
||||
@@ -324,7 +346,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: iterable):
|
||||
mock.patch.object(module, 'tqdm', side_effect=lambda iterable, **kwargs: FakeTqdm(iterable, **kwargs)):
|
||||
module.main(cfg)
|
||||
finally:
|
||||
os.chdir(previous_cwd)
|
||||
@@ -404,7 +426,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: iterable):
|
||||
mock.patch.object(module, 'tqdm', side_effect=lambda iterable, **kwargs: FakeTqdm(iterable, **kwargs)):
|
||||
module.main(cfg)
|
||||
finally:
|
||||
os.chdir(previous_cwd)
|
||||
@@ -422,3 +444,159 @@ 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)
|
||||
|
||||
Reference in New Issue
Block a user