Logic
499 lines
17 KiB
Markdown
499 lines
17 KiB
Markdown
# 实验操作指南
|
||
|
||
## 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` 当前主流程如下：
|
||
|
||
1. 读取 Hydra 配置并打印完整 cfg
|
||
2. 根据配置确定数据集图像 resize：
|
||
- 默认用 `data.image_resize_shape`
|
||
- 如果 `agent.vision_backbone.dataset_image_resize_shape` 存在，则优先用 backbone 覆盖值
|
||
3. 通过 `build_train_val_datasets()` 构建 dataset / train dataset / val dataset
|
||
- 若设置 `train.val_episode_indices`：按显式 episode 切出 held-out val
|
||
- 否则按 `train.val_split` 做随机划分
|
||
4. 用 `DataLoader` 建 train / val loader
|
||
5. 从 `dataset_dir/dataset_stats.pkl` 读取归一化统计
|
||
6. 实例化当前指定的 agent：`instantiate(cfg.agent, dataset_stats=...)`
|
||
7. 可选加载：
|
||
- `train.pretrained_ckpt`：微调起点
|
||
- `train.resume_ckpt`：断点续训（支持显式路径或 `auto`）
|
||
8. 建 optimizer 与 scheduler
|
||
9. 训练循环里按 `log_freq` 记录 train loss / lr
|
||
10. 按 `save_freq` 保存 `checkpoints/vla_model_step_*.pt`，并在有 val loader 时计算 `val/loss`
|
||
11. 若设置了显式 held-out val，并且 `train.action_mse_val_freq_epochs > 0`，则按 epoch 计算 held-out `action MSE`
|
||
12. 按 `train.rollout_val_freq_epochs` 跑训练期 rollout 验证
|
||
13. 最后写：
|
||
- `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`
|
||
|
||
常规验证有两种来源：
|
||
|
||
1. **随机划分验证集**
|
||
- 用 `train.val_split > 0`
|
||
- 训练脚本会在保存 checkpoint 时计算 `val/loss`
|
||
|
||
2. **显式 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()` 的逻辑是：
|
||
|
||
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 与返回字段
|
||
|
||
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`。数据集更新后需要重算：
|
||
|
||
```bash
|
||
/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 两种常见模式
|
||
|
||
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 / 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
|
||
|
||
```bash
|
||
/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 通用训练模板
|
||
|
||
```bash
|
||
/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 验证的训练模板
|
||
|
||
```bash
|
||
/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`：
|
||
|
||
```bash
|
||
/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 单次离线评估
|
||
|
||
```bash
|
||
/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 并行离线评估
|
||
|
||
```bash
|
||
/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
|
||
|
||
```bash
|
||
/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 重算统计文件
|
||
|
||
```bash
|
||
/home/droid/.conda/envs/roboimi/bin/python roboimi/vla/scripts/calculate_stats.py \
|
||
--dataset_dir /home/droid/project/diana_sim/sim_transfer
|
||
```
|
||
|
||
### 8.9 监控日志
|
||
|
||
```bash
|
||
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 当作当前主流程的唯一标准
|