AURA: 恒定VRAM下机器人策略的动作门控记忆

TL;DR AURA-Mem 为机器人 VLA 策略引入动作门控循环记忆，以恒定 4KB 推理状态和最高 9.19 倍更少写入，在长程任务中持平甚至超越 KV-cache 基线。

问题

问题背景

具身智能正推动视觉-语言-动作 (VLA) 模型从云端走向边缘机器人，这类模型需处理长达数万步的观测序列以实现闭环控制，对推理记忆机制提出严苛的资源约束。

现有方法局限

主流推理缓存 KV-cache 为数据中心批量短请求设计，在机器人单集长时运行中暴露出根本性缺陷：

• VRAM 无界增长：缓存规模随序列长度线性膨胀，而边缘硬件的高带宽内存 (HBM) 与闪存极其稀缺。
• 写操作成为瓶颈：持续写入不仅消耗带宽，更因闪存写入耐久度有限而直接威胁硬件寿命，写带宽而非计算可能成为主导约束。 现有 KV 压缩或驱逐方法多针对语言任务，以均匀或启发式规则丢弃历史，忽视动作闭环中信息的非均匀重要性；恒定内存替代方案 (如线性注意力、状态空间模型) 虽控制显存，却以重建损失训练，未直接优化下游动作影响，在关键观测上容易遗漏。

技术挑战与重要性

在电池供电的移动机器人上，VRAM 与功耗预算极度紧张。减少写操作不仅能延长存储寿命，还能显著降低发热与能耗。核心难点在于：记忆模块必须学会“何时值得写入”，即仅在当前观察会改变未来动作时才触发存储，这要求训练信号与闭环动作误差直接对齐。业界对将数十亿参数 VLA 模型压缩至常量 VRAM 并维持成功率高度关注，AURA-Mem 正是瞄准这一瓶颈。

行业类比

类似自动驾驶感知系统在车规级芯片上处理连续视频流，仅缓存与驾驶决策相关的关键片段，AURA-Mem 为机器人策略提供了“动作感知的选择性记忆”，杜绝无差别写入。

核心洞察

• Action-gated constant memory for embodied agents: Unlike datacenter-oriented KV-cache designs that amortize memory across many short requests, AURA-Mem targets a single long, non-resetting episode on edge hardware where VRAM and flash writes are the bottleneck. It maintains a fixed-size recurrent state (O(1) bytes) and uses a learned gate that writes only when an observation would change the next action. This directly reduces memory writes by up to 9.19× compared to always-write baselines, addressing flash write endurance and bandwidth limits critical for real-world robot deployment, while keeping inference VRAM constant regardless of horizon length.
• Action-error signal as the gating objective: Previous gating mechanisms often rely on reconstruction loss or handcrafted surprise metrics, which may store information irrelevant to the policy's task. AURA-Mem trains its write gate directly on a closed-loop action-error signal: the KL divergence between the compressed policy's action distribution and the full-context ungated policy. This makes the memory selective based on action utility, filtering out distractors that are surprising but do not alter behavior. Budget-matched random or periodic write schedules fail to recover this benefit, proving that the action-surprise signal is key to efficient memory usage without sacrificing policy performance.

方法

AURA-Mem 接受由视觉-语言-动作（VLA）骨干（如冻结的 OpenVLA-OFT）编码的连续观测序列作为输入，维持一个恒定大小的循环快速权重状态 W，并逐步骤输出动作。其核心创新在于一个可学习的动作-决策门控：每步先由骨干生成查询、键、值向量，门控根据当前观测计算一个动作-效用惊喜信号——即该观测对未来动作预测的影响程度。若惊喜值超过阈值，门控触发写入，用键值对更新 W；否则跳过写入，W 保持原状。之后从 W 读取上下文表示，与当前查询融合后预测动作。整个推理过程中，W 的字节数固定（4224 bytes），不随序列长度增长，实现 O(1) 推理状态 VRAM 占用。

训练时，门控与动作预测联合优化：采用动作-信息瓶颈（action-IB）目标，在最小化动作预测误差的同时，显式惩罚写入次数，从而迫使门控仅对动作有关键影响的观测做出反应。这种端到端训练直接针对下游任务（动作闭环）误差，而非传统的记忆重建或压缩指标。状态尺寸 N（W 的行数）作为弹性超参数，可在准确率与写入带宽之间进行帕累托权衡。

与现有方法的关键差异在于：AURA-Mem 的记忆写入决策由动作误差驱动，而非基于观察自重建或固定调度，因此能在不牺牲闭环成功率的前提下，将写入次数降至同类 O(1) 基线的 1/5 以下，且无需增长的内存或 KV 缓存。

实验

实验设计

