273 lines
10 KiB
Markdown
273 lines
10 KiB
Markdown
# IMF-AttnRes Policy Migration Design
|
||
|
||
**Date:** 2026-04-01
|
||
**Status:** Approved in chat, written spec pending review
|
||
|
||
## Goal
|
||
|
||
将 `/home/droid/project/diffusion_policy` 中提交 `185ed659` 的 IMF-AttnRes diffusion policy 迁移到当前 `roboimi` 仓库，作为当前 DiT / Transformer diffusion policy 的替代训练选项；同时迁移其训练目标与一步推理机制，并保持 RoboIMI 现有的仿真环境、三相机视觉输入、数据集格式、训练脚本和 rollout 验证工作流可继续使用。
|
||
|
||
## Non-Goals
|
||
|
||
- 不迁移 external repo 中与当前任务无关的 obs encoder、dataset、env wrapper、PushT 专用逻辑。
|
||
- 不强行复刻 external repo 中全部目录结构；仅迁移当前 RoboIMI 训练所必需的模型、loss、inference 语义。
|
||
- 不在本次工作中同时保留旧 DiT 为默认训练目标；旧配置继续可用，但新模型单独提供 config 入口。
|
||
|
||
## User-Confirmed Requirements
|
||
|
||
1. 迁移对象是 `185ed659` 中的 **IMF-AttnRes 模型相关代码**。
|
||
2. 不只是迁移骨架，还要迁移：
|
||
- **训练目标**
|
||
- **一步推理机制**
|
||
3. 视觉输入与当前 RoboIMI diffusion policy 一致：
|
||
- 使用三个相机图像作为条件输入
|
||
- 图像观测必须作为条件，而不是拼进输出预测目标
|
||
4. 当前任务里，IMF policy 用来替代现有 DiT/Transformer diffusion policy 训练。
|
||
5. 训练参数沿用最近一次训练的大体设置（后续由训练命令显式覆盖），但推理方式改为 IMF 的 one-step 机制。
|
||
6. 用户接受 IMF 中“全注意力 / 非因果注意力”的实现约束。
|
||
|
||
## External Source of Truth
|
||
|
||
迁移语义以 external repo 的以下文件为准：
|
||
|
||
- `diffusion_policy/model/diffusion/attnres_transformer_components.py`
|
||
- `diffusion_policy/model/diffusion/imf_transformer_for_diffusion.py`
|
||
- `diffusion_policy/policy/imf_transformer_hybrid_image_policy.py`
|
||
- 参考配置：`image_pusht_diffusion_policy_dit_imf_attnres_full.yaml`
|
||
|
||
其中最关键的差异是：该策略并非 DDPM/DDIM 多步去噪，而是 IMF 训练目标 + one-step 推理。
|
||
|
||
## Current RoboIMI Baseline
|
||
|
||
当前 RoboIMI 中与该任务直接相关的基线如下：
|
||
|
||
- 视觉编码：`ResNetDiffusionBackbone`
|
||
- 三相机：`r_vis`, `top`, `front`
|
||
- 每个时间步将相机特征与 `qpos` 拼接为 per-step condition
|
||
- 策略主体：`VLAAgent`
|
||
- `compute_loss()` 使用 DDPM 噪声预测损失
|
||
- `predict_action()` 使用 DDIM 多步采样
|
||
- 在线控制通过动作队列机制在 `select_action()` 中按 chunk 触发预测
|
||
- 训练脚本：`roboimi/demos/vla_scripts/train_vla.py`
|
||
- 支持 GPU 训练、SwanLab 日志、headless rollout 验证
|
||
|
||
因此，本次迁移的核心不是换视觉 backbone，而是替换 **head + loss + inference semantics**。
|
||
|
||
## Recommended Integration Approach
|
||
|
||
采用 **最小侵入式集成**：
|
||
|
||
1. **保留当前 RoboIMI 的视觉编码、数据读取、rollout/eval、训练脚本主框架**。
|
||
2. **新增 IMF 专用 head 模块**，在 RoboIMI 内本地实现：
|
||
- AttnRes 组件
|
||
- IMF transformer 主体
|
||
3. **新增 IMF 专用 agent**，复用当前 `VLAAgent` 的：
|
||
- 归一化逻辑
|
||
- 相机顺序管理
|
||
- 观测缓存 / 动作 chunk 缓存
|
||
- rollout 接口
|
||
但覆盖：
|
||
- `compute_loss()`
|
||
- `predict_action()`
|
||
4. **新增独立 Hydra config**，让 IMF policy 作为新的 agent 选项，不破坏已有 resnet_transformer / gr00t_dit 配置。
|
||
|
||
这样做的原因：
|
||
|
||
- 迁移 IMF 语义时不必把当前 DDPM agent 搅乱；
|
||
- rollout / eval / checkpoint 逻辑仍然可复用；
|
||
- 便于和现有 Transformer / DiT 直接做 A/B 对比训练。
|
||
|
||
## Architecture
|
||
|
||
### 1. Observation / Conditioning Path
|
||
|
||
沿用当前 RoboIMI 的视觉路径：
|
||
|
||
- 输入观测：`images={r_vis, top, front}` + `qpos`
|
||
- `ResNetDiffusionBackbone` 对每个相机编码，得到 per-camera feature
|
||
- `state_encoder` 编码 `qpos`
|
||
- 将三相机特征与 state feature 按时间步拼接，形成 `per_step_cond`
|
||
|
||
这里不迁移 external repo 的 obs_encoder 实现；我们只对齐 **“图像作为条件 token 输入 transformer”** 这一语义。
|
||
|
||
### 2. Condition Tokenization
|
||
|
||
对齐 external IMF transformer 的 token 使用方式：
|
||
|
||
- action trajectory token：由 `(B, pred_horizon, action_dim)` 通过线性层映射到 `n_emb`
|
||
- time token：两个标量 `r` 与 `t`，分别通过 sinusoidal embedding + linear projection 得到 token
|
||
- observation token：`per_step_cond` 通过线性层映射到 `n_emb`
|
||
- 最终 token 序列为：
|
||
- `[r_token, t_token, obs_cond_tokens..., action_tokens...]`
|
||
|
||
在当前任务中，obs token 数量等于 `obs_horizon`，且图像观测始终作为条件输入。
|
||
|
||
### 3. IMF-AttnRes Backbone
|
||
|
||
在 RoboIMI 内新增 AttnRes backbone 实现，保持 external commit 的关键语义：
|
||
|
||
- `RMSNorm` / `RMSNormNoWeight`
|
||
- RoPE
|
||
- Grouped Query Self-Attention
|
||
- SwiGLU FFN
|
||
- AttnRes operator / residual source aggregation
|
||
- `AttnResTransformerBackbone`
|
||
|
||
并保持：
|
||
|
||
- **full attention**（不使用因果注意力）
|
||
- `backbone_type='attnres_full'`
|
||
- 输出仅切回 action token 部分，再经过最终 norm + head 得到 velocity-like 输出
|
||
|
||
### 4. Training Objective
|
||
|
||
训练目标从当前 DDPM epsilon prediction 改为 external IMF 目标：
|
||
|
||
给定真实轨迹 `x` 与随机噪声 `e`：
|
||
|
||
1. 采样 `t ~ U(0,1)`、`r ~ U(0,1)`，并排序为 `t >= r`
|
||
2. 构造插值状态：
|
||
- `z_t = (1 - t) x + t e`
|
||
3. 用模型计算：
|
||
- `v = f(z_t, t, t, cond)`
|
||
4. 对 `g(z, r, t) = f(z, r, t, cond)` 做 JVP，得到：
|
||
- `u, du_dt`
|
||
5. 构造 compound velocity：
|
||
- `V = u + (t - r) * du_dt`
|
||
6. 目标为：
|
||
- `target = e - x`
|
||
7. 用 action 维度上的 MSE 作为最终损失
|
||
|
||
RoboIMI 现有 batch 中的 `action_is_pad` 仍要保留支持；如果存在 padding，只在有效 action 上计算损失。
|
||
|
||
### 5. One-Step Inference
|
||
|
||
推理改为 external IMF 的一步采样语义：
|
||
|
||
1. 从标准高斯初始化 action trajectory `z_t`
|
||
2. 计算 `u = f(z_t, r=0, t=1, cond)`
|
||
3. 一步更新：
|
||
- `x_hat = z_t - (t-r) * u = z_t - u`
|
||
4. 反归一化得到动作序列
|
||
|
||
这意味着：
|
||
|
||
- `num_inference_steps` 对 IMF policy 固定为 `1`
|
||
- 不再调用 DDIM scheduler 的多步 `step()`
|
||
- 在线控制中仍沿用当前 chunk 机制：
|
||
- 动作队列为空时触发一次 `predict_action_chunk()`
|
||
- 取预测序列中 `[obs_horizon-1 : obs_horizon-1+num_action_steps]` 这一段入队
|
||
|
||
也就是说，**触发模型前向的规则不变，改变的是每次触发后的动作序列生成方式**。
|
||
|
||
## API / Code Structure
|
||
|
||
计划中的主要代码边界如下：
|
||
|
||
- `roboimi/vla/models/heads/attnres_transformer_components.py`
|
||
- IMF AttnRes 基础组件
|
||
- `roboimi/vla/models/heads/imf_transformer1d.py`
|
||
- RoboIMI 版本 IMF transformer head
|
||
- 对外暴露 `forward(sample, r, t, cond=None)`
|
||
- 暴露 `get_optim_groups()` 供 AdamW 分组使用
|
||
- `roboimi/vla/agent_imf.py`
|
||
- 复用 `VLAAgent` 的观测处理 / normalization / queue 基础设施
|
||
- 覆盖 IMF 的训练损失与 one-step 预测逻辑
|
||
- Hydra config
|
||
- `roboimi/vla/conf/head/imf_transformer1d.yaml`
|
||
- `roboimi/vla/conf/agent/resnet_imf_attnres.yaml`
|
||
|
||
训练脚本主流程尽量不改；只要求它能 instantiate 新 agent 并继续使用当前 rollout / checkpoint / swanlab 逻辑。
|
||
|
||
## Compatibility Decisions
|
||
|
||
## Initial Config Defaults To Preserve
|
||
|
||
为避免迁移时语义漂移，首版 IMF 配置默认值明确固定为：
|
||
|
||
- `backbone_type: attnres_full`
|
||
- `n_head: 1`
|
||
- `n_kv_head: 1`
|
||
- `n_cond_layers: 0`
|
||
- `time_as_cond: true`
|
||
- `causal_attn: false`
|
||
- `num_inference_steps: 1`
|
||
|
||
这些默认值与 external `185ed659` 的 IMF-AttnRes 使用方式保持一致；后续调参可以覆盖，但首版迁移必须先以该语义跑通。
|
||
|
||
### Reuse From RoboIMI
|
||
|
||
保留：
|
||
|
||
- 三相机数据读取方式
|
||
- ResNet visual backbone
|
||
- qpos / action normalization
|
||
- 训练循环、优化器、scheduler、SwanLab、headless rollout
|
||
- `select_action()` 的在线 chunk 执行方式
|
||
|
||
### Replace With External IMF Semantics
|
||
|
||
替换：
|
||
|
||
- transformer head 实现
|
||
- diffusion training objective
|
||
- inference sampling semantics
|
||
|
||
### Intentionally Not Mirrored 1:1
|
||
|
||
不强行与 external repo 一致的部分：
|
||
|
||
- external repo 的整体 policy 基类继承体系
|
||
- external repo 的 obs encoder 模块树
|
||
- external repo 的 normalizer / mask generator 框架
|
||
|
||
原因是当前 RoboIMI 已有稳定的数据接口和 rollout 流程，直接嫁接进去更稳。
|
||
|
||
## Testing / Verification Strategy
|
||
|
||
迁移完成后至少验证以下内容：
|
||
|
||
1. **单元 / 冒烟验证**
|
||
- IMF head 前向 shape 正确
|
||
- IMF agent `compute_loss()` 在真实 batch 上可前向、反向
|
||
- IMF agent `predict_action()` 能输出 `(B, pred_horizon, action_dim)`
|
||
2. **训练链路验证**
|
||
- 使用 GPU 跑一个短训练任务，确认：
|
||
- dataloader 正常
|
||
- optimizer / lr scheduler 正常
|
||
- SwanLab 正常记录配置和训练指标
|
||
3. **rollout 验证**
|
||
- 训练中周期性 headless rollout 能跑通
|
||
- 环境仍按 EE-style `step()` 接收动作
|
||
4. **最终交付**
|
||
- 用用户指定的同类超参数启动正式训练
|
||
|
||
## Risks and Mitigations
|
||
|
||
### Risk 1: JVP 在 CUDA 注意力内核上不稳定
|
||
|
||
缓解：沿用 external repo 的策略，在 JVP 路径上切换到 math SDP kernel，必要时 fallback 到 `torch.autograd.functional.jvp`。同时，JVP 的切线构造与 `u, du_dt` 计算流程必须严格对齐 external source，不在本次迁移中自行改写其数学语义。
|
||
|
||
### Risk 2: Optimizer 参数分组遗漏新模块
|
||
|
||
缓解：IMF head 提供 `get_optim_groups()`，并在训练脚本中按“只要 head 提供该接口就使用”的策略统一处理，而不是绑定旧 `head_type`。
|
||
|
||
### Risk 3: 现有 rollout 逻辑假定 DDIM 多步采样
|
||
|
||
缓解：保持 `select_action()` / `predict_action_chunk()` 接口不变，只替换 `predict_action()` 内部实现，确保 eval 代码无需理解 IMF 细节。
|
||
|
||
### Risk 4: 训练命令参数与新 config 不一致
|
||
|
||
缓解：新增独立 agent config，并保留此前训练参数作为显式 CLI override 模板。
|
||
|
||
## Success Criteria
|
||
|
||
以下条件全部满足，视为本次迁移成功：
|
||
|
||
1. RoboIMI 中新增 IMF-AttnRes policy，可通过 Hydra config 单独启用。
|
||
2. 训练时使用 external IMF 的 loss，而不是当前 DDPM epsilon loss。
|
||
3. 推理时使用 one-step IMF 采样，而不是 DDIM 多步采样。
|
||
4. 三相机图像始终作为条件输入参与模型前向。
|
||
5. 在线 rollout 能在 headless 仿真环境中跑通。
|
||
6. 能按最近一次实验参数模板成功启动训练。
|