PaperlayerView exact revision on arXiv

Source-linked AI summary

Divide-and-Remember: Recursive Action-Relevant Memory for Long-Horizon VLA Policies

Xuehui Yu, Eason Yu, Meiyi Wang, Haozhe Du, Stefano V. Albrecht, Harold Soh

arXiv:2610.00982v1cs.ROcs.AI

TL;DR

Long-horizon manipulation can require history because the current observation alone may not determine the action, while existing memory designs show inconsistent gains. D&R optimizes action-relevant memory through recursive token selection. It leads across RoboMME’s memory demands and attains the highest success rate in real-robot tests.

  • Problem

    Long-horizon manipulation can require history because the current observation alone may not determine the correct action, and existing memory designs show inconsistent gains across tasks.

    Source: §1 · pp. 1–2

  • Method

    D&R recursively selects history tokens with a lightweight selector, optimizing memory for conditional mutual information between the action and memory given the current observation.

    Source: §1 · p. 2

  • Results

    D&R leads across all four RoboMME memory demands, averaging 38.6% success versus 32.2% for the best latent memory, and achieves the highest success rate in real-robot tests.

    Source: §1 · pp. 2–3

  • Takeaways & Limitations

    A policy-sufficient memory preserves the optimal action distribution when conditioned on the current observation and memory instead of the full history.

    Source: §4 · p. 4

    D&R ties latent memory on state-like facts and falls behind on StopCube and ButtonUnmask, where direct history tokens are less compact than reformulated latent memory.

    Source: §7 · p. 9

Abstract

from arXiv · show

Vision-language-action (VLA) models struggle on history-dependent manipulation tasks, where the current observation alone does not determine the action, and the policy needs a memory of the history. Existing memory methods decide what to remember by design, for example, keeping frames with large pixel changes, and show inconsistent gains across tasks. We view what to remember as an optimisation problem. From the POMDP formulation of imitation learning, we show that the optimal memory maximises the conditional mutual information $I(a_t; m_t \mid o_t)$ between the action and the memory given the current observation. Intuitively, this means preserving the action-relevant information in the history that is not already contained in the current observation. Based on our analysis, we propose Divide-and-Remember (D&R), a recursive memory method that learns a memory function $m_t = M(h_t)$ and scales to long contexts while staying compute-light. It involves two strategies: (1) the selection over the full history is divided recursively into subproblems of top-$K$ selection over $2K$ tokens, so that fixed-size, lightweight selectors learned end-to-end support an unbounded history; (2) all recursion blocks share one selector, which captures the selection rule common to every block and keeps the method efficient. On RoboMME, a benchmark of 16 long-horizon manipulation tasks that require remembering when, where, what, and how to act, D&R achieves a state-of-the-art average success rate with consistent gains across all four suites under a budget of only 64 tokens; real-robot experiments show the same gain. Code, checkpoints and more results are at https://dnr-memory.github.io/

1 INTRODUCTION

The paper frames memory as preserving historical information needed to act when the current observation is insufficient. It proposes D&R, which recursively selects action-relevant history tokens, and reports gains on RoboMME and real-robot tasks.

1 INTRODUCTION

  • When a task’s memory length exceeds the policy’s context length, the same current observation can require different actions, so the history is needed to resolve the ambiguity.The introduction illustrates this with two demonstrations of different cube-moving manners followed by the same observation.Source: §1 · p. 1, Figure 1 · p. 2
  • Existing memory designs choose what to retain by different task-specific rules, and their gains are inconsistent across tasks.Latent representations track temporal facts but lose spatial detail, while perceptual snapshots preserve spatial facts yet can miss robot behavior and are not optimized for control.Source: §1 · pp. 1–2
  • The proposed objective is to maximize I(at; mt | ot), preserving historical information relevant to the action that the current observation does not contain.Pixel-change heuristics instead use a loose upper bound and can spend memory on changes unrelated to control, while latent memory optimizes the objective in a new latent space and loses spatial detail.Source: §1 · p. 2
  • D&R recursively selects K tokens from each 2K-token history slice, merges neighboring selections, and reapplies one lightweight selector trained end-to-end with the policy’s action loss.This composes selections over the full history while keeping the selector fixed-size and lightweight.Source: §1 · p. 2
  • Under a 64-token budget, D&R achieves a state-of-the-art average success rate of 38.6%, versus 32.2% for the best latent memory and 17.9% for the memoryless policy.On RoboMME, D&R also leads across every memory demand; real-robot experiments report the highest success rate.Source: §1 · pp. 2–3

