Memory Retrieval in Visuomotor Policies for Long-Horizon Robot Control

01 Motivation

在家庭这类部分可观测环境中运行的通用机器人，必须能回忆不再出现在当前感知输入中的信息——物体先前被放在哪、家电什么时候被打开、已经完成了哪些任务——才能完成长程任务。纯反应式（reactive）视觉运动策略在这类任务上必然失败。General-purpose robots operating in partially observable environments such as homes must be able to recall information that is no longer present in the current perceptual input — where an object was previously placed, when an appliance was turned on, which tasks have already been completed — in order to finish long-horizon tasks. Purely reactive visuomotor policies are bound to fail on such tasks.

“General-purpose robots operating in partially observable environments, such as homes, require memory to support autonomy.”

基于 attention 的 query-key 关联检索是一种通用、无需针对模态手工设计的记忆机制，但论文指出直接把长上下文 transformer 接入模仿学习会带来两个根本性挑战：Attention-based query-key associative retrieval is a general memory mechanism that needs no modality-specific hand design, but the paper points out that plugging a long-context transformer directly into imitation learning brings two fundamental challenges:

• 虚假相关（spurious correlations）：策略可能学会关注那些在训练数据中和专家动作相关、但并非任务因果相关的历史信息，这类相关性在测试时往往失效。Spurious correlations: the policy may learn to attend to history that correlates with expert actions in the training data but is not causally relevant to the task, and such correlations usually break down at test time.
• 误差累积（error accumulation）：离线模仿学习中的小预测误差会在闭环执行中与环境交互复合放大，噪声进一步污染存入记忆的表征，导致 model drift 和级联失败（cascading failures）。Error accumulation: small prediction errors from offline imitation learning compound through interaction with the environment during closed-loop execution, and the resulting noise further corrupts the representations written into memory, causing model drift and cascading failures.

02 Method

HALO 用 query-key-value 形式的注意力实现记忆检索：观测编码器 gθobs 和动作编码器 gθact 把历史中的图像、本体感知与动作映射到统一潜空间，策略骨干 fθ 中的检索模块对编码后的历史做注意力，再分别用动作头和文本头预测低层动作与 VQA 答案。为解决上述两个挑战，HALO 引入两个组件：HALO implements memory retrieval as query-key-value attention: an observation encoder gθobs and an action encoder gθact map the images, proprioception and actions in the history into a shared latent space; a retrieval module inside the policy backbone fθ attends over the encoded history, and an action head and a text head then predict low-level actions and VQA answers respectively. To address the two challenges above, HALO introduces two components:

1. 用 video question-answering 蒸馏 VLM 先验1. Distilling VLM priors with video question-answering

标准端到端模仿学习只用动作预测目标训练检索机制。HALO 额外用自动生成的 VQA 目标共同训练同一个策略骨干：从演示轨迹中生成探测任务相关信息的问题-答案对 (u, v)（例如物体位置、物体关系、事件顺序、子目标进度），策略需要检索到历史中支持答案的具体片段才能答对，这就为"应该检索什么"提供了直接监督。动作预测目标和 VQA 目标共享同一套编码器与骨干，仅预测头不同，因此 VQA 提供的先验会迁移到动作预测上：Standard end-to-end imitation learning trains the retrieval mechanism with the action prediction objective alone. HALO additionally co-trains the same policy backbone with an automatically generated VQA objective: question-answer pairs (u, v) that probe task-relevant information (object locations, object relations, event order, subgoal progress) are generated from demonstration trajectories, and the policy can answer correctly only if it retrieves the specific segment of history that supports the answer, which gives direct supervision for “what should be retrieved”. The action prediction objective and the VQA objective share the same encoders and backbone and differ only in the prediction head, so the priors supplied by VQA transfer to action prediction:

“This joint training ensures that VQA provides informative priors about what information from the history should be retrieved, while imitation learning grounds retrieval in expert action prediction.”

