feat(vla): add SmolVLA conditioning and experiment artifacts

2026-07-31 10:11:04 +08:00
parent acbd7c605a
commit 5ae9f5fa48
175 changed files with 317471 additions and 70 deletions
@@ -0,0 +1,454 @@
+# 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