# 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 模式下被实际调用