VQA 数据的生成是多阶段流水线（Fig. 4，正文 III-C）：(i) 用目标检测器 + VLM 把轨迹片段转成包含物体身份/位置与活动描述的文本摘要；(ii) 用 LLM 从摘要与任务指令生成聚焦特定帧的问题-答案对，并允许返回 "N/A" 以避免偏向轨迹早期信息；(iii) 用另一个 LLM 对每对 QA 按正确性与相关性打 0–5 分并过滤掉约 20% 低质量数据。论文报告该 stage-wise 流水线比直接单次视频提示生成 QA 的基线准确率高 60 个百分点，幻觉少 20 个百分点。The VQA data is produced by a multi-stage pipeline (Fig. 4, Sec. III-C): (i) an object detector plus a VLM turns trajectory segments into text summaries containing object identities and locations together with activity descriptions; (ii) an LLM generates question-answer pairs focused on specific frames from the summaries and the task instruction, and is allowed to return “N/A” so as not to bias toward the early part of the trajectory; (iii) another LLM scores every QA pair from 0 to 5 on correctness and relevance and filters out about 20% of the low-quality data. The paper reports that this stage-wise pipeline is 60 percentage points more accurate than a baseline that prompts the video once to generate QA directly, and hallucinates 20 percentage points less.

2. Top-K 稀疏注意力降低 model drift2. Top-K sparse attention reduces model drift

标准注意力对全部历史条目 {mi} 加权求和，即使无关信息也会获得非零权重。HALO 在每个时间步只保留 query-key 相似度最高的 k 个条目（Ik = TopK({si}, k)），把注意力权重严格限制在这个子集内，其余条目权重为零。由于 top-k 不可导，训练时使用 straight-through estimator：若被检索的元素有助于降低动作/VQA 损失，其 query-key 相似度会被强化，反之则在后续被淘汰的概率更高。Standard attention takes a weighted sum over all history entries {mi}, so even irrelevant information receives non-zero weight. At every timestep HALO keeps only the k entries with the highest query-key similarity (Ik = TopK({si}, k)) and confines the attention weights strictly to this subset, giving all remaining entries zero weight. Because top-k is not differentiable, training uses a straight-through estimator: if a retrieved element helps lower the action/VQA loss its query-key similarity is reinforced, otherwise it becomes more likely to be dropped later.

论文强调 HALO 把全部历史都存入记忆（不做压缩），只是在检索时做选择性稀疏化，因此相比 Scene Memory Transformer、Token Merging 等压缩式记忆，能保留低层控制所需的细粒度信息，同时通过稀疏化抑制噪声：The paper stresses that HALO stores the entire history in memory (without compression) and only performs selective sparsification at retrieval time, so compared with compressive memories such as Scene Memory Transformer and Token Merging it preserves the fine-grained detail that low-level control needs while suppressing noise through sparsification:

“HALO stores all information in memory, enabling access to fine-grained details and performing selective retrieval to reduce the impact of noise.”

03 Experiments

评测覆盖 REMEMBENCH 仿真长程操纵基准（4 类任务、每任务 50 次 rollout）与 5 个真实世界任务（跨固定臂与移动双平台，每任务 20 次 rollout，含一个人机协作场景），任务分别要求空间、关系、数量、事件时刻等不同类型的记忆。The evaluation covers the REMEMBENCH simulated long-horizon manipulation benchmark (4 task categories, 50 rollouts per task) and 5 real-world tasks (spanning a fixed arm and a mobile platform, 20 rollouts per task, including one human-robot collaboration scenario); the tasks require spatial, relational, numerical and event-time memory respectively.

仿真基准：REMEMBENCH（Table I，成功率）Simulation benchmark: REMEMBENCH (Table I, success rate)

HALO 平均成功率 41%，比 SAM2Act++（物体中心记忆）高 21 个百分点，比 ReMemBer（文本摘要）高 23 个百分点，比 Hand-Designed Features 高 12 个百分点，比 SMT 高 7 个百分点、比 Token Merging 高 12 个百分点。论文也如实指出："no single method wins across all tasks"——例如 Hand-Designed Features 在空间任务（0.68）上仍略优于 HALO（0.64），SAM2Act++ 擅长计数但在事件时刻任务上很差（0.04）。HALO reaches 41% average success, 21 points above SAM2Act++ (object-centric memory), 23 points above ReMemBer (text summaries), 12 points above Hand-Designed Features, 7 points above SMT and 12 points above Token Merging. The paper also notes plainly that “no single method wins across all tasks” — Hand-Designed Features, for instance, is still slightly better than HALO on the spatial task (0.68 against 0.64), and SAM2Act++ is strong at counting but very weak on the event-time task (0.04).

