JiajunLI/roboimi

Fork 0

Files

T

Logic 5ae9f5fa48 feat(vla): add SmolVLA conditioning and experiment artifacts

2026-07-31 10:11:04 +08:00

17 KiB

Raw Blame History

实验操作指南

0. 适用范围与当前起点

这份文档面向 当前 main 分支 的通用训练 / 验证 / 评估流程，不绑定某一条历史实验分支。

当前默认 Hydra 起点：

agent：resnet_transformer
data：simpe_robot_dataset
eval：eval

当前仓库里常见可选 agent 配置包括：

resnet_transformer
resnet_diffusion
resnet_gr00t_dit
resnet_imf_attnres
resnet_imf_attnres_multitoken
siglip2_imf_attnres
lewm_imf_attnres

说明：

experiment_suites/ 下的历史目录仍然可以作为参考，但正文以当前 main 的通用流程为准。
如果你只是想快速开始，优先从默认 resnet_transformer 配置出发，再按实验需求替换 agent。

1. 关键文件与入口

路径	作用
`roboimi/demos/vla_scripts/train_vla.py`	主训练入口；负责数据集、checkpoint、`val/loss`、held-out action MSE、训练期 rollout 验证、SwanLab
`roboimi/demos/vla_scripts/eval_vla.py`	单次 rollout / 离线评估入口；支持 headless、summary、trajectory image、video artifact、多进程并行 rollout
`roboimi/vla/conf/config.yaml`	全局 Hydra 训练配置
`roboimi/vla/conf/data/simpe_robot_dataset.yaml`	默认数据集配置；数据路径、相机名、图像 resize 都在这里
`roboimi/vla/conf/eval/eval.yaml`	eval 默认配置；`eval.ckpt_path`、`eval.num_episodes`、artifact 开关、多进程 rollout 配置都在这里
`roboimi/vla/conf/agent/*.yaml`	可选 agent 配置集合
`roboimi/vla/data/simpe_robot_dataset.py`	HDF5 懒加载数据集；支持 `episode_indices` 过滤与 `available_episode_indices` 元信息
`roboimi/vla/scripts/calculate_stats.py`	重算 `dataset_stats.pkl`
`experiment_suites/`	历史实验目录；可参考，但不作为当前主流程的唯一事实来源

2. 三台机器与环境

机器	GPU	Python	常用数据集路径	备注
本地 `droid-z790eagleax`	1× RTX 5090 32GB	`/home/droid/.conda/envs/roboimi/bin/python`	`/home/droid/project/diana_sim/sim_transfer`	适合 smoke、单条主跑、本地调参
5880 节点 `100.73.14.65`	2× RTX 5880 Ada 48GB	`/home/droid/miniforge3/envs/roboimi/bin/python`	`/home/droid/sim_dataset/sim_transfer`	适合 2 条并行主跑
L20 节点 `100.119.99.14`	8× NVIDIA L20 46GB	`/home/droid/miniforge3/envs/roboimi/bin/python`	`/data/simtransfer/current`	适合大规模 grid；建议数据与 run 放 `/data`

连接：

5880：ssh droid@100.73.14.65
L20：ssh droid@100.119.99.14

说明：

repo / worktree 路径可能会随当前 checkout 或同步目录变化，以你当前机器上的实际 repo 路径为准。
如果训练目录会被 Hydra 自动切走，建议显式设置 hydra.run.dir=... 以固定输出目录。

3. 训练流怎么走

train_vla.py 当前主流程如下：

读取 Hydra 配置并打印完整 cfg
根据配置确定数据集图像 resize：
- 默认用 data.image_resize_shape
- 如果 agent.vision_backbone.dataset_image_resize_shape 存在，则优先用 backbone 覆盖值
通过 build_train_val_datasets() 构建 dataset / train dataset / val dataset
- 若设置 train.val_episode_indices：按显式 episode 切出 held-out val
- 否则按 train.val_split 做随机划分
用 DataLoader 建 train / val loader
从 dataset_dir/dataset_stats.pkl 读取归一化统计
实例化当前指定的 agent：instantiate(cfg.agent, dataset_stats=...)
可选加载：
- train.pretrained_ckpt：微调起点
- train.resume_ckpt：断点续训（支持显式路径或 auto）
建 optimizer 与 scheduler
训练循环里按 log_freq 记录 train loss / lr
按 save_freq 保存 checkpoints/vla_model_step_*.pt，并在有 val loader 时计算 val/loss
若设置了显式 held-out val，并且 train.action_mse_val_freq_epochs > 0，则按 epoch 计算 held-out action MSE
按 train.rollout_val_freq_epochs 跑训练期 rollout 验证
最后写：
- 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 runtime output dir 下
checkpoint 默认在：<hydra_output_dir>/checkpoints/
如果不显式指定 hydra.run.dir，Hydra 会自动生成目录

