# 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