Blame for docs/bridge_memory_plan.md

# Bridge 外部记忆库：现状、消融与适配方案

检查日期：2026-09-24。本文是代码/权重审计与实施方案，**尚未启动 Bridge
CTE 或检索头训练，也没有替换正在准备的全量训练缓存**。

## 当前记忆如何影响训练

当前使用 `/home/gpu4/tianyi/zeva-release/weights/stage1/` 中的
`zeva_cte.pt` 和 `train_memory_effect_v3.pt`，以及 `stage3/best.pt`。

记忆库、CTE、检索头本身被冻结。`prepare_widowx_features.py` 从初始图像和
指令检索得到 `behavior_global`，缓存后作为 Stage-2 输入。
它经过 `behavior_global_projector` 和 `behavior_pbd` 影响动作预测及梯度，
因此“冻结记忆库”**不等于它不影响权重更新**。监督动作仍来自 Bridge 样本，
不是 RoboCasa 的动作标签。可训练参数名单见 `scripts/widowx_config.py:25`。

关键实现（路径相对 `/home/gpu4/tianyi/cosmos-framework`）：

- `cosmos_framework/model/generator/omni_mot_model.py:698`：接入缓存行为特征。
- `cosmos_framework/model/generator/mot/cosmos3_vfm_network.py:718`：全局记忆前缀投影。
- `cosmos_framework/model/generator/mot/cosmos3_vfm_network.py:1015`：记忆参与动作先验。
- `cosmos_framework/model/generator/omni_mot_model.py:1734`：动作先验监督损失。

这是**跨域、不匹配的 RoboCasa 记忆**，但仅凭这一点不能断言它是成功率低的
唯一原因。当前静态检索库也不要与 Zeva 的在线 PIM 混为一谈；本地 WidowX
配置继承的 Atomic5 配方未启用 PIM/online-memory。

## 实际权重内容

`train_memory_effect_v3.pt` 有 2,495 条记录，均来自 RoboCasa：
TurnOnMicrowave 527、TurnOnElectricKettle 504、CloseToasterOvenDoor 490、
CoffeeSetupMug 487、OpenStandMixerHead 487。没有 Bridge 条目。

每条包含 `episode_id`、`source_index`、`instruction`、`length`，
以及归一化的 128 维 `retrieval_key` 和 256 维 `behavior_value`。
库格式为 `cosmos_behavior_memory_effect_v3`。
现有检索头输入 4,096 维，隐藏层 1,024 维，输出归一化的 128 维查询；
余弦检索后对 top-5 的 256 维 value 加权。

检索头 checkpoint 的 Recall@1=1.0 是其 RoboCasa 五类验证集指标，
不是 Bridge 的指标，也不能证明在 Bridge 中可用。

## 本次消融

`scripts/widowx_policy_server.py` 增加 `--memory-mode bank_off`：
不加载外部库和检索头，令全局输入为零；保留真实当前轨迹的 CTE phase/effect
历史及全部已训练策略权重。投影层和 PBD 的已学偏置仍保留，
因此这是“去除外部记忆内容”，不是“关闭全部 Zeva 模块”。

固定旧训练最后保存的 `step_000290000.pt`，每组种子 0–19、150 个环境步、
30 步去噪、guidance=3。按每次执行 4/16 个动作，各做外部库 off/on 对照。
模型仍预测 32 个动作；执行 16 个动作后再规划，在 5 Hz 下对应 3.2 秒。
所有实际执行动作都会进入历史，CTE 的四步边界不随执行长度改变。

此消融仅发生在评测时：不能消除已有权重对旧记忆的依赖，也不等同于
“从头关闭记忆训练”。关闭更好可提示旧条件有害，关闭更差也可能是输入
分布改变；小样本成功率均需结合视频、抓取/移动指标和更多种子判断。

运行入口和输出说明见 `scripts/WIDOWX.md` 的 External-bank ablation 一节。

## Bridge 构建方案（以下步骤需要新增/补齐实现）

本地已找到 CTE 模型/损失、检索头模型/对比损失、库加载与推理代码，
**未找到可以直接运行的 Bridge Stage-1 训练器、记忆库导出器或 Stage-3
训练器**。不要把下面的步骤理解为已有完整官方命令。

1. **整理任务标签和划分数据。** 本地 IPEC 导出有 53,192 个 episode；
   按 `meta/episodes.jsonl` 的任务文本统计，38,660 个有非空文本，
   14,532 个无非空文本（27.32%，分母是 episode，不是采样窗口）。
   共有 19,973 个非空任务字符串，其中 18,380 个只对应一个 episode。
   先做可审计的同义任务归并，保留物体、方向、目标容器等差别；
   “放胡萝卜到盘里”与“从盘里拿走”不能合并。空文本不能全部视为同一种任务。
   先按 episode/原始采集 session 划分 train/val，验证和仿真评测数据不入库。
   建议先选约 5–10 个有足够重复样本的任务试跑，再扩展全量。