4. 验证流怎么走

4.1 常规 `val/loss`

常规验证有两种来源：

随机划分验证集
- 用 train.val_split > 0
- 训练脚本会在保存 checkpoint 时计算 val/loss
显式 held-out episode 验证集
- 用 train.val_episode_indices=[...]
- 训练集 = 全部 episode - held-out episode
- 验证集 = held-out episode
- checkpoint 保存时同样会计算 val/loss

说明：

若 train.val_split=0.0 且未设置 train.val_episode_indices，则不会有 val/loss。

4.2 held-out action MSE

如果想对固定 held-out episode 做更稳定、可复现的动作预测误差验证，推荐使用：

train.val_split=0.0
train.val_episode_indices=[100]（或你指定的 episode 列表）
train.action_mse_val_freq_epochs=1

当前 main 中这套逻辑会：

每隔 action_mse_val_freq_epochs 个 epoch
对 held-out val loader 调 agent.predict_action_chunk(...)
与 batch 中的 action 计算 masked MSE
如果存在 action_is_pad，会自动 mask 掉 padding 部分

日志 key：

控制台 / train_vla.log：held-out action MSE
SwanLab：val/action_mse

重要约束：

train.action_mse_val_freq_epochs > 0 必须搭配 train.val_episode_indices
如果只设了 action_mse_val_freq_epochs 而没有设 val_episode_indices，训练会直接报错

4.3 rollout 验证

训练内 rollout 验证由：

train_vla.py -> run_rollout_validation() -> eval_vla._run_eval()

当前训练内 rollout 会强制：

headless=true
verbose_action=false
record_video=false
save_trajectory_image=true
trajectory_image_camera_name=front
save_summary_json=true

当前这套路径是配置驱动的：

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_output_dir>/rollout_artifacts/<checkpoint_stem>/

常见文件：

rollout_summary.json
rollout_front_ep01_trajectory.png ...

日志重点看：

Epoch X rollout 平均奖励
最佳模型已更新

5. 数据集加载与 `val_episode_indices` 机制

5.1 数据集格式

SimpleRobotDataset 读取 dataset_dir 下的 *.hdf5 / 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 文件
在 worker 内做 HDF5 文件句柄 LRU 缓存
根据文件名中的 episode_XXX 建立 available_episode_indices

5.3 `val_episode_indices` 怎么切

build_train_val_datasets() 的逻辑是：

先 instantiate 一次完整 dataset
读取 dataset.available_episode_indices
检查 train.val_episode_indices 是否都存在
用 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 与返回字段

dataset 侧 resize 默认来自：

data.image_resize_shape
如果 backbone 额外覆盖，则优先 agent.vision_backbone.dataset_image_resize_shape

当前通用 batch 返回字段包括：

observation.state
observation_is_pad
observation.<cam>
action
action_is_pad
task

5.5 统计文件

训练和推理都默认依赖 dataset_stats.pkl。数据集更新后需要重算：

/home/droid/.conda/envs/roboimi/bin/python roboimi/vla/scripts/calculate_stats.py \
  --dataset_dir /home/droid/project/diana_sim/sim_transfer

远端只要把：

Python 路径
--dataset_dir

替换成对应机器上的实际路径即可。

6. SwanLab 行为

当前默认配置里：

train.use_swanlab=false

如果要开启，通常显式设置：

train.use_swanlab=true
train.swanlab_project=roboimi-vla
train.swanlab_run_name=<run_name>

train_vla.py 当前会记录：

初始化时上传 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 能力，入口就是：

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 两种常见模式

单机单卡，多 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
单机多卡，多 server 分摊 worker
- 典型：5880 / 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
多 worker 时不支持同时导出：
- eval.record_video=true
- eval.save_trajectory=true
- eval.save_trajectory_npz=true
eval.save_trajectory_image=true 可以开，适合并行 reward + 定性检查一起做

8. 常用命令模板

下面的模板都以 当前 main 通用流程 为准。远端机器只需要替换：

Python 路径
repo 路径
data.dataset_dir
hydra.run.dir

8.1 本地 smoke train

/home/droid/.conda/envs/roboimi/bin/python roboimi/demos/vla_scripts/train_vla.py \
  agent=resnet_transformer \
  data.dataset_dir=/home/droid/project/diana_sim/sim_transfer \
  train.device=cuda \
  train.batch_size=8 \
  train.max_steps=1000 \
  train.num_workers=4 \
  train.save_freq=200 \
  hydra.run.dir=/home/droid/project/roboimi/runs/smoke-resnet-transformer

8.2 通用训练模板