2 RELATED WORK

Related approaches differ in what they store and how they select it, while classical POMDP solutions condition on either the full history or a belief state. The paper frames these designs with a shared information objective and presents D&R as a bounded, readable approximation.

2 RELATED WORK

  • Latent memory compresses history and tracks temporal facts but loses spatial detail; perceptual memory retains spatial and object detail but uses hand-designed selection, while parametric memory is slow to write and hard to read back exactly.These are the three memory categories described in the section.Source: §2 · p. 3
  • Conditioning on the full history makes the policy grow exponentially with history length, whereas recurrent belief-state updates grow linearly in compute and activation memory and cannot be parallelized over time.The section describes both routes as costly for long histories or repeated VLA forward passes.Source: §2 · p. 3
  • D&R approximates the shared objective I(at; mt | ot) in pretrained token space, using a bounded statistic whose retained facts can be read back exactly.The objective unifies the memory designs discussed in the paper.Source: §2 · p. 3

3 PRELIMINARIES

The preliminaries describe partially observable control and imitation learning from expert demonstrations. They define a memory problem by whether a task requires more history than the policy can access.

3 PRELIMINARIES

  • A POMDP has states, actions, observations, transition dynamics, and an observation function; the agent acts from partial or noisy observations, with the dynamics and expert policy unknown.The transition and observation functions describe how states change and how observations arise from states.Source: §3 · p. 3
  • In imitation learning, the policy is trained on expert observation–action demonstrations to maximize action log-likelihood under the expert data distribution.A memoryless visuomotor policy maps observations to actions; a VLA policy may also condition on a language instruction.Source: §3 · p. 3
  • A policy faces a memory problem exactly when task memory length exceeds its context length, leaving the input insufficient to determine the correct action.Task memory length is the shortest history window from which the expert policy can be realized.Source: §3 · pp. 3–4
  • In the memory-augmented graphical model, memory makes the action independent of the full history given the current observation and memory.The comparison is with imitation learning in a POMDP, where the action depends on history.Source: Figure 2 · p. 4

4 OPTIMAL MEMORY ANALYSIS

The analysis defines optimal memory as the history information about the action that remains after conditioning on the current observation. Such memory is policy-sufficient: the optimal policy can condition on it instead of the full history.

4 OPTIMAL MEMORY ANALYSIS

  • Maximising I(at; mt | ot) preserves the history information about the action that is not already supplied by the current observation.The analysis frames this as making the action independent of the remaining history once observation and memory are known.Source: §4 · p. 4
  • The optimal memory is a policy-sufficient statistic, making the optimal policy's action distribution identical whether it conditions on the full history or on memory plus the current observation.Equivalently, the action is independent of the remaining history given the observation and memory.Source: §4 · p. 4
  • Belief states may also predict reward and transitions, but requiring a history statistic to carry more information raises model, compute, and memory demands.The section contrasts action sufficiency with richer statistics that serve broader predictive purposes.Source: §4 · p. 4
  • Perceptual selection by frame difference maximises H(ok | ot), an upper bound on action-relevant information that can be loose for lighting changes or barely visible actions.The passage gives a lighting change and a button press as examples of the mismatch.Source: §4 · p. 4
  • The tractable objective trains policy and memory together through the negative action loss and a constant, without an auxiliary loss or hand-designed selection rule.The flow-matching loss is used as a surrogate for action log-likelihood; it trains the same conditional action distribution but does not bound its likelihood.Source: §4 · p. 5

5 DIVIDE-AND-REMEMBER