2. **固定 Bridge 输入契约。** 使用真实前视 `image_0`、5 Hz、现有
   XYZ/RPY 六维归一化与夹爪正开负闭约定、每四个控制步一个 CTE 边界。
   CTE 输入包括潜变量图像 `[B,T,48,H,W]`、动作 `[B,T-1,4,7]`、
   帧/动作有效掩码，动作只包含已经执行的历史。
   要记录分辨率、RGB 约定、VAE 和所有权重指纹，并保持训练/服务一致。
   当前外部检索 readout 仍沿用 RoboCasa 的 20 FPS/拼接视图提示模板
   （上游 server 的 `format_action_prompt`/`extract_batch`），真正适配时应与
   Bridge 单视图、5 Hz 一并修正并重新训练检索头，不能单独修改线上提示。

3. **训练或适配 Bridge CTE。** 可从发布 CTE 初始化，保留模型定义和
   `cte_losses.py` 中的视觉/动作预测、语义聚类、阶段和 effect 目标。
   每次优化后调用 `update_ema_target()`。不足 16 个控制步的轨迹没有完整
   effect 窗口，应按掩码跳过对应目标，而不是删除整条轨迹。
   语义损失需要多个任务、每类多个样本；未知标签应显式屏蔽语义目标，
   仍可用于其他有监督/自监督目标。不要直接照搬发布模型的 500 步停止点。
   同样的 16 控制步在 Bridge 是 3.2 秒，在 RoboCasa 20 Hz 下只有 0.8 秒，
   因而时间尺度也需要适配。

4. **从训练轨迹导出 Bridge key/value。** CTE 有 128 维 `retrieval`
   和 256 维 `causal_interaction_state`。一种维度兼容的初始方案是分别对
   有效时刻池化并 L2 归一化，保存 key/value 和任务/episode ID；
   **这只是建议实现，不是已核实的官方导出池化公式**，需要检索验证。
   训练演示的完整轨迹可以用于离线建库，部署查询只能使用当时已知信息。
   特征生成时建议排除同一 episode 的自检索，并记录 CTE SHA256、
   输入契约、split 和数据版本，防止验证泄漏和自匹配捷径。

5. **适配 Stage-3 检索头。** 冻结策略 readout，从初始图像与指令提取
   4,096 维特征，训练 4,096→1,024→128 的 head 查询新 key 空间。
   利用 `bidirectional_supervised_contrastive_loss` 的同任务多正样本目标；
   batch 中必须有多个任务，否则单一任务的对比学习缺少负类。
   之前 159 条单任务集可以做行为过拟合测试，但不适合独立证明多任务检索有效。
   用未入库的验证轨迹报告 Recall@1/5、top-k 任务纯度、任务宏平均结果。
   readout checkpoint、提示模板和 CTE/bank/head 必须成套固定。

6. **重新生成 Stage-2 条件并训练/评测。** 给 prepare、server、cache contract
   三处增加统一的 CTE/bank/head 路径配置，生成新的缓存目录并严格校验。
   当前缓存只有旧库检索后的 global 和 phase/effect，并不含重建新库所需的
   完整 CTE key/value 或原始 VAE 潜变量，不能改名后充作 Bridge 库。
   新向量空间需要重新适配 Stage-2；不要悄悄给旧 290k checkpoint 换库后
   就认定它已经完成 Bridge 记忆训练。独立比较 bank-off、RoboCasa bank、
   Bridge bank，固定动作执行长度和去噪步数。

## 成本与建议顺序

首先完成当前 off/on × 4/16 步评测，再决定是否投入完整的 Stage-1/3 适配。
一个更快的诊断版是先用冻结发布 CTE 编码 Bridge 建库、仅适配检索头；
它仍含跨域 CTE 限制，应明确标为 provisional，不能当作完整适配结论。

检查时工作盘仅约 19 GiB 空闲。按当前每帧 `48×30×52`、492,160 个边界
帧估算，仅全量 FP16 VAE 缓存就约 **68.6 GiB**，尚未算训练中间量；
53,192 条 FP32 key/value 原始向量本身仅约 **78 MiB**。
因此应先做多任务子集/流式编码，或准备新磁盘；不能因为库文件小就直接
启动全量潜变量缓存，也不要自动删除现有 checkpoint 腾空间。

官方背景资料：[Zeva 项目](https://air-embodied-brain.github.io/Zeva/)、
[发布仓库](https://github.com/dsdfasfaac/zeva-test)。具体本地差异以以上代码、
已安装版本和实际 checkpoint 内容为准。