# VLA Training + Headless Rollout + SwanLab Implementation Plan

> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking.

**Goal:** 补齐当前 `Transformer1D` 训练依赖，在 `/home/droid/project/diana_sim/sim_transfer` 上启动训练，接入 SwanLab 标量日志，并提供训练期可选的 headless rollout validation 路径。

**Architecture:** 保持现有 `train_vla.py` / `eval_vla.py` 主体不变，只做最小必要改造：补依赖、补 stats 生成入口、在训练脚本里加轻量 SwanLab logger 和可选 checkpoint-time rollout wrapper、在环境侧把图像更新与 GUI 显示解耦。训练默认仍走当前 `resnet_transformer + Transformer1D` 路线，rollout validation 作为薄封装默认关闭。

**Tech Stack:** Python, mamba/conda, pip, PyTorch, Hydra, diffusers, torchvision, einops, SwanLab, MuJoCo, OpenCV, unittest

---

### Task 0: 执行前提与分支/环境确认

**Files:**
- Verify only

- [ ] **Step 1: 确认当前分支仍是目标分支**

Run: `git branch --show-current`
Expected: `feat-align-dp-transformer-ee`

- [ ] **Step 2: 记录当前 Python 解释器与环境名**

Run: `/home/droid/.conda/envs/roboimi/bin/python - <<'PY'
import sys
print(sys.executable)
PY`
Expected: 输出 `/home/droid/.conda/envs/roboimi/bin/python`

- [ ] **Step 3: 记录当前数据集目录存在性与 episode 数量**

Run: `/usr/bin/zsh -lc 'echo DATASET=/home/droid/project/diana_sim/sim_transfer; find /home/droid/project/diana_sim/sim_transfer -maxdepth 1 -name "episode_*.hdf5" | wc -l'`
Expected: 输出目录路径与 `100`

### Task 1: 补齐训练依赖并把 resolved versions 写回环境定义

**Files:**
- Modify: `environment.yml`
- Verify only: local `roboimi` env

- [ ] **Step 1: 写出缺失依赖的最小清单**

需要补齐：
- `diffusers`
- `torchvision`
- `einops`
- `swanlab`

- [ ] **Step 2: 先用 dry-run 解析候选版本，确认不会升级 Torch**

Run:
```bash
/home/droid/.conda/envs/roboimi/bin/python -m pip install --dry-run \
  diffusers torchvision einops swanlab
```
Expected: 输出候选版本；若显示会升级/替换 `torch`，则停止并改用显式兼容版本

- [ ] **Step 3: 安装与当前 Torch 兼容的缺失依赖到现有环境**

Run:
```bash
/home/droid/.conda/envs/roboimi/bin/python -m pip install \
  diffusers torchvision einops swanlab
```
Expected: 安装成功，且不替换当前 `torch==2.4.0`

- [ ] **Step 4: 运行 import 验证**

Run:
```bash
/home/droid/.conda/envs/roboimi/bin/python - <<'PY'
mods=['torch','hydra','omegaconf','diffusers','torchvision','einops','cv2','h5py','swanlab','mujoco']
for m in mods:
    __import__(m)
    print('OK', m)
PY
```
Expected: 每个模块都输出 `OK <module>`

- [ ] **Step 5: 记录实际安装版本**

Run:
```bash
/home/droid/.conda/envs/roboimi/bin/python - <<'PY'
import diffusers, torchvision, einops, swanlab
print('diffusers', getattr(diffusers,'__version__',''))
print('torchvision', getattr(torchvision,'__version__',''))
print('einops', getattr(einops,'__version__',''))
print('swanlab', getattr(swanlab,'__version__',''))
PY
```
Expected: 输出四个包的 resolved versions

- [ ] **Step 6: 将 resolved versions 写回 `environment.yml`**

把新增依赖补到 `environment.yml` 的现有依赖列表（若使用 `pip:` 段则更新该段）里，使用 Step 5 得到的**实际 resolved versions**，避免环境漂移，并避免重复 package 条目。

- [ ] **Step 7: 语法检查环境定义文件仅作结构确认**

Run: `python - <<'PY'
from pathlib import Path
text = Path('environment.yml').read_text()
assert 'diffusers' in text
assert 'torchvision' in text
assert 'einops' in text
assert 'swanlab' in text
print('environment.yml updated')
PY`
Expected: `environment.yml updated`

