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
|
||||
Reference in New Issue
Block a user