D&R selects a bounded subset of history tokens to approximate the action-relevant memory objective. Recursive fixed-size selection supports long histories and aligns training and deployment, while a shared selector is trained through the policy's action loss.

5 DIVIDE-AND-REMEMBER

  • D&R selects at most K tokens from the history in the pretrained VLA token space, turning memory construction into a discrete optimisation problem.The full-history mapping is trained end-to-end as an approximation because the objective is not available in closed form.Source: §5 · p. 5, §5.1 · p. 5

5.1 MEMORY AS BUDGETED TOKEN SELECTION

  • A fixed-size selector repeatedly keeps K of 2K tokens, composing these selections recursively so one lightweight selector can handle an unbounded history.The recursive construction addresses the growing-history input while keeping each selector's input size fixed.Source: §5.1 · p. 6
  • Because the recursion composes over token spans, training and deployment build the same tree for a given episode prefix, avoiding the offline–online mismatch.The split is left-aligned, and kept tokens are reselected at successive levels.Source: §5.1 · p. 6
  • Hierarchical filtering re-scores only tokens retained at the previous level, dropping low-importance information at the leaves and concentrating computation on survivors.The passage describes higher levels as holding increasingly important and sparse information.Source: §5.1 · p. 6

5.2 LEARNING THE SELECTOR

  • A transformer scores each token and a linear head produces selection scores; the selector uses a sigmoid relaxation for gradients and Gumbel noise for exploration during training.The threshold is set at the detached midpoint between the K-th and (K+1)-th largest scores.Source: §5.2 · p. 6
  • One selector is shared across recursion nodes, and its only training signal is the policy's flow-matching action loss on the selected tokens.At the root, the selection gates mask attention over the candidates passed to the policy.Source: §5.2 · p. 6

6 EXPERIMENTS

The experiments test whether D&R improves across tasks with different memory demands and examine where memory methods help. On RoboMME, methods are compared under a 64-token budget; ablations vary the candidate pool, budget, training objective, and candidate tokens.

6.2 EXPERIMENTAL RESULTS

  • 38.6% average success rate across four RoboMME suites, versus 32.2% for HAMLET, 27.9% for FrameSamp, and 17.9% for memoryless π0.5.D&R retrieves a grasp pose from the MoveCube video more than 300 steps earlier.Source: §6.2 · p. 7
  • D&R leads on frame-like facts, reaching 47.5% on Motion-Centric versus 31.0% for FrameSamp and 24.8% for HAMLET; latent memory is competitive on state-like facts.D&R reaches 31.7% on Long-Horizon Video, the only memory above the memoryless policy at 28.5%; HAMLET leads Time-Sensitive at 39.3% versus 31.3% for D&R and ties D&R on Dynamic Scene-Change and Event-Salient.Source: §6.2 · p. 8, Table 1 · p. 8
  • D&R achieves the highest success at essentially the same cost per policy query as memory budgets vary from 64 to 512 tokens.Training takes about 20 hours for D&R versus about 10 days for RB-VLA for 80k steps.Source: §6.2 · p. 8, Figure 6 · p. 8
  • D&R succeeds in 35 of 40 real-robot trials across four tasks, versus 22 for FrameSamp and 3 for π0.5.The largest margin is on PutBottles, the longest task, which lasts up to 4 minutes.Source: Table 2 · p. 9, §6.2 · p. 9

6.3 ABLATION

  • A larger candidate pool helps state-like facts but can hinder retrieval of frame-like facts; Motion-Centric peaks at 47.5% with 512 candidates and drops to 36.2% with 1024.At K = 64, Time-Sensitive rises from 20.7% to 32.4% as the pool grows from 128 to 1024 tokens. D&R leads at every tested memory budget, but reaches 50.1% at 256 tokens and dips to 46.4% at 512, which the authors attribute to convergence at 80k steps.Source: Figure 6 · p. 8, §6.3 · p. 9
  • Adding a future-prediction loss brings no gain, while adding instruction text tokens costs 6.6 points overall and 14.5 on imitation.Action tokens are excluded because the action-conditioned policy might learn to output the previous action, which is spuriously correlated with the next action.Source: §6.3 · p. 9

