roboimi/docs/superpowers/specs/2026-04-05-phase2-full-attnres-vision-design.md

Phase-2 Full-AttnRes Vision Design

Goal

在当前 roboimi IMF policy 中，把视觉 backbone 里原先由 ResNet BasicBlock/Bottleneck 提供的残差单元全部替换为 AttnRes 风格单元，同时尽量保持现有 agent / cond / rollout / 训练脚本接口不变。

User requirement interpretation

这里按最严格解释执行：

• 不是“在 ResNet 后面再加一个 AttnRes 模块”
• 也不是“只在某几个 stage 加 AttnRes 混合”
• 而是：视觉主干网络中原本依赖 ResNet residual block 的地方，统一改成 AttnRes residual operator 驱动的 block
• 最终仍然输出与现有 ResNetDiffusionBackbone 相同的每相机特征接口，以便复用 SpatialSoftmax -> Linear -> ReLU、多相机拼接、state concat、IMF head 条件输入

Recommended design

Option A (recommended)

保留 ResNet 的宏观 stage/stem 结构与通道/步幅规划，但把每个 stage 内的 BasicBlock/Bottleneck 替换为新的 AttnResImageBlock2D：

• 输入仍是 (B, C, H, W) feature map
• block 内先把空间维 flatten 成 token 序列 (B, H*W, C)
• 用二维位置编码 / 可学习位置偏置 + AttnRes self-attention + AttnRes FFN 完成 block 变换
• 再 reshape 回 (B, C, H, W)
• stage 间下采样仍由显式 stride/downsample path 完成

优点：

• 最接近“ResNet 中所有残差都由 AttnRes 代替”的要求
• 保留现有视觉输出接口和 cond_dim，不用改 agent/head/data pipeline
• 仍可沿用现有多相机编码器框架

缺点：

• 需要新写 2D 版 AttnRes image block，而不是直接复用 1D IMF head block

Option B

完全移除 ResNet stage，换成 patchify + ViT/AttnRes 图像 transformer，再接 SpatialSoftmax/MLP。

优点：实现概念更统一。
缺点：已经不算“把 ResNet 中残差替换掉”，而是直接换 backbone，和用户要求不完全一致。

Option C

保留现有 ResNet block，只在 block 外层加 AttnRes mixing。

不推荐，因为不满足“所有残差均由 AttnRes 替代”。

Concrete architecture choice

采用 Option A：

1. 保留 stem（conv/bn-or-gn/relu/maxpool）与 stage 边界
2. 新增 AttnResImageBlock2D
3. 新增 AttnResResNetLikeBackbone2D，负责堆叠 stage/block
4. 在 ResNetDiffusionBackbone 中增加可选 backbone mode，例如：
   • vision_backbone_mode: resnet
   • vision_backbone_mode: attnres_resnet
5. resnet_imf_attnres agent 配置新增一个 Phase-2 变体，默认打开 attnres_resnet
6. 仍保持：
   • 每相机输出 64
   • 多相机总视觉输出 3 * 64
   • 与 state 拼接后 cond_dim = 208

Files likely to change

• roboimi/vla/models/backbones/resnet_diffusion.py
• roboimi/vla/conf/backbone/resnet_diffusion.yaml
• roboimi/vla/conf/agent/resnet_imf_attnres.yaml
• new: roboimi/vla/models/backbones/attnres_resnet2d.py
• tests:
  • new: tests/test_attnres_resnet2d_backbone.py
  • update/add wiring test for agent cond dims

Test plan

1. New backbone instantiates and forwards (B,T,C,H,W) multi-camera input
2. Output shape unchanged vs current backbone
3. output_dim == 64
4. 3-camera cond path still yields 208
5. Phase-2 config instantiates full IMF agent successfully
6. One short CPU smoke forward for compute_loss

Phase-2 experiment plan

固定使用 Phase-1 最优组合：

• pred_horizon=16
• num_action_steps=8

比较：

1. baseline: current IMF head-only AttnRes + original ResNet vision backbone
2. phase2: IMF head AttnRes + full AttnRes-replaced vision backbone

训练超参保持与 Phase-1 最优设置一致，先跑一组 50k step 对比。