From cb79e00546527dee2b35b8e6fa9abb25ab4feca2 Mon Sep 17 00:00:00 2001
From: Logic <Logic@gitea.ideaopen.cn>
Date: Mon, 30 Mar 2026 18:50:12 +0800
Subject: [PATCH] docs: add VLA training headless swanlab design spec

---
 ...30-vla-training-headless-swanlab-design.md | 241 ++++++++++++++++++
 1 file changed, 241 insertions(+)
 create mode 100644 docs/superpowers/specs/2026-03-30-vla-training-headless-swanlab-design.md

diff --git a/docs/superpowers/specs/2026-03-30-vla-training-headless-swanlab-design.md b/docs/superpowers/specs/2026-03-30-vla-training-headless-swanlab-design.md
new file mode 100644
index 0000000..f232125
--- /dev/null
+++ b/docs/superpowers/specs/2026-03-30-vla-training-headless-swanlab-design.md
@@ -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 模式下被实际调用