JiajunLI/roboimi
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
5ae9f5fa482fc7c5b1c659d7defe5eb7ed190c03
roboimi/docs/superpowers/plans/2026-03-30-vla-training-headless-swanlab.md
T

14 KiB
Raw Blame History

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:

/home/droid/.conda/envs/roboimi/bin/python -m pip install --dry-run \
  diffusers torchvision einops swanlab

Expected: 输出候选版本；若显示会升级/替换 torch，则停止并改用显式兼容版本

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

Run:

/home/droid/.conda/envs/roboimi/bin/python -m pip install \
  diffusers torchvision einops swanlab

Expected: 安装成功，且不替换当前 torch==2.4.0

  • Step 4: 运行 import 验证

Run:

/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:

/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:

/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:

/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:

/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:

/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:

/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:

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:

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:

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
Reference in New Issue View Git Blame Copy Permalink