JiajunLI/roboimi
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
4cd33258d257236a707ad344f023bbe57cc0ec15
roboimi/docs/superpowers/specs/2026-03-30-vla-training-headless-swanlab-design.md

9.1 KiB
Raw Blame History

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

命令模板：

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

命令模板：

/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 开关

命令模板：

/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 模式下被实际调用
Reference in New Issue View Git Blame Copy Permalink