AURA: Action-Gated Memory for Robot Policies at Constant VRAM

Published 1 Jun 2026 in cs.AI, cs.AR, cs.DC, cs.PF, and cs.RO | (2606.02775v1)

Abstract: The KV-cache is the right memory for datacenters but the wrong memory for robots. Datacenter inference batches many short requests and resets them, amortizing an attention cache across a crowd. Embodied agents instead run one long, non-resetting episode on bandwidth-limited edge hardware, where high-bandwidth memory and flash are scarce, flash has finite write endurance, and memory writes rather than compute can become the binding constraint. AURA-Mem (Action-Utility Recurrent Adaptive Memory) targets this regime. It wraps a frozen vision-language-action backbone with a constant-size recurrent memory and a learned gate that writes only when the current observation would change the next action: memory that knows when to stay silent. Unlike reconstruction-based memory, the gate is trained directly against a closed-loop action-error signal. Its inference state is fixed at 4,224 bytes regardless of horizon, while a KV-cache grows to 6,061 times larger at 100,000 steps. On a controlled synthetic benchmark, AURA-Mem matches the best O(1) baseline in accuracy while using 5.19-6.13 times fewer writes, and up to 9.19 times fewer writes on easier configurations. Budget-matched random and periodic schedules do not recover this gain, isolating the benefit to the action-surprise signal. On a trained closed-loop OpenVLA-OFT 7B panel on LIBERO-Long (n=60 episodes per arm), the gate does not hurt success: AURA-Mem matches the ungated base policy (0.233) and slightly exceeds an always-write KV arm (0.217), while using 7.0 times fewer writes and constant memory. We also instantiate an approximate-information-state value-loss bound as a methodology demonstration; at this scale, the bound is vacuous rather than a guarantee.

Abstract PDF Upgrade to Chat

Authors (1)

Josef Chen

Summary

The paper introduces a constant VRAM memory architecture using an action-error gate to selectively update a fast-weight memory, reducing unnecessary writes.
It integrates a fixed Vision-Language-Action backbone with a recurrent fast-weight tensor to guarantee an O(1) inference memory footprint regardless of episode length.
Empirical results demonstrate up to 9.19× reduction in memory writes while maintaining state-of-the-art performance on both synthetic and real robotic policy tasks.

AURA: Action-Gated Constant Memory for Bandwidth-Efficient Robot Policies

Motivation and Problem Formulation

Physical robot deployments increasingly face severe memory constraints: unlike datacenter inference, where keyword-value (KV) caches are amortized and reset across many short batched requests, embodied agents operate in continuous, unbounded episodes where memory grows linearly with control steps. This growth saturates finite edge device memory and, more critically, memory bandwidth becomes the inference bottleneck. DRAM/Flash bandwidth is a systemic choke point, with high-bandwidth memory scarce and write endurance limited; every memory write directly impacts hardware lifetime and energy cost. Robots do not require full reconstruction of past states—only sufficient state to enable optimal actions.

AURA (Action-Utility Recurrent Adaptive Memory) directly addresses this challenge: it implements a constant-size fast-weight memory module, wrapped around a frozen Vision-Language-Action (VLA) backbone, with a learned action-error gate controlling memory writes. The gate operates on an action-prediction error signal and is explicitly trained against closed-loop action loss, aligning memory management strictly with utility for control, not generic reconstruction.

Figure 1: AURA system overview—action-driven write gate and bounded fast-weight memory encasing a frozen VLA policy backbone.

Architecture Overview

The AURA memory system comprises a single recurrent fast-weight tensor $W \in \mathbb{R}^{d_k \times d_v}$ , whose shape remains unchanged across episode length, guaranteeing O(1) inference-memory footprint. At each tick, observations are embedded using a fixed VLA model (OpenVLA-OFT 7B), projected to latent $z_t$ , and then processed through linear projections to yield query $q$ , key $k$ , and value $v$ . The read operation is $o_t = q_t^\top W_{t-1}$ , followed by a decision on whether to write, based on the magnitude of the reconstruction error $\|k_t^\top W_{t-1} - v_t\|_2^2$ ("surprise").

The write gate—a trainable MLP—is fed the current latent, prior memory read, and surprise scalar, outputting a soft write probability $p_t$ , binarized with a straight-through estimator for differentiability. Only when surprise exceeds a threshold does the memory update via a closed-form gradient step. Write frequency is regularized toward a target rate $\rho$ with a sparsity penalty, giving explicit control over DRAM/HBM traffic.

Figure 2: AURA per-step datapath—memory update is conditional on action-prediction surprise; O(1) VRAM is maintained at all steps.

Constant Memory vs. Growing KV-Cache

AURA's carried state is provably invariant to sequence length—unlike KV-caches that append new rows per step, escalating memory linearly and saturating hardware on long episodes. For $d_k = d_v = 32$ , the carried inference state is $z_t$ 0 bytes; a comparable KV-cache at 100,000 steps reaches $z_t$ 1 MB, $z_t$ 2 larger. This distinction is fundamental: only AURA provides robust guarantees on inference VRAM at any horizon.