/home/droid/.conda/envs/roboimi/bin/python roboimi/demos/vla_scripts/train_vla.py \
  agent=resnet_transformer \
  data.dataset_dir=/home/droid/project/diana_sim/sim_transfer \
  train.device=cuda \
  train.batch_size=16 \
  train.lr=0.0001 \
  train.max_steps=100000 \
  train.num_workers=4 \
  train.save_freq=2000 \
  train.use_swanlab=true \
  train.swanlab_project=roboimi-vla \
  train.swanlab_run_name=<run_name> \
  hydra.run.dir=/home/droid/project/roboimi/runs/<run_name>

8.3 带 held-out episode 验证的训练模板

/home/droid/.conda/envs/roboimi/bin/python roboimi/demos/vla_scripts/train_vla.py \
  agent=resnet_transformer \
  data.dataset_dir=/home/droid/project/diana_sim/sim_transfer \
  train.device=cuda \
  train.batch_size=16 \
  train.lr=0.0001 \
  train.max_steps=100000 \
  train.num_workers=4 \
  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 \
  train.use_swanlab=true \
  train.swanlab_project=roboimi-vla \
  train.swanlab_run_name=<run_name> \
  hydra.run.dir=/home/droid/project/roboimi/runs/<run_name>

8.4 断点续训模板

train_vla.py 支持通过 CLI override 使用 train.resume_ckpt：

/home/droid/.conda/envs/roboimi/bin/python roboimi/demos/vla_scripts/train_vla.py \
  agent=resnet_transformer \
  data.dataset_dir=/home/droid/project/diana_sim/sim_transfer \
  train.device=cuda \
  train.resume_ckpt=auto \
  hydra.run.dir=/home/droid/project/roboimi/runs/<run_name>

也可以把 train.resume_ckpt 换成显式 checkpoint 路径。

8.5 单次离线评估

/home/droid/.conda/envs/roboimi/bin/python roboimi/demos/vla_scripts/eval_vla.py \
  agent=resnet_transformer \
  data.dataset_dir=/home/droid/project/diana_sim/sim_transfer \
  train.device=cuda eval.device=cuda \
  eval.ckpt_path=/home/droid/project/roboimi/runs/<run_name>/checkpoints/vla_model_best.pt \
  eval.num_episodes=10 \
  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/roboimi_eval_front

8.6 并行离线评估

/home/droid/.conda/envs/roboimi/bin/python roboimi/demos/vla_scripts/eval_vla.py \
  agent=resnet_transformer \
  data.dataset_dir=/home/droid/project/diana_sim/sim_transfer \
  train.device=cuda eval.device=cuda \
  eval.ckpt_path=/home/droid/project/roboimi/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/roboimi_parallel_eval

8.7 训练内启用并行 GPU rollout

/home/droid/.conda/envs/roboimi/bin/python roboimi/demos/vla_scripts/train_vla.py \
  agent=resnet_transformer \
  data.dataset_dir=/home/droid/project/diana_sim/sim_transfer \
  train.device=cuda \
  train.batch_size=16 \
  train.lr=0.0001 \
  train.max_steps=100000 \
  train.num_workers=4 \
  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.use_swanlab=true \
  train.swanlab_project=roboimi-vla \
  train.swanlab_run_name=<run_name> \
  hydra.run.dir=/home/droid/project/roboimi/runs/<run_name>

8.8 重算统计文件

/home/droid/.conda/envs/roboimi/bin/python roboimi/vla/scripts/calculate_stats.py \
  --dataset_dir /home/droid/project/diana_sim/sim_transfer

8.9 监控日志

tail -f runs/<run_name>/train_vla.log

如果你显式设置了 hydra.run.dir，优先直接 tail 对应绝对路径下的日志文件。

9. 操作建议

优先固定 hydra.run.dir，避免输出目录分散在 Hydra 自动生成路径里
如果只需要常规验证，使用：
- train.val_split > 0
如果需要稳定、可复现的 held-out 指标，使用：
- train.val_split=0.0
- train.val_episode_indices=[...]
- train.action_mse_val_freq_epochs=1
train.num_workers 是 DataLoader worker，不等于 rollout 并行度
rollout 并行优先看：
- train.rollout_num_workers
- train.rollout_cuda_devices
- eval.num_workers
- eval.cuda_devices
数据集更新后记得重算 dataset_stats.pkl
若同时设置了：
- train.pretrained_ckpt
- train.resume_ckpt 当前逻辑会优先走 resume_ckpt
建议 run name 中带上关键信息：
- agent / horizon / batch size / lr / host / gpu / date
experiment_suites/ 下的历史目录适合作为参考资料，但不要默认把某一条旧 suite 当作当前主流程的唯一标准

17 KiB Raw Blame History Unescape Escape