7 CONCLUSION, LIMITATIONS, AND FUTURE WORK

D&R achieves state-of-the-art success across different memory demands, with gains transferring to a real robot. Its direct use of history tokens is less effective on some state-like tasks, while shared multi-task training may leave redundancy and specialization trade-offs.

7 CONCLUSION, LIMITATIONS, AND FUTURE WORK

  • D&R achieves state-of-the-art success rates across different memory demands, and its gain transfers to a real robot.The conclusion describes gains across different memory demands and reports transfer to a real robot.Source: §7 · p. 9
  • D&R ties latent memory on state-like facts and falls behind on StopCube and ButtonUnmask, where latent memory can reformulate or aggregate knowledge more compactly.The authors identify D&R’s direct storage and use of history tokens as the reason for this difference.Source: §7 · p. 9
  • Training all memory methods in the same task-agnostic multi-task setting may lead them to acquire shared knowledge redundantly, leaving room to improve parameter efficiency while retaining specialization.The authors frame specialization around focusing on distinctive aspects of tasks.Source: §7 · p. 9

AI USE STATEMENT

The appendix states the authors’ limited uses of generative AI, then reviews POMDP history representations: full-history policies scale exponentially, while belief states are sufficient but can grow with the state space. It discusses compact representations and policy similarity as alternatives to reward-based sufficiency measures.

AI USE STATEMENT

  • Generative AI tools were used for manuscript editing and checking, literature search, and code refinement, but not for research ideation, method implementation, or experimental data or results.The authors say they reviewed the AI-assisted work and take responsibility for the final content.Source: AI USE STATEMENT · p. 10

A FROM FIRST PRINCIPLES OF POMDP THEORY

  • Conditioning on the full observation–action history yields a policy whose size grows exponentially with time, so it is practical only for short histories.The authors therefore seek history representations that scale to high-dimensional VLA observations and large models.Source: A FROM FIRST PRINCIPLES OF POMDP THEORY · p. 12

A.1 BELIEF

  • A belief state is sufficient for the state posterior, but its dimension grows with the number of states, motivating a more compact representation for large or continuous state spaces.The belief state can be updated recursively from the previous belief, action, and new observation.Source: §A.1 · p. 13
  • Bisimulation measures reward and transition similarity, whereas policy similarity substitutes policy differences for reward differences while retaining transition sufficiency.Reward grounding can be too restrictive or permissive for policy similarity, and imitation learning has no reward available.Source: §A.1 · p. 14

A.2 MEMORY

The memory formulation seeks a history representation sufficient for the expert action given the current observation, optimizing action-relevant information rather than requiring prediction of future observations. The appendix also documents the shared training setup and D&R variants, including a world-model prediction objective and pool-size and ablation sweeps.

A.2 MEMORY

  • A memory is a deterministic function of history that makes the expert action independent of the remaining history given the current observation and memory.Unlike a belief state, it need not predict the next observation.Source: §A.2 · p. 14
  • The optimal memory maximizes the action-relevant information retained from history beyond the current observation, while discarding the least relevant information.A memory maximizing this objective is policy-sufficient: the expert action distribution is unchanged when history is replaced by memory.Source: §A.2 · p. 14
  • D&R keeps policy sufficiency alone under a limited budget, whereas the figure contrasts it with methods that also measure reward or transition sufficiency.The figure identifies bisimulation as measuring reward and transition, policy similarity as measuring policy and transition, and belief as covering all three.Source: Figure 8 · p. 14
  • Policy and memory can be trained end-to-end with a single action loss, without an auxiliary objective or supervision specifying what to remember.The objective is optimized through a variational lower bound using the policy in place of the true demonstration action distribution.Source: §A.2 · p. 15

B.1 EXPERIMENTAL SETUP

  • The shared training recipe uses the released π0.5 configuration, freezes the SigLIP vision backbone, caches visual tokens, and trains models for 80K steps.Training times vary by method, from about 20 hours for perceptual and latent variants to 10 days for RB-VLA.Source: §B.1 · p. 16