Figure 3: Comparison of state growth—AURA’s constant-shape memory vs. linearly expanding KV-cache across episode length.

Formal Action-Sufficiency Bound and Empirical Certification

AURA instantiates the Approximate Information State (AIS) framework [subramanian2022ais], which provides a value-loss bound for compressed recurrent states in POMDPs:

$z_t$ 3

where $z_t$ 4 quantifies reward/action prediction loss and $z_t$ 5 characterizes transition mis-prediction, and $z_t$ 6 denotes the Minkowski functional of the surrogate value function. Empirically, AURA yields strong action-prediction sufficiency ( $z_t$ 7), but instantiated value-loss bounds are vacuous at current scale (e.g., $z_t$ 8 vs. trivial span $z_t$ 9), appearing as methodological rather than practical guarantees.

Experimental Evidence—Bandwidth and Memory Efficiency

On synthetic associative recall tasks with distractor streams, AURA achieves parity with the best O(1) baseline (fixed_size_state) in accuracy, while reducing memory writes by $q$ 0– $q$ 1 (up to $q$ 2 in closed-loop LIBERO-Long rollouts), and with a fixed inference-state VRAM (batch $q$ 3, $q$ 4 bytes). Random and periodic write schedules at matched bandwidth do not recover accuracy ( $q$ 5 vs. $q$ 6 at $q$ 7), confirming the gain derives from the action-surprise gating, not from the write budget alone.

Figure 4: Write-bandwidth vs. accuracy frontier; AURA yields 9.19 $q$ 8 write reduction at statistically indistinguishable performance.

Figure 5: Success rate vs. state budget $q$ 9; AURA matches fixed-size baselines at all budgets, outperforming naive and collapsed variants.

Figure 6: Accuracy vs. memory-footprint frontier; AURA achieves full performance with smallest, constant memory.

Comparisons and Deployment Regimes

AURA’s mechanism holds across batch deployment regimes: datacenter serving (batch- $k$ 0) amortizes KV-traffic, but robots (batch- $k$ 1) face unbounded state growth. The growing-KV side can match accuracy only by paying proportional VRAM and write-bandwidth costs. Real robot panel experiments show the gate mechanism maintains task success while suppressing unnecessary writes and memory footprint, substantiating O(1) memory claims under closed-loop, zero-shot rollouts.

Figure 7: Carried-state growth—AURA remains flat while KV-cache expands exponentially with episode horizon.

Figure 8: Batch- $k$ 2 vs batch- $k$ 3 deployment implications for memory footprint.

Figure 9: Closed-loop 3-arm panel—AURA matches base and KV performance with $k$ 4 fewer writes and invariant memory cost.

Ablations and Mechanism Insights

Gate selectivity experiments confirm robust alignment to salient events: write probability on decision-relevant steps ( $k$ 5) vastly exceeds non-informative distractor steps ( $k$ 6). Sweeping the write target $k$ 7 establishes that below the event density threshold, accuracy collapses, while moderate $k$ 8 captures essential information efficiently.

Figure 10: Action-error gate selectivity—high probability on salient steps, low on fillers.

Figure 11: Write-rate vs accuracy—sharp threshold near task’s information density.

Limitations

AURA’s claims are bounded in scope:

O(1)-VRAM applies strictly to carried inference state, not training activations.
Accuracy claim is parity with the best O(1) baseline; no superiority.
9.19 $k$ 9 gain measured on near-saturated tasks; hard-task results anchor bandwidth-accuracy tradeoffs.
Gradient-active parameter asymmetry exists due to a trainable gate-MLP.
AIS certificate is vacuous at deployed scale (methodology only).
Information-bottleneck contribution is not decisively isolated.
No real-robot, energy, or wall-clock latency measurements are reported.

Theoretical and Practical Implications

AURA advances efficient memory design for embodied agents by integrating action-aligned gating and constant-size memory. The established write-bandwidth frontier is hardware-relevant: direct reduction of memory writes translates to extended flash endurance, reduced DRAM/HBM traffic, and lowered deployment costs for robot policies scaling under real-world constraints. The action-sufficient compressed state principle unifies information representation with control utility.

On the theory side, empirical instantiation of the AIS bound is a step toward measurable certificates for recurrent policy memory, though practical guarantees remain a goal for future scale.

Conclusion

AURA establishes a solution to the memory wall in robot policy deployment by coupling a learned action-error gate with bounded fast-weight memory. It achieves write-bandwidth efficiency (up to $v$ 0 fewer writes), constant inference-state VRAM, and direct alignment of memory writes with control utility, outperforming naive write-triggers and competing O(1) baselines on both synthetic and real policy tasks. Although the theoretical certificate is loose at current scale and the accuracy claim is parity (not superiority), the practical implications for bandwidth-limited edge deployment are substantial. Future directions include hardware-level deployment, tighter certificate bounds, and capacity-controlled ablations for isolating gating effects. AURA sets the organizing principle for robot memory: write only what enables action.