### Task 2: 让统计脚本支持外部数据目录并生成 dataset stats

**Files:**
- Modify: `roboimi/vla/scripts/calculate_stats.py`
- Test: `tests/test_calculate_stats_cli.py`

- [ ] **Step 1: 写 failing test，验证统计脚本可接受外部 `--dataset_dir` 并输出目标路径**

Test file should:
- 用临时目录创建最小 HDF5 episode
- 调用脚本入口/函数时传入外部目录
- 断言输出 `dataset_stats.pkl` 出现在该目录
- 断言 pickle 内包含 `action_mean/qpos_mean/...`

- [ ] **Step 2: 跑测试确认它先失败**

Run: `/home/droid/.conda/envs/roboimi/bin/python -m unittest tests.test_calculate_stats_cli -v`
Expected: FAIL（当前脚本写死默认目录）

- [ ] **Step 3: 最小实现 `--dataset_dir` 支持**

要求：
- 保留现有统计逻辑
- 仅增加 CLI 参数解析
- 输出仍写入 `<dataset_dir>/dataset_stats.pkl`

- [ ] **Step 4: 重新跑测试确认转绿**

Run: `/home/droid/.conda/envs/roboimi/bin/python -m unittest tests.test_calculate_stats_cli -v`
Expected: PASS

- [ ] **Step 5: 用真实数据集生成 stats**

Run:
```bash
/home/droid/.conda/envs/roboimi/bin/python roboimi/vla/scripts/calculate_stats.py \
  --dataset_dir /home/droid/project/diana_sim/sim_transfer
```
Expected: 生成 `/home/droid/project/diana_sim/sim_transfer/dataset_stats.pkl`

- [ ] **Step 6: 验证 stats 文件结构**

Run:
```bash
/home/droid/.conda/envs/roboimi/bin/python - <<'PY'
import pickle
path='/home/droid/project/diana_sim/sim_transfer/dataset_stats.pkl'
with open(path,'rb') as f:
    stats=pickle.load(f)
for k in ['action_mean','action_std','action_min','action_max','qpos_mean','qpos_std','qpos_min','qpos_max']:
    assert k in stats, k
print('stats_ok')
PY
```
Expected: `stats_ok`

### Task 3: 增加 SwanLab 训练日志集成

**Files:**
- Modify: `roboimi/demos/vla_scripts/train_vla.py`
- Modify: `roboimi/vla/conf/config.yaml`
- Test: `tests/test_train_vla_swanlab_logging.py`

- [ ] **Step 1: 写 failing test，验证训练脚本在 `train.use_swanlab=true` 时会初始化 SwanLab 并记录标量**

Test should:
- stub `swanlab`
- 调用训练脚本中抽取出的非 Hydra helper（如 `_run_training(cfg)`）的最小路径（`max_steps=0` 或很小）
- 断言调用了：
  - `swanlab.init(project='roboimi-vla', ...)`
  - 至少一次 `swanlab.log({...})`
- 断言当 `use_swanlab=false` 时不会 import/初始化 SwanLab
- 断言当 `use_swanlab=true` 且 import `swanlab` 失败时会 fail fast
- 断言当 `use_swanlab=true` 且认证/登录状态不可用时会 fail fast
- 断言会记录最终/最佳 checkpoint 路径到 log/summary

- [ ] **Step 2: 跑测试确认它先失败**

Run: `/home/droid/.conda/envs/roboimi/bin/python -m unittest tests.test_train_vla_swanlab_logging -v`
Expected: FAIL（当前无 SwanLab 集成）

- [ ] **Step 3: 在配置中增加最小 SwanLab 契约**

在 `roboimi/vla/conf/config.yaml` 添加：
- `train.use_swanlab: true`
- `train.swanlab_project: roboimi-vla`
- 可选 `train.swanlab_run_name: null`

- [ ] **Step 4: 从 Hydra 入口提取可测试的训练 helper**

要求：
- 新增类似 `_run_training(cfg)` 的普通函数
- `main()` 只做 Hydra 入口转发
- 测试只调用 helper，不直接调用 Hydra-decorated `main(cfg)`