B.2 VARIANTS OF D&R

  • D&R+pred adds a world-model loss to the action loss while leaving the selector, reduction tree, and policy unchanged, making memory features predictive of future latents as well as sufficient for action.The variant omits inverse-dynamics and KL terms because the static candidate pool lacks per-step actions and a stochastic belief.Source: §B.2 · p. 16
  • Table 4 reports D&R success rate (%) averaged over four tasks per suite as mean ± standard deviation across evaluation runs (checkpoints × seeds 0/7/42), sweeping |H| from 8 to 64 frames at K = 64 tokens.Table 6 ablates the training signal and candidate set at |H| = 256 against D&R at the same pool size.Source: §B.2 · p. 16, Table 4 · p. 17

B.3.1 FULL RESULTS ON ALL TASKS

The full RoboMME results are reported as success rates under a 64-token memory budget, with suite-level and per-task breakdowns. The four suites represent counting, tracking hidden objects, referring to earlier moments, and reproducing demonstrated sequences.

B.3.1 FULL RESULTS ON ALL TASKS

  • The four suites test accumulating event counts, tracking an object after it is hidden, referring back to an earlier episode moment, and reproducing a demonstrated sequence.The passage says the results are reported by suite and by task.Source: §B.3.1 · p. 17
  • Table 7 reports every method’s success rate under the 64-token budget, with suite averages, an average across all 16 tasks, and per-task means and standard deviations over seeds.The caption identifies the π0.5 row as the number reported by Dai et al. (2026).Source: Table 7 · p. 18

B.3.2 ABLATION: MEMORY BUDGET

The appendix reports memory-budget comparisons on RoboMME and extends memory visualizations across all tasks.

B.3.2 ABLATION: MEMORY BUDGET

  • Table 8 reports average success rates overall and by suite as the memory budget grows from 64 to 512 tokens.Some FrameSamp and TokenDrop entries are reported by Dai et al. (2026); the remaining entries are the authors’ runs.Source: Table 8 · p. 19, §B.3.2 · p. 19
  • Table 9 gives per-task RoboMME success rates at a 512-token budget, with results grouped by the memory each method stores.The π0.5, FrameSamp, and TTT rows are reported by Dai et al. (2026); bold marks the highest rate in each column.Source: Table 9 · p. 19, §B.3.2 · p. 19

B.3.3 MEMORY VISUALISATION ON ALL TASKS

B.4 REAL ROBOT EXPERIMENT

The real-robot evaluation uses four tasks on a Franka Panda, with 10 rollouts per task. D&R has few reported failures across the tasks, while failure reviews describe distinct errors for the baselines.

B.4 REAL ROBOT EXPERIMENT

  • The four tasks test counting with human interventions, locating hidden cubes, recalling picked cubes, and reproducing demonstrated paths.PutBottles, TrackCube, RepickCube, and DrawPattern mirror BinFill, VideoUnmask, VideoRepick, and PatternLock, respectively.Source: §B.4 · p. 26
  • Figure 25 separates earlier-history frames from the front-camera execution observation for each real-world task.For PutBottles, earlier history includes episode frames with a human intervention.Source: Figure 25 · p. 26
  • The setup uses a 7-DoF Franka Panda, two RealSense cameras, 100 teleoperated demonstrations per task, and 10 evaluation rollouts per method.The memory uses the front view, while the wrist view is the current observation; real-robot memory budgets are larger than in simulation because episodes are longer.Source: §B.4 · p. 26
  • Across 10 rollouts per task, D&R had one PutBottles failure, lifted the correct TrackCube cup in all ten trials (one logged failure followed an early stop), one RepickCube failure, and two DrawPattern failures.The DrawPattern failures involved a wrong movement direction.Source: §B.4 · pp. 26–27
  • For PutBottles, π0.5 often miscounted and FrameSamp sometimes remembered the wrong count; for TrackCube, both baselines had wrong-cup failures, while D&R chose correctly in all ten trials.One TrackCube D&R rollout was logged as a failure because it stopped before the automatic put-back.Source: §B.4 · p. 26

Related papers

Loading 2610.00982v1…