feat(vla): add SmolVLA conditioning and experiment artifacts
This commit is contained in:
Logic
@@ -0,0 +1,256 @@
|
||||
# Align ResNet Transformer Diffusion To External Repo Implementation Plan
|
||||
|
||||
> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development to implement this plan task-by-task. Do not switch to inline execution. Steps use checkbox (`- [ ]`) syntax for tracking.
|
||||
|
||||
**Goal:** 将当前仓库的 `resnet_transformer` 对齐到 external repo `/home/droid/project/diffusion_policy` 的原生 DDPM Transformer Diffusion（非 PMF、非 DiT、非 UNet）实现，并采用 external 中已存在的 full-attention / nocausal 变体（`causal_attn=false`），固定为三相机图像条件输入，同时保持 EE action 语义的推理执行路径正确。
|
||||
|
||||
**Architecture:** 保留当前仓库较轻量的训练脚本和数据集组织，但把 Transformer denoiser 本体对齐到 external repo 的 `TransformerForDiffusion` 实现与接口语义；视觉编码器保留当前仓库的 ResNet+SpatialSoftmax 路线，仅保证其输出维度与三相机条件输入兼容。评估路径统一按 EE action 调 `env.step(action)`，训练/推理配置显式固定三相机 `r_vis/top/front` 且图像永远作为条件输入。
|
||||
|
||||
**Tech Stack:** Python, PyTorch, diffusers DDPM/DDIM, Hydra/OmegaConf, unittest, h5py, OpenCV
|
||||
|
||||
---
|
||||
|
||||
### Task 0: 执行前提与分支约束
|
||||
|
||||
**Files:**
|
||||
- Verify only
|
||||
|
||||
- [ ] **Step 1: Confirm execution stays on the feature branch**
|
||||
|
||||
Run: `git branch --show-current`
|
||||
Expected: `feat-align-dp-transformer-ee`
|
||||
|
||||
If the current branch is not `feat-align-dp-transformer-ee`, create/switch to it before continuing:
|
||||
|
||||
Run: `git checkout -b feat-align-dp-transformer-ee || git checkout feat-align-dp-transformer-ee`
|
||||
Expected: working branch becomes `feat-align-dp-transformer-ee`
|
||||
|
||||
- [ ] **Step 2: Confirm implementation is executed with subagents**
|
||||
|
||||
Use `superpowers:subagent-driven-development` for implementation tasks and reviews. Do not switch to inline/manual execution unless the human explicitly changes course.
|
||||
|
||||
- [ ] **Step 3: Confirm the current branch is the only target branch for this work**
|
||||
|
||||
Do not create or switch to another branch/worktree during implementation unless the human explicitly asks for it. All changes for this migration stay on `feat-align-dp-transformer-ee`.
|
||||
|
||||
### Task 1: 验证 EE action 推理执行语义已固定
|
||||
|
||||
**Files:**
|
||||
- Verify: `roboimi/vla/eval_utils.py`
|
||||
- Verify: `roboimi/demos/vla_scripts/eval_vla.py`
|
||||
- Verify: `tests/test_eval_vla_execution.py`
|
||||
|
||||
- [ ] **Step 1: Verify the test still passes on the feature branch**
|
||||
|
||||
Run: `mamba run -n roboimi python -m unittest tests.test_eval_vla_execution -v`
|
||||
Expected: PASS
|
||||
|
||||
- [ ] **Step 2: Verify the real eval script no longer executes joint-action stepping**
|
||||
|
||||
Run:
|
||||
```bash
|
||||
python - <<'PY'
|
||||
from pathlib import Path
|
||||
src = Path('roboimi/demos/vla_scripts/eval_vla.py').read_text()
|
||||
assert 'execute_policy_action(env, action)' in src
|
||||
assert 'env.step_jnt(action)' not in src
|
||||
print('eval_vla execution path verified')
|
||||
PY
|
||||
```
|
||||
Expected: PASS
|
||||
|
||||
### Task 2: 通过对拍测试把当前 Transformer head 对齐到 external repo 的 `TransformerForDiffusion`
|
||||
|
||||
**Files:**
|
||||
- Modify: `roboimi/vla/models/heads/transformer1d.py`
|
||||
- Modify: `roboimi/vla/conf/head/transformer1d.yaml`
|
||||
- Test: `tests/test_transformer1d_external_alignment.py`
|
||||
|
||||
- [ ] **Step 1: Confirm the alignment test exists and captures the intended parity contract**
|
||||
|
||||
Write a test that imports external repo's `TransformerForDiffusion` from:
|
||||
`/home/droid/project/diffusion_policy/diffusion_policy/model/diffusion/transformer_for_diffusion.py`
|
||||
and verifies the local `Transformer1D` can load the external model's `state_dict` and produce numerically identical outputs for the same inputs when configured equivalently. Use the full-attention / nocausal configuration (`causal_attn=False`) as the target behavior.
|
||||
|
||||
Required assertions:
|
||||
- same parameter key structure (or compatible `load_state_dict(strict=True)`)
|
||||
- same output shape `(B, T, action_dim)`
|
||||
- `torch.allclose(local_out, external_out, atol=1e-6, rtol=1e-5)`
|
||||
- local model exposes `get_optim_groups(weight_decay=...)` like external repo
|
||||
- use an explicit import path / loader that does not depend on installing the external repo as a package
|
||||
- account for external `ModuleAttrMixin`-style optimizer grouping expectations, including `_dummy_variable` / no-decay bookkeeping if needed for strict state-dict parity
|
||||
- set both models to `eval()` and fix the random seed inside the test so dropout does not create false mismatches
|
||||
|
||||
- [ ] **Step 2: If the test still fails on this branch, use that red state as the TDD starting point**
|
||||
|
||||
Run: `mamba run -n roboimi python -m unittest tests.test_transformer1d_external_alignment -v`
|
||||
Expected: either FAIL before implementation on a fresh checkout, or PASS if Task 2 has already been completed on this branch.
|
||||
|
||||
- [ ] **Step 3a: Port API and state-dict compatibility first**
|
||||
|
||||
Match constructor args, parameter names, embeddings, masks, and strict `state_dict` layout with external `TransformerForDiffusion`.
|
||||
|
||||
- [ ] **Step 3b: Port `_init_weights` and optimizer grouping**
|
||||
|
||||
Match external `_init_weights`, `get_optim_groups`, and `configure_optimizers`.
|
||||
|
||||
- [ ] **Step 3c: Port forward semantics**
|
||||
|
||||
Match external `forward(sample, timestep, cond)` behavior under the full-attention / nocausal configuration.
|
||||
|
||||
- [ ] **Step 3d: Port the minimal external implementation**
|
||||
|
||||
Update `roboimi/vla/models/heads/transformer1d.py` so it matches external repo's native DDPM transformer implementation semantics:
|
||||
- same constructor arguments and defaults where relevant
|
||||
- same token accounting (`time_as_cond`, `obs_as_cond`, `T_cond`)
|
||||
- same parameter naming/layout for embeddings, encoder, decoder, masks
|
||||
- same `_init_weights`
|
||||
- same `get_optim_groups` / `configure_optimizers`
|
||||
- same `forward(sample, timestep, cond)` behavior
|
||||
|
||||
Do **not** port PMF / IMF / DiT branches.
|
||||
|
||||
- [ ] **Step 4: Run test to verify it passes**
|
||||
|
||||
Run: `mamba run -n roboimi python -m unittest tests.test_transformer1d_external_alignment -v`
|
||||
Expected: PASS with strict state-dict load and matching outputs.
|
||||
|
||||
### Task 3: 将当前 Agent 与配置固定到“三相机图像作为条件”的 Transformer diffusion 路线
|
||||
|
||||
**Files:**
|
||||
- Modify: `roboimi/vla/agent.py`
|
||||
- Modify: `roboimi/vla/conf/agent/resnet_transformer.yaml`
|
||||
- Modify: `roboimi/vla/conf/data/simpe_robot_dataset.yaml`
|
||||
- Modify: `roboimi/vla/conf/eval/eval.yaml`
|
||||
- Modify: `roboimi/vla/models/backbones/resnet_diffusion.py`
|
||||
- Test: `tests/test_resnet_transformer_agent_wiring.py`
|
||||
|
||||
- [ ] **Step 1: Write the failing test**
|
||||
|
||||
Write a wiring test that instantiates the Transformer agent config and checks:
|
||||
- `head_type == "transformer"`
|
||||
- `cfg.data.camera_names == cfg.eval.camera_names == ["r_vis", "top", "front"]`
|
||||
- `num_cams == 3`
|
||||
- Transformer `cond_dim == single_cam_feat_dim * 3 + obs_dim`
|
||||
- `predict_action(...)` accepts image conditions and returns `(B, pred_horizon, action_dim)`
|
||||
- test setup must not download weights; override `pretrained_backbone_weights=null` or stub the backbone
|
||||
- assert camera order used for conditioning is tied to the required three cameras, not a stray constant
|
||||
|
||||
- [ ] **Step 2: Run test to verify it fails**
|
||||
|
||||
Run: `mamba run -n roboimi python -m unittest tests.test_resnet_transformer_agent_wiring -v`
|
||||
Expected: FAIL if configs/head wiring do not yet guarantee three-camera conditional setup.
|
||||
|
||||
- [ ] **Step 3: Implement minimal alignment**
|
||||
|
||||
Make the transformer path explicit and stable:
|
||||
- keep image observations always as condition (`cond_dim > 0`, `obs_as_cond` semantics)
|
||||
- keep exactly three cameras: `r_vis`, `top`, `front`
|
||||
- keep current ResNet+SpatialSoftmax backbone unless a test proves incompatibility
|
||||
- make the camera feature order deterministic and aligned to the required three-camera list, not generic key sorting
|
||||
- ensure `agent.per_step_cond_dim` and config `head.cond_dim` stay consistent
|
||||
- ensure eval config uses the same three camera names as training
|
||||
- ensure transformer config follows the external full-attention variant (`causal_attn=false`) instead of the external default causal setting
|
||||
|
||||
- [ ] **Step 4: Run test to verify it passes**
|
||||
|
||||
Run: `mamba run -n roboimi python -m unittest tests.test_resnet_transformer_agent_wiring -v`
|
||||
Expected: PASS
|
||||
|
||||
### Task 4: 让训练脚本在 Transformer 路线上尽量遵循 external repo 的 optimizer/head 使用方式
|
||||
|
||||
**Files:**
|
||||
- Modify: `roboimi/demos/vla_scripts/train_vla.py`
|
||||
- Test: `tests/test_train_vla_transformer_optimizer.py`
|
||||
|
||||
- [ ] **Step 1: Write the failing test**
|
||||
|
||||
Write a test that builds a transformer agent and verifies the training script prefers the head/model supplied optimizer grouping when available (via `get_optim_groups`) instead of blindly using one flat `AdamW(agent.parameters(), ...)`.
|
||||
The test must also prove that every remaining trainable non-head parameter is included exactly once in the optimizer (no silent drops, no duplicates).
|
||||
|
||||
- [ ] **Step 2: Run test to verify it fails**
|
||||
|
||||
Run: `mamba run -n roboimi python -m unittest tests.test_train_vla_transformer_optimizer -v`
|
||||
Expected: FAIL because current training script uses flat optimizer construction.
|
||||
|
||||
- [ ] **Step 3: Implement minimal optimizer alignment**
|
||||
|
||||
Update `train_vla.py` so that:
|
||||
- for transformer head paths, if `agent.noise_pred_net` exposes `get_optim_groups`, build optimizer groups like external repo
|
||||
- explicitly include the remaining trainable non-head parameters (for example the ResNet backbone's non-frozen projection / pooling layers) in the optimizer instead of accidentally dropping them
|
||||
- keep the rest of the simple training loop intact (no EMA unless required later)
|
||||
- do not over-port external workspace abstractions
|
||||
|
||||
- [ ] **Step 4: Run test to verify it passes**
|
||||
|
||||
Run: `mamba run -n roboimi python -m unittest tests.test_train_vla_transformer_optimizer -v`
|
||||
Expected: PASS
|
||||
|
||||
### Task 5: 端到端实例化与最小推理验收
|
||||
|
||||
**Files:**
|
||||
- Verify only
|
||||
|
||||
- [ ] **Step 1: Run all focused unit tests**
|
||||
|
||||
Run:
|
||||
```bash
|
||||
mamba run -n roboimi python -m unittest \
|
||||
tests.test_eval_vla_execution \
|
||||
tests.test_transformer1d_external_alignment \
|
||||
tests.test_resnet_transformer_agent_wiring \
|
||||
tests.test_train_vla_transformer_optimizer -v
|
||||
```
|
||||
Expected: all PASS
|
||||
|
||||
- [ ] **Step 2: Run syntax checks on changed training/eval/model files**
|
||||
|
||||
Run:
|
||||
```bash
|
||||
mamba run -n roboimi python -m py_compile \
|
||||
roboimi/vla/eval_utils.py \
|
||||
roboimi/vla/models/heads/transformer1d.py \
|
||||
roboimi/vla/agent.py \
|
||||
roboimi/demos/vla_scripts/train_vla.py \
|
||||
roboimi/demos/vla_scripts/eval_vla.py
|
||||
```
|
||||
Expected: no syntax errors
|
||||
|
||||
- [ ] **Step 3: Run a local instantiation smoke test**
|
||||
|
||||
Run:
|
||||
|
||||
```bash
|
||||
/home/droid/.conda/envs/roboimi/bin/python - <<'PY'
|
||||
import torch
|
||||
from hydra import compose, initialize_config_dir
|
||||
from hydra.utils import instantiate
|
||||
from pathlib import Path
|
||||
|
||||
config_dir = str((Path.cwd() / "roboimi" / "vla" / "conf").resolve())
|
||||
with initialize_config_dir(version_base=None, config_dir=config_dir):
|
||||
cfg = compose(config_name="config", overrides=[
|
||||
"agent=resnet_transformer",
|
||||
"agent.vision_backbone.pretrained_backbone_weights=null",
|
||||
])
|
||||
|
||||
agent = instantiate(cfg.agent, dataset_stats=None)
|
||||
images = {
|
||||
"r_vis": torch.rand(1, cfg.agent.obs_horizon, 3, 224, 224),
|
||||
"top": torch.rand(1, cfg.agent.obs_horizon, 3, 224, 224),
|
||||
"front": torch.rand(1, cfg.agent.obs_horizon, 3, 224, 224),
|
||||
}
|
||||
qpos = torch.rand(1, cfg.agent.obs_horizon, cfg.agent.obs_dim)
|
||||
out = agent.predict_action(images, qpos)
|
||||
assert out.shape == (1, cfg.agent.pred_horizon, cfg.agent.action_dim), out.shape
|
||||
print("smoke_shape=", out.shape)
|
||||
PY
|
||||
```
|
||||
|
||||
that:
|
||||
- instantiates `agent=resnet_transformer`
|
||||
- constructs fake three-camera input (`r_vis`, `top`, `front`)
|
||||
- calls `agent.predict_action(...)`
|
||||
- verifies output shape is `(1, pred_horizon, 16)`
|
||||
|
||||
Expected: PASS
|
||||
@@ -0,0 +1,202 @@
|
||||
# VLA Experiment Sweep Execution 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:** Probe the largest safe shared batch size, prepare the 10 unique VLA sweep runs, and launch the serial GPU experiment queue with SwanLab logging and 20-epoch headless rollout validation.
|
||||
|
||||
**Architecture:** Use the existing `train_vla.py` / `eval_vla.py` path without changing model code unless launch-blocking issues appear. A small external launcher contract will manage per-run directories, logs, pids, and serial execution while Hydra places trainer-local `checkpoints/` under each run directory.
|
||||
|
||||
**Tech Stack:** zsh, mamba env `roboimi`, Hydra overrides, PyTorch CUDA, SwanLab, existing `train_vla.py`
|
||||
|
||||
---
|
||||
|
||||
### Task 1: Reconfirm schedule math and current runtime assumptions
|
||||
|
||||
**Files:**
|
||||
- Modify: `docs/superpowers/specs/2026-03-30-vla-experiment-sweep-design.md` (only if execution reveals a spec mismatch)
|
||||
- Test: none
|
||||
|
||||
- [ ] **Step 1: Print dataset size and epoch math for the baseline batch size**
|
||||
|
||||
Run:
|
||||
```bash
|
||||
python - <<'PY'
|
||||
num_samples = 70000
|
||||
for batch_size in (32, 48, 64, 80):
|
||||
steps_per_epoch = num_samples // batch_size
|
||||
print(batch_size, steps_per_epoch, 20 * steps_per_epoch, 200 * steps_per_epoch)
|
||||
PY
|
||||
```
|
||||
|
||||
Expected: integer epoch-step mappings for candidate batch sizes.
|
||||
|
||||
- [ ] **Step 2: Reconfirm no stale training process is already active**
|
||||
|
||||
Run:
|
||||
```bash
|
||||
ps -ef | grep -E 'train_vla.py|python .*roboimi/demos/vla_scripts/train_vla.py' | grep -v grep || true
|
||||
```
|
||||
|
||||
Expected: no active training process unless intentionally started by this session.
|
||||
|
||||
- [ ] **Step 3: Record the execution assumptions**
|
||||
|
||||
Capture in session notes:
|
||||
- dataset path
|
||||
- camera list
|
||||
- current branch
|
||||
- no active conflicting training process
|
||||
|
||||
|
||||
### Task 2: Probe the largest safe shared batch size
|
||||
|
||||
**Files:**
|
||||
- Create: `runs/batch-probe-<timestamp>/` (runtime artifact)
|
||||
- Test: probe command output only
|
||||
|
||||
- [ ] **Step 1: Write the candidate probe command for the largest model**
|
||||
|
||||
Use overrides:
|
||||
- `data.dataset_dir=/home/droid/project/diana_sim/sim_transfer`
|
||||
- `data.camera_names=[r_vis,top,front]`
|
||||
- `agent.head.n_emb=384`
|
||||
- `agent.head.n_layer=12`
|
||||
- `agent.head.n_head=4`
|
||||
- `agent.vision_backbone.pretrained_backbone_weights=IMAGENET1K_V1`
|
||||
- `agent.vision_backbone.freeze_backbone=true`
|
||||
- `agent.vision_backbone.use_separate_rgb_encoder_per_camera=true`
|
||||
- `train.device=cuda`
|
||||
- `train.val_split=0.0`
|
||||
- `train.seed=42`
|
||||
- `train.use_swanlab=false`
|
||||
- `train.rollout_val_freq_epochs=0`
|
||||
- `train.rollout_validate_on_checkpoint=false`
|
||||
- `train.max_steps=4`
|
||||
|
||||
- [ ] **Step 2: Run the probe for `batch_size=32`**
|
||||
|
||||
Expected pass condition:
|
||||
- no CUDA OOM
|
||||
- no crash/abort
|
||||
- no NaN/Inf loss
|
||||
- 4 steps complete
|
||||
|
||||
- [ ] **Step 3: Repeat for `batch_size=48`, `64`, then optionally `80`**
|
||||
|
||||
Stop increasing once one candidate fails.
|
||||
|
||||
- [ ] **Step 4: Compute the shared LR and 200-epoch step budget**
|
||||
|
||||
Using:
|
||||
```text
|
||||
lr = 1e-4 * (batch_size / 16)
|
||||
max_steps = 200 * floor(70000 / batch_size)
|
||||
```
|
||||
|
||||
- [ ] **Step 5: Write down the chosen shared batch size and derived max_steps**
|
||||
|
||||
This becomes the execution contract for all runs.
|
||||
|
||||
|
||||
### Task 3: Materialize the 10 unique run matrix
|
||||
|
||||
**Files:**
|
||||
- Create: `runs/vla-sweep-<timestamp>/manifest.txt` (runtime artifact)
|
||||
- Create: `runs/vla-sweep-<timestamp>/launch_queue.sh` (runtime artifact)
|
||||
- Test: shell syntax / manifest inspection
|
||||
|
||||
- [ ] **Step 1: Enumerate the 2-point pretraining comparison**
|
||||
|
||||
Canonical runs:
|
||||
1. `sim-transfer-baseline-pretrain-on-emb128-layer4`
|
||||
2. `sim-transfer-pretrain-off`
|
||||
|
||||
- [ ] **Step 2: Enumerate the remaining 8 unique architecture-sweep runs**
|
||||
|
||||
Skip the duplicate `(pretrain-on, emb128, layer4)` because it reuses the baseline run.
|
||||
|
||||
- [ ] **Step 3: For each run, define exact Hydra overrides**
|
||||
|
||||
Every run must pin:
|
||||
- dataset path
|
||||
- three cameras
|
||||
- `train.device=cuda`
|
||||
- `train.val_split=0.0`
|
||||
- `train.num_workers=12`
|
||||
- `train.seed=42`
|
||||
- `train.rollout_val_freq_epochs=20`
|
||||
- `train.rollout_validate_on_checkpoint=false`
|
||||
- `train.rollout_num_episodes=3`
|
||||
- `train.use_swanlab=true`
|
||||
- a unique `train.swanlab_run_name`
|
||||
- shared batch size
|
||||
- shared LR
|
||||
- derived `max_steps`
|
||||
- `agent.head.n_head=4`
|
||||
- `agent.vision_backbone.freeze_backbone=true`
|
||||
- `agent.vision_backbone.use_separate_rgb_encoder_per_camera=true`
|
||||
|
||||
- [ ] **Step 4: Write the queue launcher script**
|
||||
|
||||
Launcher responsibilities:
|
||||
- create run directory
|
||||
- export `LD_LIBRARY_PATH` CUDA/cuDNN fixup
|
||||
- export `MPLCONFIGDIR=/tmp/mpl`
|
||||
- set `hydra.job.chdir=true`
|
||||
- set `hydra.run.dir=<run_dir>`
|
||||
- tee logs to `<run_dir>/train.log`
|
||||
- write `<run_dir>/train.pid`
|
||||
- run jobs serially
|
||||
|
||||
- [ ] **Step 5: Verify the queue script syntax**
|
||||
|
||||
Run:
|
||||
```bash
|
||||
zsh -n runs/vla-sweep-<timestamp>/launch_queue.sh
|
||||
```
|
||||
|
||||
Expected: exit 0.
|
||||
|
||||
|
||||
### Task 4: Launch the first probe-selected production run and verify startup
|
||||
|
||||
**Files:**
|
||||
- Create: `runs/vla-sweep-<timestamp>/<run-name>/` (runtime artifact)
|
||||
- Test: live log inspection / process inspection
|
||||
|
||||
- [ ] **Step 1: Start the first run in the queue**
|
||||
|
||||
Use the generated launcher contract.
|
||||
|
||||
- [ ] **Step 2: Inspect the first 100-200 log lines**
|
||||
|
||||
Expected:
|
||||
- config resolves correctly
|
||||
- dataset loads
|
||||
- agent initializes on CUDA
|
||||
- training loop starts
|
||||
- SwanLab initializes
|
||||
|
||||
- [ ] **Step 3: Confirm process identity and run directory**
|
||||
|
||||
Collect:
|
||||
- PID
|
||||
- run dir
|
||||
- log path
|
||||
- SwanLab run name
|
||||
|
||||
- [ ] **Step 4: If startup is clean, continue the remaining serial queue**
|
||||
|
||||
If startup fails, stop and debug before launching more runs.
|
||||
|
||||
|
||||
### Task 5: Report the launched sweep contract back to the user
|
||||
|
||||
**Files:**
|
||||
- Test: log evidence only
|
||||
|
||||
- [ ] **Step 1: Summarize the chosen batch size, LR, max_steps, and rollout cadence**
|
||||
|
||||
- [ ] **Step 2: Provide the active run directory / pid / first-run status**
|
||||
|
||||
- [ ] **Step 3: State the remaining queued runs by name**
|
||||
@@ -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
|
||||
@@ -0,0 +1,43 @@
|
||||
# Held-out Episode Validation 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:** Add optional held-out episode validation to main training flow, including explicit val episode selection and periodic action MSE evaluation without merging LEWM-only model features.
|
||||
|
||||
**Architecture:** Extend the generic dataset with optional episode filtering metadata, extend train config with explicit held-out validation knobs, and wire train_vla to choose between random split and explicit episode split. Reuse agent.predict_action_chunk for action MSE so the metric stays model-agnostic.
|
||||
|
||||
**Tech Stack:** Python, Hydra/OmegaConf, PyTorch, unittest
|
||||
|
||||
---
|
||||
|
||||
### Task 1: Add failing tests for dataset episode filtering and config semantics
|
||||
|
||||
**Files:**
|
||||
- Modify: `tests/test_simple_robot_dataset_image_loading.py`
|
||||
- Modify: `tests/test_train_vla_rollout_validation.py`
|
||||
|
||||
- [ ] Add dataset tests for `episode_indices` and `available_episode_indices`.
|
||||
- [ ] Add training tests for explicit held-out episode splitting and fail-fast config validation.
|
||||
- [ ] Run focused tests and verify they fail for the expected missing behavior.
|
||||
|
||||
### Task 2: Implement minimal dataset and training support
|
||||
|
||||
**Files:**
|
||||
- Modify: `roboimi/vla/data/simpe_robot_dataset.py`
|
||||
- Modify: `roboimi/vla/conf/config.yaml`
|
||||
- Modify: `roboimi/demos/vla_scripts/train_vla.py`
|
||||
|
||||
- [ ] Add config keys `train.val_episode_indices` and `train.action_mse_val_freq_epochs`.
|
||||
- [ ] Add optional dataset filtering by episode index plus `available_episode_indices` metadata.
|
||||
- [ ] Add explicit train/val dataset builder and held-out action MSE computation.
|
||||
- [ ] Log `val/action_mse` only when explicit held-out episode validation is configured.
|
||||
|
||||
### Task 3: Verify focused coverage
|
||||
|
||||
**Files:**
|
||||
- Test: `tests/test_simple_robot_dataset_image_loading.py`
|
||||
- Test: `tests/test_train_vla_rollout_validation.py`
|
||||
- Test: `tests/test_train_vla_swanlab_logging.py`
|
||||
|
||||
- [ ] Run focused unittest targets for dataset filtering, held-out MSE, and SwanLab logging.
|
||||
- [ ] Fix any regressions with minimal code changes.
|
||||
@@ -0,0 +1,242 @@
|
||||
# VLA Experiment Sweep Design
|
||||
|
||||
## Goal
|
||||
|
||||
Run a matched-configuration training sweep on `/home/droid/project/diana_sim/sim_transfer` for the aligned Transformer diffusion policy, comparing:
|
||||
|
||||
1. ResNet ImageNet pretraining vs no pretraining under matched settings.
|
||||
2. Transformer width/depth choices across a 3×3 grid.
|
||||
|
||||
The sweep should use the full dataset, headless rollout validation every 20 epochs, SwanLab logging, and a total budget of about 200 epochs per run.
|
||||
|
||||
## Current Context
|
||||
|
||||
- Dataset: `100` episodes, `70000` frame-level samples.
|
||||
- Cameras used as image conditions: `r_vis`, `top`, `front`.
|
||||
- Current policy family: ResNet visual encoder + Transformer DDPM head.
|
||||
- Current training path already supports:
|
||||
- full-dataset training (`val_split=0.0`)
|
||||
- checkpoint-time and epoch-time rollout hooks
|
||||
- headless eval path
|
||||
- SwanLab logging
|
||||
- best-checkpoint selection by rollout reward once rollout metrics exist
|
||||
- Current trainer behavior that matters for execution:
|
||||
- step checkpoints are written to `checkpoints/` relative to the Hydra job working directory
|
||||
- epoch-time rollout validation is controlled by `train.rollout_val_freq_epochs`
|
||||
- checkpoint-time rollout validation is separately controlled by `train.rollout_validate_on_checkpoint`
|
||||
- per-run log files, pid files, and run directories must be created by the external launcher, not by the trainer itself
|
||||
|
||||
## Derived Schedule
|
||||
|
||||
For a chosen global batch size `B`, define:
|
||||
|
||||
- `steps_per_epoch = floor(70000 / B)` because training currently uses `drop_last=True` when the train set is larger than the batch.
|
||||
- `rollout_freq_steps = 20 * steps_per_epoch`
|
||||
- `max_steps = 200 * steps_per_epoch`
|
||||
|
||||
Example at `B = 32`:
|
||||
|
||||
- `steps_per_epoch = 2187`
|
||||
- `rollout_freq_steps = 43740`
|
||||
- `max_steps = 437400`
|
||||
|
||||
This schedule must be recomputed after the batch-size probe chooses the final shared batch size.
|
||||
|
||||
## Experiment Strategy
|
||||
|
||||
### Recommended approach
|
||||
|
||||
Use one **shared batch size across all runs**, chosen by probing the largest model that will appear in the sweep. This preserves fairness while still using as much GPU memory as safely available.
|
||||
|
||||
### Why not per-run max batch size
|
||||
|
||||
Allowing each model to use a different batch size would add another changing variable, making the 11-run comparison harder to interpret.
|
||||
|
||||
### Unique executions vs logical comparisons
|
||||
|
||||
There are `11` logical comparison points in the user request (`2` for the pretraining comparison and `9` for the emb/layer grid), but only `10` **unique** training configurations because:
|
||||
|
||||
- the pretrained baseline in Section A (`pretrained_backbone_weights=IMAGENET1K_V1`, `n_emb=128`, `n_layer=4`) is also one cell of the Section B sweep.
|
||||
|
||||
This design will therefore run **10 unique training jobs**, and the pretrained-on baseline will be reused as the `(128, 4)` architecture-sweep point. This avoids wasting one full 200-epoch run on an identical configuration.
|
||||
|
||||
The reused baseline has one canonical identity throughout execution:
|
||||
|
||||
- run name stem: `sim-transfer-baseline-pretrain-on-emb128-layer4`
|
||||
- this same run is counted both as:
|
||||
- the pretrained-on arm of the 2-point pretraining comparison
|
||||
- the `(n_emb=128, n_layer=4)` point of the 9-point architecture sweep
|
||||
|
||||
## Batch-Size Probe
|
||||
|
||||
Before launching the full sweep:
|
||||
|
||||
1. Probe the largest planned model:
|
||||
- `n_emb=384`
|
||||
- `n_layer=12`
|
||||
- ResNet pretrained enabled
|
||||
2. Try candidate batch sizes in ascending order:
|
||||
- `32`
|
||||
- `48`
|
||||
- `64`
|
||||
- optionally `80` if `64` is clearly safe
|
||||
3. For each candidate, run a short GPU smoke training with:
|
||||
- real dataset
|
||||
- three cameras
|
||||
- no SwanLab
|
||||
- no rollout validation
|
||||
- `max_steps=4`
|
||||
4. A candidate **passes** if all 4 steps finish without:
|
||||
- CUDA OOM
|
||||
- process crash / abort
|
||||
- NaN / Inf loss
|
||||
5. Choose the **highest passing** candidate and use it for all sweep runs.
|
||||
|
||||
Learning rate should scale linearly using the current training rule:
|
||||
|
||||
- config baseline: `lr=1e-4` at `batch_size=16`
|
||||
- equivalent rule: `lr = 1e-4 * (batch_size / 16)`
|
||||
- equivalently: `lr = 2e-4 * (batch_size / 32)`
|
||||
|
||||
## Experiment Matrix
|
||||
|
||||
Total logical comparison points: `11`
|
||||
|
||||
Total unique training executions: `10`
|
||||
|
||||
### A. ResNet pretraining comparison (2 runs)
|
||||
|
||||
Fixed Transformer settings:
|
||||
|
||||
- `n_emb=128`
|
||||
- `n_layer=4`
|
||||
- `n_head=4`
|
||||
- `freeze_backbone=true`
|
||||
- `use_separate_rgb_encoder_per_camera=true`
|
||||
|
||||
Interpretation note:
|
||||
|
||||
- This comparison keeps the current frozen-backbone recipe unchanged.
|
||||
- Therefore, the `pretrained_backbone_weights=null` run measures **random frozen visual features vs ImageNet frozen visual features**.
|
||||
- It does **not** test end-to-end training of the ResNet from scratch; that would be a separate experiment.
|
||||
|
||||
Runs:
|
||||
|
||||
1. `pretrained_backbone_weights=IMAGENET1K_V1`
|
||||
2. `pretrained_backbone_weights=null`
|
||||
|
||||
### B. Transformer architecture sweep (9 runs)
|
||||
|
||||
Fixed visual backbone setting:
|
||||
|
||||
- `pretrained_backbone_weights=IMAGENET1K_V1`
|
||||
- `freeze_backbone=true`
|
||||
- `use_separate_rgb_encoder_per_camera=true`
|
||||
- `n_head=4`
|
||||
|
||||
Rationale:
|
||||
|
||||
- The 9-point emb/layer sweep fixes `pretrained_backbone_weights=IMAGENET1K_V1`.
|
||||
- The separate 2-run comparison in Section A isolates the effect of ResNet pretraining at the baseline architecture.
|
||||
|
||||
Grid:
|
||||
|
||||
- `n_emb ∈ {128, 256, 384}`
|
||||
- `n_layer ∈ {4, 8, 12}`
|
||||
|
||||
## Shared Training Configuration
|
||||
|
||||
All full runs should use:
|
||||
|
||||
- dataset: `/home/droid/project/diana_sim/sim_transfer`
|
||||
- cameras: `[r_vis, top, front]`
|
||||
- device: `cuda`
|
||||
- `val_split=0.0`
|
||||
- `num_workers=12`
|
||||
- `seed=42`
|
||||
- headless rollout validation
|
||||
- `rollout_val_freq_epochs=20`
|
||||
- `rollout_validate_on_checkpoint=false`
|
||||
- `rollout_num_episodes=3`
|
||||
- `use_swanlab=true`
|
||||
- shared batch size from probe
|
||||
- shared learning-rate rule from probe
|
||||
- `max_steps` recomputed to represent about `200` epochs
|
||||
|
||||
Reproducibility note:
|
||||
|
||||
- `train.seed=42` will be fixed and recorded for all runs.
|
||||
- Under the current `train_vla.py`, this does **not** make the runs bitwise deterministic when `val_split=0.0`, because model initialization and shuffled training order are not fully seeded by the trainer today.
|
||||
- The sweep is therefore treated as a **matched-config comparison**, not a strict deterministic benchmark.
|
||||
|
||||
## Rollout / Validation Policy
|
||||
|
||||
- No held-out val split will be used.
|
||||
- Loss remains a fallback best-model metric until the first rollout reward exists.
|
||||
- After rollout metrics begin, the best checkpoint should be chosen by **higher average rollout reward**.
|
||||
- Rollout should run with:
|
||||
- no GUI
|
||||
- CPU eval device
|
||||
- current evaluation helper path
|
||||
|
||||
## Execution Model
|
||||
|
||||
Runs should be launched **serially**, not concurrently, because:
|
||||
|
||||
- the user currently has one GPU available
|
||||
- concurrent jobs would reduce throughput and complicate memory management
|
||||
- serial execution preserves clean logs and easier resume semantics
|
||||
|
||||
Each run gets:
|
||||
|
||||
- a dedicated run directory under `runs/`
|
||||
- its own log file
|
||||
- its own pid / launcher metadata if backgrounded
|
||||
- a unique SwanLab run name
|
||||
- `hydra.job.chdir=true`
|
||||
- `hydra.run.dir=<run_dir>` so the trainer-local `checkpoints/` land inside the run directory
|
||||
|
||||
## Naming Convention
|
||||
|
||||
Suggested run name pattern:
|
||||
|
||||
- pretraining comparison:
|
||||
- `sim-transfer-baseline-pretrain-on-emb128-layer4-bs{B}-lr{LR}-e200`
|
||||
- `sim-transfer-pretrain-off-bs{B}-lr{LR}-e200`
|
||||
- architecture sweep:
|
||||
- `(128,4)` reuses `sim-transfer-baseline-pretrain-on-emb128-layer4-bs{B}-lr{LR}-e200`
|
||||
- `sim-transfer-emb{E}-layer{L}-bs{B}-lr{LR}-e200`
|
||||
|
||||
## Success Criteria
|
||||
|
||||
The sweep is successful when:
|
||||
|
||||
1. A shared safe batch size has been found.
|
||||
2. All 10 unique runs have launchable commands and isolated run directories.
|
||||
3. Each run trains for about 200 epochs worth of steps.
|
||||
4. Rollout validation occurs every 20 epochs.
|
||||
5. SwanLab records train metrics and rollout rewards.
|
||||
6. Final comparison can be made using best rollout average reward and training traces across:
|
||||
- 2-point pretraining comparison
|
||||
- 9-point emb/layer sweep
|
||||
|
||||
## Non-Goals
|
||||
|
||||
- No change to dataset format.
|
||||
- No change to model architecture beyond the requested pretraining toggle and Transformer width/depth sweep.
|
||||
- No attempt to equalize wall-clock time across models.
|
||||
- No multi-GPU scheduling.
|
||||
|
||||
## Risks and Mitigations
|
||||
|
||||
### Risk: larger models may OOM
|
||||
|
||||
Mitigation: probe the largest model first and choose one shared safe batch size.
|
||||
|
||||
### Risk: long runs may be interrupted
|
||||
|
||||
Mitigation: use per-run directories, periodic checkpoints, and resumable launch commands.
|
||||
|
||||
### Risk: rollout validation increases total training time
|
||||
|
||||
Mitigation: keep rollout headless, CPU-based, and limited to 3 episodes every 20 epochs.
|
||||
Reference in New Issue
Block a user