- [ ] **Step 5: 在训练脚本中实现 SwanLab 初始化的 fail-fast 逻辑**

要求：
- `use_swanlab=true` 时：
  - import 失败 -> 直接报错
  - 本地未登录/认证失败 -> 直接报错
  - 成功后执行 `swanlab.init(project=cfg.train.swanlab_project, ...)`

- [ ] **Step 6: 在训练脚本中实现轻量 scalar logger**

要求：
- 仅 scalar logging，不引入自定义 callback 框架
- 训练时记录 `train/loss`, `train/lr`, `train/best_loss`, `train/step`
- 验证时记录 `val/loss`
- 训练结束时记录：
  - `train/final_checkpoint_path`
  - `train/best_checkpoint_path`
- 若库支持则显式 `finish/close`

- [ ] **Step 7: 重新跑测试确认转绿**

Run: `/home/droid/.conda/envs/roboimi/bin/python -m unittest tests.test_train_vla_swanlab_logging -v`
Expected: PASS

- [ ] **Step 8: 用提供的 API key 完成本地登录**

Run:
```bash
/usr/bin/zsh -lc 'SWANLAB_API_KEY="<user-provided>"; /home/droid/.conda/envs/roboimi/bin/swanlab login -k "$SWANLAB_API_KEY"'
```
Expected: 登录成功并保存本地凭证

### Task 4: 为评估/rollout 路径增加 headless 模式

**Files:**
- Modify: `roboimi/envs/double_base.py`
- Modify: `roboimi/envs/double_pos_ctrl_env.py`
- Modify: `roboimi/vla/conf/eval/eval.yaml`
- Modify: `roboimi/demos/vla_scripts/eval_vla.py`
- Test: `tests/test_eval_vla_headless.py`

- [ ] **Step 1: 写 failing test，验证 headless 路径不触发 GUI 调用**

Test should stub:
- `cv2.namedWindow`
- `cv2.imshow`
- `cv2.waitKey`
- viewer launch/render path

And assert:
- `eval.headless=true` 时不调用这些 GUI 接口
- 仍能获取图像观测并走到 policy action 执行前/后关键路径

- [ ] **Step 2: 跑测试确认它先失败**

Run: `/home/droid/.conda/envs/roboimi/bin/python -m unittest tests.test_eval_vla_headless -v`
Expected: FAIL（当前 env/eval 默认会开 viewer 和 cv2 窗口）

- [ ] **Step 3: 统一配置开关为 `eval.headless`**

在 `roboimi/vla/conf/eval/eval.yaml` 添加：
- `headless: false`

不要再引入第二个同义开关。

- [ ] **Step 4: 在 env 工厂中接入 `headless`**

在 `make_sim_env(...)` 中：
- `eval.headless=true` -> `is_render=False`
- 保留图像观测更新
- 不创建 MuJoCo viewer

- [ ] **Step 5: 将相机更新与 GUI 显示解耦**

在 `double_base.py`：
- 图像更新逻辑保留
- `cv2.namedWindow/imshow/waitKey` 仅在非 headless 下执行

- [ ] **Step 6: eval 脚本中在 headless 下跳过 `env.render()`**

只在 `eval.headless=false` 时调用 `env.render()`。

- [ ] **Step 7: 重新跑测试确认转绿**

Run: `/home/droid/.conda/envs/roboimi/bin/python -m unittest tests.test_eval_vla_headless -v`
Expected: PASS

### Task 5: 给训练脚本加可选 checkpoint-time rollout validation 薄封装

**Files:**
- Modify: `roboimi/demos/vla_scripts/train_vla.py`
- Possibly modify: `roboimi/demos/vla_scripts/eval_vla.py`（仅当需要提取可复用入口）
- Modify: `roboimi/vla/conf/config.yaml`
- Test: `tests/test_train_vla_rollout_validation.py`

- [ ] **Step 1: 写 failing test，验证 checkpoint 保存点可选调用 rollout validation，且会传 `eval.headless=true`**

Test should:
- stub rollout/eval helper
- 开启 `train.rollout_validate_on_checkpoint=true`
- 设置小 `save_freq`
- 断言训练脚本在 checkpoint 时调用验证 helper
- 断言调用参数带 `headless=true`

- [ ] **Step 2: 跑测试确认它先失败**

