docs: add VLA training headless swanlab design spec
This commit is contained in:
Logic
@@ -0,0 +1,241 @@
|
||||
# VLA Training + Headless Rollout + SwanLab Design
|
||||
|
||||
**Date:** 2026-03-30
|
||||
**Branch:** feat-align-dp-transformer-ee
|
||||
|
||||
## Goal
|
||||
在当前仓库中补齐默认 `resnet_transformer` / `Transformer1D` 路线的训练依赖，使用数据集 `/home/droid/project/diana_sim/sim_transfer` 启动训练；同时支持训练过程中的 SwanLab 标量日志上传，并为后续 rollout 验证提供 headless 模式，避免弹出 MuJoCo / OpenCV 图形界面。
|
||||
|
||||
## Non-Goals
|
||||
- 不重写整套训练框架
|
||||
- 不引入新的 workspace / callback 框架
|
||||
- 不在本轮做复杂的视频/媒体日志上传
|
||||
- 不修改数据集格式本身
|
||||
|
||||
## Current State
|
||||
- 默认训练配置已切到 `agent=resnet_transformer`，head 为 `Transformer1D`
|
||||
- 当前环境缺少训练所需的若干 Python 依赖：`diffusers`、`torchvision`、`einops`、`swanlab`
|
||||
- 评估环境 `make_sim_env(task_name)` 当前写死 `is_render=True`
|
||||
- 相机线程 `camera_viewer()` 默认会 `cv2.namedWindow/imshow`，即使只想拿图像也会弹窗
|
||||
- 训练脚本当前支持 train/val loss、checkpoint，但没有 SwanLab 集成
|
||||
- 数据集目录 `/home/droid/project/diana_sim/sim_transfer` 下已有 100 个 episode，但还没有 `dataset_stats.pkl`
|
||||
|
||||
## User Requirements
|
||||
1. 在现有 mamba 环境里补齐训练依赖
|
||||
2. 在 `/home/droid/project/diana_sim/sim_transfer` 上开始训练
|
||||
3. 如果训练中需要 rollout 验证，希望支持 headless，不弹 GUI
|
||||
4. 训练指标上传到 SwanLab
|
||||
5. 默认 SwanLab project 名为 `roboimi-vla`
|
||||
|
||||
## Proposed Approach
|
||||
采用“最小必要改造”方案：
|
||||
|
||||
### 1. Dependency Layer
|
||||
在现有 `roboimi` 环境中补齐缺失训练依赖，并优先保持现有环境名与脚本入口不变。
|
||||
|
||||
#### Install Plan
|
||||
- 环境：继续使用现有 mamba 环境 `roboimi`
|
||||
- 安装方式：
|
||||
- 优先使用当前 env 的 `python -m pip install`
|
||||
- 安装包：
|
||||
- `diffusers`
|
||||
- `torchvision`
|
||||
- `einops`
|
||||
- `swanlab`
|
||||
- 版本策略：
|
||||
- 优先选择与当前 `torch==2.4.0` 可兼容的最新可安装版本
|
||||
- 若出现兼容性问题，再回退到与 `torch 2.4` 对齐的稳定版本
|
||||
- 复现策略：
|
||||
- 本轮会把**实际安装成功的 resolved versions** 补写回仓库的环境定义文件，避免后续环境漂移
|
||||
|
||||
训练前验证以下 import：
|
||||
- `torch`
|
||||
- `hydra`
|
||||
- `omegaconf`
|
||||
- `diffusers`
|
||||
- `torchvision`
|
||||
- `einops`
|
||||
- `swanlab`
|
||||
- `cv2`
|
||||
- `h5py`
|
||||
- `mujoco`
|
||||
|
||||
### 2. Dataset Preparation
|
||||
直接复用现有 `SimpleRobotDataset`，仅将 `data.dataset_dir` 指向：
|
||||
- `/home/droid/project/diana_sim/sim_transfer`
|
||||
|
||||
训练前使用现有统计脚本生成：
|
||||
- `/home/droid/project/diana_sim/sim_transfer/dataset_stats.pkl`
|
||||
|
||||
统计文件生成命令目标为：
|
||||
- 从仓库根目录执行
|
||||
- 直接针对 `/home/droid/project/diana_sim/sim_transfer` 输出 stats
|
||||
- 训练脚本不再依赖默认数据目录
|
||||
|
||||
### 3. SwanLab Logging
|
||||
在训练脚本中增加一个轻量 logging 集成层：
|
||||
- 通过配置决定是否启用 SwanLab，默认启用
|
||||
- 默认 project：`roboimi-vla`
|
||||
- API key 不写入仓库，不写入配置文件，只通过本地登录状态或环境变量使用
|
||||
- 当 `train.use_swanlab=true` 时：
|
||||
- 若 `swanlab` 不可 import，训练直接 fail fast
|
||||
- 若未登录或认证失败，训练直接 fail fast
|
||||
- 每个训练日志点上传：
|
||||
- `train/loss`
|
||||
- `train/lr`
|
||||
- `train/best_loss`
|
||||
- `train/step`
|
||||
- 每次验证时上传：
|
||||
- `val/loss`
|
||||
- 训练结束时记录最终 checkpoint 路径与 best checkpoint 路径
|
||||
|
||||
### 4. Headless Rollout Design
|
||||
目标是让 rollout 验证可以“拿到图像观测，但不弹任何窗口”。
|
||||
|
||||
最小改造策略：
|
||||
- 给 `make_sim_env(...)` 增加 `headless` / `is_render` 参数
|
||||
- 给相机线程显示逻辑增加开关：
|
||||
- headless 时继续更新 `r_vis/top/front/...` 图像缓存
|
||||
- 但不执行 `cv2.namedWindow` / `cv2.imshow` / `cv2.waitKey`
|
||||
- 评估脚本中：
|
||||
- headless 时不调用 `env.render()`
|
||||
- 仍然允许 `env._get_image_obs()` 和 policy inference 正常运行
|
||||
|
||||
#### Training-Time Rollout Scope
|
||||
- 本轮**会提供一个可选的 checkpoint-time rollout validation 路径**，默认关闭
|
||||
- 启用后，在训练保存 checkpoint 时可以调用同仓库的 rollout/eval 逻辑做少量 episode 验证
|
||||
- 此路径要求支持**唯一权威开关** `eval.headless=true`，即：
|
||||
- 不弹 MuJoCo viewer
|
||||
- 不执行 `cv2.namedWindow / cv2.imshow / cv2.waitKey`
|
||||
- 仍可读取图像并完成策略推理
|
||||
- 默认情况下不增加频繁 rollout，以避免拖慢训练；只提供能力与配置开关
|
||||
|
||||
如果验证发现相机线程强依赖 GUI，我们的降级策略是：
|
||||
- 训练主流程 + SwanLab 必须先跑通
|
||||
- rollout validation 保持为显式可选能力
|
||||
- 但本轮仍要保证至少存在可调用的 headless 验证执行路径，而不是仅停留在文档层面
|
||||
|
||||
### 5. Training Execution Strategy
|
||||
分两步执行：
|
||||
|
||||
#### Step A: Smoke Run
|
||||
使用较小步数启动一次 smoke training，确认：
|
||||
- 数据集可正常读取
|
||||
- 统计文件可加载
|
||||
- 模型可实例化
|
||||
- 单步前后向正常
|
||||
- checkpoint 正常写出
|
||||
- SwanLab 成功上传标量
|
||||
|
||||
#### Step B: Real Training Run
|
||||
在 smoke run 成功后，再启动正式训练。
|
||||
|
||||
## Execution Commands
|
||||
|
||||
### A. Stats Generation
|
||||
从仓库根目录执行，生成：
|
||||
- `/home/droid/project/diana_sim/sim_transfer/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
|
||||
```
|
||||
|
||||
### B. Smoke Training Command
|
||||
从仓库根目录执行，核心覆盖项包括：
|
||||
- `data.dataset_dir=/home/droid/project/diana_sim/sim_transfer`
|
||||
- 较小 `train.max_steps`
|
||||
- 较高日志频率
|
||||
- 启用 SwanLab
|
||||
- 输出目录使用当前运行目录下的 `checkpoints/`
|
||||
|
||||
命令模板：
|
||||
```bash
|
||||
/home/droid/.conda/envs/roboimi/bin/python roboimi/demos/vla_scripts/train_vla.py \
|
||||
data.dataset_dir=/home/droid/project/diana_sim/sim_transfer \
|
||||
train.max_steps=20 \
|
||||
train.log_freq=1 \
|
||||
train.save_freq=10 \
|
||||
train.use_swanlab=true \
|
||||
train.swanlab_project=roboimi-vla \
|
||||
train.rollout_validate_on_checkpoint=false
|
||||
```
|
||||
|
||||
### C. Real Training Command
|
||||
从仓库根目录执行，核心覆盖项包括：
|
||||
- `data.dataset_dir=/home/droid/project/diana_sim/sim_transfer`
|
||||
- 正式 `train.max_steps`
|
||||
- 默认 project=`roboimi-vla`
|
||||
- 若启用 rollout validation，则传入 `eval.headless=true` 以及训练侧 rollout 开关
|
||||
|
||||
命令模板：
|
||||
```bash
|
||||
/home/droid/.conda/envs/roboimi/bin/python roboimi/demos/vla_scripts/train_vla.py \
|
||||
data.dataset_dir=/home/droid/project/diana_sim/sim_transfer \
|
||||
train.use_swanlab=true \
|
||||
train.swanlab_project=roboimi-vla \
|
||||
train.rollout_validate_on_checkpoint=true \
|
||||
eval.headless=true
|
||||
```
|
||||
|
||||
### D. Output Behavior
|
||||
- checkpoint 输出目录：当前工作目录下的 `checkpoints/`
|
||||
- 关键文件：
|
||||
- `checkpoints/vla_model_step_<N>.pt`
|
||||
- `checkpoints/vla_model_best.pt`
|
||||
- `checkpoints/vla_model_final.pt`
|
||||
|
||||
## File-Level Changes
|
||||
- `environment.yml`
|
||||
- 补写新增训练依赖，保证后续可复现
|
||||
- `roboimi/demos/vla_scripts/train_vla.py`
|
||||
- 增加 SwanLab 集成
|
||||
- 增加更明确的数据集目录覆盖支持
|
||||
- 增加可选 checkpoint-time rollout validation 入口
|
||||
- 保持当前 optimizer 对齐逻辑不变
|
||||
- `roboimi/vla/conf/config.yaml`
|
||||
- 增加/扩展训练日志、SwanLab、rollout 相关配置项
|
||||
- `roboimi/vla/conf/eval/eval.yaml`
|
||||
- 增加 `headless` 等评估控制项
|
||||
- `roboimi/envs/double_pos_ctrl_env.py`
|
||||
- `make_sim_env` 支持 headless / no-render
|
||||
- `roboimi/envs/double_base.py`
|
||||
- 相机采集与 GUI 显示解耦
|
||||
- `roboimi/vla/scripts/calculate_stats.py`
|
||||
- 改为直接支持通过命令行传入外部 `dataset_dir`
|
||||
- tests（新增）
|
||||
- 覆盖 SwanLab 可选初始化路径
|
||||
- 覆盖 headless 环境下“不弹窗但可取图”的关键逻辑
|
||||
|
||||
## Validation Plan
|
||||
1. 补齐依赖后验证 import 全通过
|
||||
2. 生成 `dataset_stats.pkl`
|
||||
3. 运行训练 smoke run
|
||||
4. 确认 SwanLab dashboard 在 project `roboimi-vla` 下有标量更新
|
||||
5. 若启用 rollout 验证：确认 headless 下不弹 GUI，且 rollout 路径能真正执行
|
||||
6. 再启动正式训练
|
||||
|
||||
## Config Contract
|
||||
本轮新增/固定的配置键以以下形式为准：
|
||||
- `train.use_swanlab: true|false`
|
||||
- `train.swanlab_project: roboimi-vla`
|
||||
- `train.rollout_validate_on_checkpoint: true|false`
|
||||
- `eval.headless: true|false`
|
||||
|
||||
## Risks and Mitigations
|
||||
- **Risk:** GUI/相机线程与离屏渲染耦合
|
||||
- **Mitigation:** 先解耦显示与图像更新；必要时把 rollout 验证降级为第二阶段
|
||||
- **Risk:** 现有 env 依赖不完整
|
||||
- **Mitigation:** 先做 import 验证，再做 smoke run
|
||||
- **Risk:** 数据集过大导致 smoke run 也很慢
|
||||
- **Mitigation:** smoke run 只跑极小步数
|
||||
- **Risk:** SwanLab API key 泄漏
|
||||
- **Mitigation:** 不写入代码/配置，只保存在本地登录态或环境变量
|
||||
|
||||
## Success Criteria
|
||||
- 训练脚本能在 `/home/droid/project/diana_sim/sim_transfer` 上启动
|
||||
- 能成功写出 checkpoint 到 `checkpoints/`
|
||||
- SwanLab 在 `roboimi-vla` 项目下能看到 train/val 标量
|
||||
- headless rollout 具备不弹 GUI 的执行路径
|
||||
- 若训练侧启用 rollout validation，则该路径可以在 headless 模式下被实际调用
|
||||
Reference in New Issue
Block a user