实验在两个合成记忆基准（noisy_long_recall 和 sparse_recall）以及真实机器人操控任务 LIBERO-Long 上评估 AURA-Mem。合成任务测试模型在长序列中记住关键观测的能力，并通过“写入预算”与随机、周期性写入计划对比。LIBERO-Long 用 OpenVLA-OFT 7B 策略进行闭环多臂操控（n=60 回合），对比 always-write KV-cache、无门控基线和预算匹配的写入策略。指标包括任务成功率、内存写入次数、显存占用。

关键发现

• AURA-Mem 在所有任务上以极少的写入匹配或超越基线：合成任务写入减少 5.19–6.13 倍（简单配置下达 9.19 倍），LIBERO-Long 上写入减少 7.0 倍。
• 在 LIBERO-Long，门控策略成功率 0.233，与无门控版本持平，且略高于 always-write KV 的 0.217；恒定 4224 字节 推断状态显存，而 KV-cache 随步数膨胀（10 万步时达 6061 倍）。
• 门控信号来自 动作误差（action-surprise），而非输入重构，能够识别“改变未来行动”的关键观测；预算匹配的随机或周期性写入无法复现相同增益。

与基线对比的深度解读

传统 KV-cache 在数据中心按批次分摊显存，但在边缘机器人上，长程单会话导致显存线性增长、Flash 写入耐久度受限。AURA-Mem 通过恒定尺寸的递归记忆 + 学习门控，将写入约束转化为行动导向的选择：只写入对下游动作有影响的观测。实验表明，这种选择不仅节省带宽，还避免了冗余信息写入引发的策略干扰——always-write KV 因保留过多噪声记忆略降成功率。与 O(1) 基线相比，AURA-Mem 在相同内存限制下，通过门控实现投入更少写入仍维持准确，证明行动监督的门控是边缘部署的关键赋能技术。

行业影响

落地场景

AURA-Mem 瞄准具身智能边缘设备：工业机器人、仓储 AMR、巡检无人机、自动驾驶汽车等需要长时间、非重置运行的场景。这类设备通常搭载电池、低带宽内存和有限写入寿命的闪存，传统 KV-cache 随步长线性增长的内存和写入量成为瓶颈。AURA-Mem 的 4,224 字节恒定推理状态和最高 9.19 倍写入量缩减使其天然适配此类受限硬件。

商业价值

• 降本：大幅降低对高带宽内存（HBM）和大容量闪存的依赖，硬件物料成本减少；写入量骤降延长闪存寿命，降低维修更换频率，尤其适合大规模机器人车队运营。
• 增收与体验：恒定 VRAM 和低写入开销释放边缘算力，可支撑更长任务（如持续监控、长途递送），提升服务连续性和可用性；策略成功率持平甚至略优（OpenVLA-OFT 7B 面板上 0.233 持平基座模型），用户体验无损。

集成接口

AURA-Mem 作为轻量即插即用记忆模块，可嫁接于现有 VLA（视觉-语言-动作）骨干（如 OpenVLA）。训练时冻结骨干，仅优化循环状态和门控参数（+41.9% 梯度活跃参数），推理时完全替代 KV-cache 的存储与写入逻辑。部署时可与边缘推理引擎（如 TensorRT、ONNX Runtime）结合，通过定制内存分配器接管状态管理，无需重构模型管道。

具体落地 Use Case

• 仓储 AMR 自主导航：机器人需持续处理 3D 视觉+自然语言指令，闪存写入次数直接决定车载存储寿命。AURA-Mem 的 7.0 倍写入降低在 LIBERO-Long 长程任务中验证，可支撑上万小时运行，减少车队维护停机时间。
• 户外无人机光伏巡检：轻量边缘 NPU 上运行视觉检测模型，内存严格受限。AURA-Mem 常数 VRAM 占用和动作门控写入减少，能延长电池航时，同时保持缺陷识别准确率，提升巡检效率与 ROI。

局限

• **推理状态恒定，训练仍为 O(T)**：论文明确声明 O(1)-VRAM 仅适用于推理阶段前向状态，训练时仍依赖长度 T 的 BPTT，内存与时间线性增长。这意味着在长时程持续学习或在线微调场景中，边缘设备仍可能面临显存瓶颈，AURA-Mem 无法解决训练侧的内存问题。
• **动作门控对噪声敏感且验证场景有限**：门控信号完全依赖动作预测误差来训练，若动作标签存在噪声或任务中动作-观察耦合不紧密，门控可能学习到次优策略。实验仅在合成环境 noisy_long_recall 和 LIBERO-Long 上进行，未在真实机器人长时间操作或动态交互下测试，泛化至遮挡、移动操纵等复杂情形仍存疑。
• **梯度参数非对称对比，隔离性不足**：AURA-Mem 比定时写入基线多 41.9% 的梯度活跃参数，作者也承认信息瓶颈贡献仅“borderline positive”，未严格消融参数量的影响。这使得写入减少的收益是否完全源于门控机制，还是额外的容量提升，难以彻底分离。