REKI: Register Memory for History-Aware Robot Control
A robot action expert needs recent visual history without repeatedly processing old frames; REKI summarizes each frame in learned registers and carries their attention keys and values through a rolling cache.
Read the paper ↗What it is
- Register memory Learned tokens summarize each frame inside the action expert. Their attention keys and values are kept so later frames can consult that summary.
- Asymmetric attention Actions can read current and earlier registers, while registers do not read action tokens. That keeps the persistent memory tied to observations.
- Rolling cache Training uses short frame windows with a separate action target per frame. At inference, the oldest cached frame is dropped as a new frame arrives.
Why it matters
A robot can use recent visual history without repeatedly sending old frames through the encoder. The tradeoff is that later actions see a learned summary of the past, not the complete past video or earlier action tokens.
Before we get into the moving parts, here is the short version: REKI gives the action-generating part of a robot policy a small learned memory for each frame. It keeps those memory tokens around so the next decision can use recent visual history, without making the vision encoder repeatedly process a growing video.
-
History: A policy that only conditions on the current camera view can miss what changed between observations. Think of reaching for a cup: the latest image may show where it is, but recent frames can tell you whether it slid while the robot moved. The canonical example here is a sequence of manipulation frames from one episode; the expert needs a way to carry useful earlier visual context forward. A register is like a notebook entry for a frame, though unlike a human notebook it is a learned vector representation, not a readable record.
-
Compute: One option is to send a longer and longer sequence of images through the vision-language encoder. REKI instead leaves the encoder operating on one frame at a time and adds memory on the action-expert side. So the question is not simply how to remember more; it is where to put the memory so old observations do not require a new encoder pass.
-
Separation: The memory should describe observations, not the changing noisy action being generated. REKI uses learned register tokens alongside the current action tokens. Registers read the current frame’s encoded context and earlier registers, while action tokens can read registers and their own current chunk. In practical terms, actions get a lookup into recent frame summaries, but the persistent summaries do not absorb the action tokens. This resembles keeping observation notes separate from a draft action; the limit is that the learned notes may omit details needed later.
How does that memory get built? Each frame’s registers look at its encoded camera views and instruction, then at registers already in the history. The result is cached as keys and values, the pieces attention needs to retrieve information later. The important mechanism is that registers do not read action tokens and skip the adaptive normalisation used for actions. That keeps their representation independent of the current action noise and flow-time step, so the same frame registers can be reused while the action is refined.
Training and running the policy use that memory differently. In training, several frames from an episode are packed together, and every frame gets its own action target using context up to that frame. Earlier registers can receive learning signal from later frames that read them. At inference, the cache rolls forward: add the new frame’s registers and drop the oldest ones when the cache is full. The subtle bit is that a current register can summarize earlier registers, so information may continue through the chain even after the original frame has fallen out of the explicit cache. What gets lost when a learned summary stands in for the actual earlier frames?
Read the deep dive
Why People Wanted It
A frame-by-frame robot policy can see the current camera views and instruction, but a decision may depend on what happened just before: whether an object moved, whether a grasp was attempted, or where the arm was heading. Re-encoding a growing video history would make the vision-language side do more work at every step. So the real question is: can the action expert keep a compact trace of past frames, while leaving the encoder’s one-frame-at-a-time job intact?
The Core Move
Give the action expert learnable register tokens that summarize each frame, then retain their attention keys and values so later frames can read them as memory. The paper specifies the attention pattern directly; no additional formula is needed here.
Why Registers Can Carry Frame History
Now the version for readers who want more of the machinery: the registers are learnable tokens prepended to the current frame’s noisy action tokens in the expert stream. Each frame’s registers cross-attend that frame’s encoded prompt and camera views, and self-attend to the registers from earlier frames. Their keys and values persist in the expert’s cache, so a later frame can consult earlier summaries.
The key constraint is that registers do not attend to action tokens and skip adaptive normalisation. That makes them depend on observations rather than the current noisy action or flow-time-conditioned action computation. The registers can therefore be shared across denoising steps. The attention rule is asymmetric: register queries read the register history; action queries read the register history plus the current frame’s action tokens. This is what lets actions use past context without letting action tokens leak into the memory representation.
Why Actions Read Registers but Registers Avoid Actions
Here is the part that gets technical: the history enters the action velocity field through attention support alone. Registers of frame read registers from frames up to ; action tokens at frame read those registers and the action tokens of their own chunk. They do not read past action tokens. This separates persistent observation-derived memory from the current action-generation stream.
Each register also has a frame position and a within-frame slot position, represented by temporal and spatial axes of the VLM’s multimodal rotary embedding. The cache is thus not an undifferentiated bag: the expert receives the register sequence with positional information. What can this restricted attention pattern preserve, and what history does it necessarily leave out?
Why Window Training Can Feed a Rolling Cache
Now for the training and inference distinction: training samples a window of frames from one episode, and each frame is supervised on its own action chunk using context through that frame. Frames in a window are encoded together in a batched pass, without unrolling recurrence. Because later frame losses read earlier registers, those losses can send gradients back through the earlier registers.
At inference, the current frame’s registers are computed with the first action denoising step, then their keys and values are reused for the remaining steps. The cache keeps registers for the preceding frames in a rolling window and drops the oldest frame when a new one arrives. Unlike isolated training windows, inference registers may themselves carry information from earlier cached registers, allowing information to persist beyond the window. How much continuity you get therefore depends on the rolling chain of summaries, not on keeping every old frame explicitly.
Minimal Causal Code
This plain loop shows the state update: make the current frame’s registers from the current context and prior registers, then append them and trim the cache. It omits the paper’s batched window training and attention implementation.
# cache contains the previous frames' register keys and values.
cache = []
for context, noisy_actions in stream:
# The register path reads current context and earlier registers,
# but not the action tokens.
registers = make_registers(context, past_registers=cache)
# Actions read past/current registers and their own current tokens.
velocity = predict_actions(
context=context,
registers=cache + [registers],
noisy_actions=noisy_actions,
)
# Persist this frame's register keys and values for later frames.
cache.append(to_key_value(registers))
cache = cache[-(K - 1):]
What You Gain, and What You Lose
You gain a compact, reusable memory path: the encoder still processes one frame at a time, while the expert can condition current actions on cached register summaries. Registers can be reused across action denoising steps because they do not depend on flow time or state modulation.
You give up direct access to all past pixels and past action tokens. History is available only through what the registers encode and what the attention pattern permits. Training also uses finite windows processed in isolation, whereas inference can carry information forward through cached registers. This is a deliberate memory tradeoff, not a guarantee that a short register cache preserves every detail of the past.
REKI remembers frames by carrying forward what learned registers let the action expert read—not by replaying the whole past.