Run: `/home/droid/.conda/envs/roboimi/bin/python -m unittest tests.test_train_vla_rollout_validation -v`
Expected: FAIL（当前无该 hook）

- [ ] **Step 3: 增加最小配置键**

在 `config.yaml` 中添加：
- `train.rollout_validate_on_checkpoint: false`
- `train.rollout_num_episodes: 1`

- [ ] **Step 4: 提取一个最小 rollout validation helper 接口**

要求：
- helper 输入至少包括 `cfg`, `ckpt_path`, `num_episodes`, `headless`
- 默认训练侧调用时强制 `headless=True`
- 优先复用现有 eval 逻辑，不引入第二套 validator 类

- [ ] **Step 5: 在 checkpoint 保存路径中接入 rollout helper**

要求：
- 不重写第二套验证框架
- 优先复用现有 eval 逻辑/工具
- 默认关闭
- 仅在 checkpoint 时少量调用
- 强制 `eval.headless=true`

- [ ] **Step 6: 重新跑测试确认转绿**

Run: `/home/droid/.conda/envs/roboimi/bin/python -m unittest tests.test_train_vla_rollout_validation -v`
Expected: PASS

### Task 6: 启动训练前的集成 smoke verification

**Files:**
- Verify only

- [ ] **Step 1: 跑所有新增/相关测试**

Run:
```bash
/home/droid/.conda/envs/roboimi/bin/python -m unittest \
  tests.test_calculate_stats_cli \
  tests.test_train_vla_swanlab_logging \
  tests.test_eval_vla_headless \
  tests.test_train_vla_rollout_validation -v
```
Expected: 全部 PASS

- [ ] **Step 2: 对关键修改文件做语法检查**

Run:
```bash
/home/droid/.conda/envs/roboimi/bin/python -m py_compile \
  roboimi/vla/scripts/calculate_stats.py \
  roboimi/demos/vla_scripts/train_vla.py \
  roboimi/demos/vla_scripts/eval_vla.py \
  roboimi/envs/double_pos_ctrl_env.py \
  roboimi/envs/double_base.py
```
Expected: 无语法错误

- [ ] **Step 3: 运行训练 smoke run**

Run:
```bash
SWANLAB_API_KEY='<user-provided>' \
/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
```
Expected:
- 训练启动成功
- 产生 `checkpoints/vla_model_step_10.pt` 或 `vla_model_final.pt`
- 本地日志中无 ImportError

- [ ] **Step 4: 运行一个最小 headless rollout-validation smoke run**

Run:
```bash
SWANLAB_API_KEY='<user-provided>' \
/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=10 \
  train.log_freq=1 \
  train.save_freq=5 \
  train.use_swanlab=true \
  train.swanlab_project=roboimi-vla \
  train.rollout_validate_on_checkpoint=true \
  train.rollout_num_episodes=1 \
  eval.headless=true
```
Expected:
- 到达 checkpoint-time rollout 调用
- 不弹 MuJoCo viewer
- 不执行 `cv2.namedWindow/imshow/waitKey`

- [ ] **Step 5: 验证 checkpoint 文件已写出**

Run: `/usr/bin/zsh -lc 'ls -lah checkpoints | sed -n "1,120p"'`
Expected: 存在 `vla_model_step_10.pt` 或 `vla_model_final.pt`

- [ ] **Step 6: 验证 SwanLab 已收到标量**

验证方式：
- 终端日志中确认 `swanlab.init` / run URL / run id
- 若工具支持，确认 dashboard 中 `roboimi-vla` 项目下出现本次 run

### Task 7: 启动正式训练

**Files:**
- Verify only

- [ ] **Step 1: 用真实配置启动正式训练**

Run:
```bash
SWANLAB_API_KEY='<user-provided>' \
/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
```
Expected:
- 训练持续运行
- checkpoint 周期性写出
- SwanLab 周期性收到 train/val 标量
- 若 checkpoint-time rollout 打开，则不弹 GUI

- [ ] **Step 2: 记录启动信息并向用户汇报**

汇报内容至少包括：
- 使用的数据集路径
- 训练命令/关键 overrides
- checkpoint 输出目录
- SwanLab project 名称
- rollout validation 是否已启用以及是否 headless