真实世界任务（Table II，20 次 rollout / 任务）Real-world tasks (Table II, 20 rollouts per task)

真实世界上 HALO 平均比标准 Transformer 高 19 个百分点，在需要存储多个物体、依赖人类同伴完成部分子任务、以及长达 8 分钟加热等信息密集任务上提升尤为明显。论文进一步在 "Retrieve Object" 任务上拆分了失败类型：相较全量注意力，HALO 把操纵失败降低 8%、记忆失败降低 25%（均为绝对值），支持"稀疏化降低 model drift、VQA 先验改善记忆检索"的假设。In the real world HALO beats the standard Transformer by 19 points on average, with especially clear gains on information-dense tasks such as storing multiple objects, relying on a human partner to finish part of the subtasks, and heating for as long as 8 minutes. The paper further breaks down the failure types on the “Retrieve Object” task: compared with full attention, HALO reduces manipulation failures by 8% and memory failures by 25% (both absolute), which supports the hypothesis that “sparsification reduces model drift and VQA priors improve memory retrieval”.

消融实验Ablations

• VLM 先验的作用（VQA supervision）：相较 HALO w/o VQA，加入 VQA 监督使平均成功率提升 10%（绝对值），在需要回忆物体放置或数量信息的任务上尤为明显。Effect of the VLM priors (VQA supervision): relative to HALO w/o VQA, adding VQA supervision raises average success by 10% (absolute), most clearly on tasks that require recalling where objects were placed or how many of them there were.
• Top-K 稀疏化的作用：相较全量注意力，Top-K 提升 9% 绝对性能，减少了决策中引入的无关/噪声信息。Effect of Top-K sparsification: relative to full attention, Top-K improves absolute performance by 9%, cutting the irrelevant and noisy information fed into the decision.
• 检索条目数 k 的选择：k = 8 时性能最佳（52% 成功率）；k = 4 会遗漏有用上下文（40%）；k 过大则引入无关历史（k = 12 为 48%，k = 16 为 33%，k = 24 为 44%）。论文将自适应地选择检索量列为未来工作方向。Choice of the number of retrieved entries k: k = 8 performs best (52% success); k = 4 misses useful context (40%); too large a k pulls in irrelevant history (k = 12 gives 48%, k = 16 gives 33%, k = 24 gives 44%). The paper lists adaptively choosing how much to retrieve as a direction for future work.
• Co-training vs. 两阶段训练：VQA 与动作预测联合训练达到 64% 成功率，优于"先做 VQA 预训练再做动作微调"（44%）和完全不用 VQA（42%），分别高出 20 与 22 个百分点——说明分阶段训练时 VQA 学到的知识会在微调阶段丢失，而联合训练能持续塑造检索行为。Co-training vs. two-stage training: training VQA jointly with action prediction reaches 64% success, better than “VQA pre-training followed by action fine-tuning” (44%) and than not using VQA at all (42%), by 20 and 22 points respectively — indicating that under staged training the knowledge learned from VQA is lost during fine-tuning, whereas co-training keeps shaping retrieval behavior.
• 替代稀疏化/记忆机制对比（Table III）：HALO（Retrieve Object 0.64 / Return to Container 0.32）优于 gated attention（0.45/0.28，最强替代方案，−11%）、hierarchical attention（0.28/0.22）、strided/window attention（0.20/0.28、0.13/0.16），以及压缩式非关联记忆 LSTM（0.14/0.12，−20%）、Mamba（0.20/0.18，−14%）、Transformer-XL（0.12/0.20，−12%）。固定访问模式的稀疏化方法容易错过发生在不可预测时刻的任务相关事件，而 Top-K 按内容相关性而非位置选择记忆条目。Comparison with alternative sparsification and memory mechanisms (Table III): HALO (Retrieve Object 0.64 / Return to Container 0.32) outperforms gated attention (0.45/0.28, the strongest alternative, −11%), hierarchical attention (0.28/0.22), strided and window attention (0.20/0.28 and 0.13/0.16), as well as the compressive non-associative memories LSTM (0.14/0.12, −20%), Mamba (0.20/0.18, −14%) and Transformer-XL (0.12/0.20, −12%). Sparsification schemes with a fixed access pattern easily miss task-relevant events that occur at unpredictable moments, whereas Top-K selects memory entries by content relevance rather than by position.

04 Limitations