MEMI: Reverse-Engineering What Robots Should Remember

2. Method

2.1 Memory modeling as an inverse problem

We consider imitation learning from a dataset of robot demonstrations \(\mathcal E\). Each episode contains a task instruction \(q\) and observation–action pairs \(\{(o_t,a_t)\}_{t=1}^{T}\). Depending on the setting, actions are either language subgoals executed by a fixed low-level policy [5] [6] or end-effector waypoint poses executed directly. In long-horizon tasks, the current observation is often not enough to choose the next action, which also depends on what happened earlier.

A common design summarizes the observation history in an explicit memory state \(m_t\). Here, we consider memory as field–value pairs, following structured state representations [7] [8]. In StopCube, for example, the textual field passage_count = 2 records that the cube has crossed the target twice, while a visual field holds keyframes selected from past frames.

Given memory annotations \(\tilde m_{1:T}\), our goal is to learn a policy that updates the memory and predicts the next action:

The policy \(\pi_\theta\) is instantiated as a reasoning model guided by a policy prompt \(\theta\). Its two outputs, the memory \(m_t\) and the action \(\hat a_t\), are supervised by the annotated memory \(\tilde m_t\) and the demonstrated action \(a_t\). During training, the policy conditions on the annotated memory \(\tilde m_{t-1}\) via teacher forcing, starting from the memory \(\tilde m_0=m_0\) at the initial time step. At deployment, the policy conditions on its own memory \(m_{t-1}\) instead.

With the annotations fixed, this is a forward problem (Figure 3, left): learning to maintain and use a prescribed memory to reproduce demonstrated decisions, as in much of the prior work [2] [3] [4]. Unlike actions, however, the memory behind each decision is never observed, so a prescribed annotation may miss what the policy needs to remember.

We thus explicitly treat the memory annotation as a latent variable, inferred by an annotator \(\alpha_\phi\) from complete demonstrations:

The annotator is also instantiated as a reasoning model, guided by an annotator prompt \(\phi\). This prompt specifies what to remember: which fields the memory contains and what each one records. Unlike the policy, the annotator sees each episode in hindsight, including future observations and actions. Optimizing \(\phi\) is thus an inverse problem (Figure 3, right): inferring what to remember from the decisions that memory must explain. This formulation draws inspiration from inverse optimal control [16] and inverse reinforcement learning [14], which infer a cost or reward that explains demonstrated behavior. We optimize the policy prompt \(\theta\) jointly with \(\phi\), so that the policy keeps pace as its memory supervision changes.

2.2 Three principles of effective memory

The inverse problem in Eq. (2) is inherently underdetermined: many memory annotations can explain the same decisions, but not all provide useful memory. Some omit information needed for decisions, include information the policy cannot obtain online, or retain fields that are never used. We address these issues through three principles of effective memory, each measured with a criterion for evaluating how well it is satisfied.

• Sufficiency. Memory should preserve the information needed to reproduce demonstrated decisions. In StopCube, passage_count and last_side alone are insufficient, because they cannot tell the policy when the cube will arrive. We measure sufficiency by the action discrepancy
  \[ \mathcal L_{\mathrm{suf}}(\phi,\theta) = \frac{1}{N} \sum_{t\in\mathcal E} \mathcal D_a\!\left(\hat a_t,a_t\right), \tag{3} \]
  where \(t\in\mathcal E\) ranges over all time steps of the demonstrations, \(N\) is their total number, and \(\mathcal D_a\) measures the discrepancy between predicted and demonstrated actions. An error points either to information missing from the memory or to a policy prompt that fails to use it.
• Predictability. Memory should be reliably reconstructible from experience. Because the annotator sees the future, it could write hindsight into memory, for instance by copying the next demonstrated action. Under teacher forcing, such a memory would make sufficiency trivial, yet the policy could never produce it at deployment. We therefore require the policy to reconstruct each annotation from what it has seen so far:
  \[ \mathcal L_{\mathrm{pred}}(\phi,\theta) = \frac{1}{N} \sum_{t\in\mathcal E} \mathcal D_m\!\left(m_t,\tilde m_t\right), \tag{4} \]
  where \(\mathcal D_m\) measures the discrepancy between memory states. A mismatch points either to an annotation that cannot be reconstructed online or to a policy prompt that does not explain how to update it.
• Necessity. Memory should exclude information irrelevant to downstream decisions. Extra fields add context that the policy must maintain and sift through without improving its decisions. Let \(\mathcal F_\phi\) be the set of fields specified by the annotator prompt. Along with its memory and action, the policy reports the fields \(U_t\subseteq\mathcal F_\phi\) that it uses for its decision at step \(t\). We penalize fields that are never reported as used:
  \[ \mathcal L_{\mathrm{nec}}(\phi,\theta) = \left| \mathcal F_\phi \setminus \bigcup_{t\in\mathcal E}U_t \right|. \tag{5} \]
  Reported usage serves as a proxy for a field’s contribution to decisions.

Together, the three criteria define a joint objective over the annotator and policy prompts:

The objective couples what the annotator asks the policy to remember with how the policy maintains and uses that memory.

2.3 Self-improving memory with textual feedback

Solving Eq. (6) amounts to searching over two text prompts. Its terms, however, are not differentiable with respect to these prompts, so gradient-based optimization does not directly apply. We thus optimize Eq. (6) with a textual analogue of stochastic gradient descent [9], in which textual critiques take the place of gradients. Each iteration has three stages (Figures 4 and 5). A forward pass annotates and evaluates a minibatch of episodes. A backward pass turns failures into textual critiques. A revision stage then updates both prompts. Besides the annotator and the policy, whose prompts are learned, MEMI uses reasoning models with fixed prompts for the remaining roles. An initializer first examines a few training demonstrations and writes the initial prompts \(\phi^{(0)}\) and \(\theta^{(0)}\).

Forward. At iteration \(k\), we sample a fresh minibatch of training episodes \(\mathcal B^{(k)}\). The annotator \(\alpha_{\phi^{(k)}}\) annotates memory on each episode. The policy \(\pi_{\theta^{(k)}}\) then predicts memory updates and actions at every step under teacher forcing, reporting which fields it uses. Each criterion in Eqs. (3)–(5) yields its own failures: steps with action errors, steps with memory mismatches, and fields that are never used.

1:def memi(E):                                    # E: episodes2:    annotator, policy = initialize(E)           # initial prompts3:    optimizer = TextualSGD([annotator.prompt, policy.prompt])4:    for batch in minibatches(E):                # fresh minibatch5:        label = annotator(batch)                # in hindsight6:        loss = evaluate(policy, batch, label)   # Eqs. (3)–(6)7:        loss.backward()                         # critique errors8:        optimizer.step()                        # revise prompts9:    return annotator, policy

1:def evaluate(policy, batch, label):2:    loss, used = 0, set()3:    for q, o_curr, a_curr, m̃_prev, m̃_curr in steps(batch, label):4:        m_curr, â_curr, U_curr = policy(q, o_curr, m̃_prev)5:        loss += dist(â_curr, a_curr)   # sufficiency6:        loss += dist(m_curr, m̃_curr)   # predictability7:        used |= U_curr8:    loss += len(fields(label) - used)  # necessity9:    return loss

Backward. Each failed step is first diagnosed on its own, attributing the fault to the memory annotation, the policy prompt, or both. These critiques are then merged across episodes into prompt-level critiques \(d\phi\) and \(d\theta\). In StopCube, for example, the merged critiques ask \(\phi\) to retain timestamped cube positions. They also ask \(\theta\) to estimate the cube’s speed from those positions before choosing when to press.

Revise. An optimizer updates both prompts from \(d\phi\) and \(d\theta\), conditioning on the history of earlier prompts and their feedback, to produce \(\phi^{(k+1)}\) and \(\theta^{(k+1)}\). The new annotator prompt changes the memory annotations of the next iteration. The new policy prompt changes how the policy maintains and uses them.

The entire loop runs offline on recorded demonstrations, without environment interaction. After \(K\) iterations, MEMI returns \(\phi^{(K)}\) and \(\theta^{(K)}\). The annotator is needed only during training. At deployment, the policy \(\pi_{\theta^{(K)}}\) starts from \(m_0\) and conditions on its own memory, using \(m_{t-1}\) in place of \(\tilde m_{t-1}\) in Eq. (1), and updates it as new observations arrive. Because the memory annotations are themselves learned, robot memory can self-improve beyond its initial supervision.

FORWARD · annotate and evaluate

1Sample episodes

\(o_t\)\(a_t\)\(t-1\)prepare\(t\)remain static\(t+1\)press⋮⋮

2annotator.forward

Annotator prompt ϕ⁽⁰⁾

\(o_t\)\(\tilde m_t\)passage_count=0last_side=-1passage_count=1last_side=1passage_count=1last_side=1⋮⋮

3policy.forward

Policy prompt θ⁽⁰⁾

at every timestep

\((q, o_{t-1}, \tilde m_{t-2})\)
\(\to (m_{t-1}, \hat{a}_{t-1}, U_{t-1})\)

\((q, o_t, \tilde m_{t-1})\)
\(\to (m_t, \hat{a}_t, U_t)\)

\((q, o_{t+1}, \tilde m_t)\)
\(\to (m_{t+1}, \hat{a}_{t+1}, U_{t+1})\)

4Loss

\(\mathcal L_{\mathrm{suf}}\) SufficiencyPredicted vs. recorded subgoal

\(\mathcal L_{\mathrm{pred}}\) PredictabilityPredicted vs. annotated memory

\(\mathcal L_{\mathrm{nec}}\) NecessityFields that are never used

Failed steps and memory-usage reports

BACKWARD · diagnose and revise

7optimizer.step

Revise ϕ and θ with the optimization history

and the current feedback

Revised ϕRevised θ

↺ ϕ goes back to step 2 and θ to step 3

6annotator.backward

Merge feedback across episodes

\(d\phi\)What to retain and when

\(d\theta\)Memory updates and subgoals

5policy.backward

Diagnose each failed prediction

\(d\tilde m\)Critiques of memory annotations

\(d\theta\)Faults in the policy prompt θ

0:00 Replay

3. Experiments

We evaluate MEMI on RoboMME [1], a simulated benchmark of tabletop manipulation tasks that a robot cannot solve from its current observation alone. RoboMME provides 100 expert demonstrations per task for training and 50 test episodes per task. MEMI optimizes its prompts only on the training demonstrations, over five iterations that each draw a fresh minibatch of five episodes. We evaluate the resulting prompts on the test episodes and report the success rate averaged over the 16 tasks.

The policy is one of four reasoning models: the frontier GPT-6 Astra and Claude Fable 5.1, and the cost-efficient GPT-6 Luna and Claude Haiku 5.5. Within each model family, the stronger model serves as the initializer, annotator, and optimizer: GPT-6 Astra for the Astra and Luna policies, and Claude Fable 5.1 for the Fable and Haiku policies. We compare three conditions for each policy:

• FullHist gives the policy the full observation history in place of memory, following recent methods that keep all past observations in the model’s context [10].
• Init keeps the initializer’s prompts fixed, serving as a fixed-annotation baseline.
• MEMI uses the policy prompts jointly optimized with the annotator prompts, so both what to remember and how to act are learned from the policy’s failures.

We consider two action types that differ in what the policy outputs, a language subgoal (Plan, Section 3.1) or an end-effector waypoint pose (Act, Section 3.2). We name each method by its condition, the reasoning model, and the action type, as [condition]-[model]-[action type], for example FullHist-Luna-Plan, Init-Astra-Plan, and MEMI-Astra-Act; only Astra runs both action types.

3.1 MEMI as a High-Level Planner

We first test whether MEMI helps policies of varying capacity learn what to remember and how to act. To this end, we run all four reasoning models as a high-level planner with the Plan action type. The policy receives textual observations listing each visible object and its position \([y,x]\) in the front camera image, extracted using the simulator’s semantic segmentation. We follow RoboMME’s grounded subgoal protocol, where the policy selects a language subgoal from a fixed list and, when needed, predicts the target object’s pixel position in the front camera image. RoboMME’s ground-truth low-level controller then executes the subgoal using its internal motion planner.

Across all four models, MEMI improves on both FullHist and Init (Figure 6). The initial memory already outperforms the full observation history by 1.1% to 15.9%, and MEMI adds a further 7.0% to 18.8% on top of it. The cost-efficient models gain the most: MEMI raises Luna-Plan from 66.6% to 85.4% and Haiku-Plan from 73.9% to 90.8%, on par with humans operating the same low-level controller (90.5%) [1]. The frontier models approach the ceiling, with Fable-Plan rising from 88.6% to 97.1% and Astra-Plan from 92.4% to 99.4%. This confirms that MEMI self-improves memory beyond its initial annotations and is effective across reasoning models of varying capacity.

3.2 MEMI as an End-to-End Policy

We next test whether MEMI still helps when the policy directly predicts memory and waypoint actions, without subgoal annotations or a low-level controller. To this end, we run GPT-6 Astra with the Act action type, following Embodied-Astra [10]. Astra serves as both the memory annotator and the end-to-end policy, each guided by its own prompt, and MEMI iteratively optimizes both prompts. The policy reads front and wrist camera images and the robot state, and predicts end-effector waypoints and gripper commands itself, with no low-level controller. Here the memory can also hold keyframes: fields that name earlier frames by labels such as E17, which the policy sees again for as long as their labels stay in memory. Because the training demonstrations record waypoints instead of subgoals, sufficiency compares the predicted waypoint with the recorded one.

Across the 16 tasks, MEMI raises Astra-Act from 88.3% to 97.0% (Figure 6), while FullHist-Astra-Act reaches 85.3%. Moreover, MEMI-Astra-Act exceeds the best published result on RoboMME among policies that predict their own low-level actions (92.1%) [10]. This shows that MEMI is agnostic to the action type and adds substantial gains even without a low-level controller.

3.3 Qualitative Analyses

We further examine how MEMI turns the policy’s failures into a better memory: which revisions change what the policy remembers, which change how it acts, and how each affects success. To this end, we trace the StopCube run of MEMI-Luna-Plan iteration by iteration (Figure 5). The initial memory only counts the cube’s passes and records the side it was last seen on, so Luna-Plan often presses the button too early. Iteration 2 adds the cube’s past positions to the memory, and success rises from 34% to 84%. Later iterations refine how the policy uses them, and success reaches 94%. This shows that what to remember and how to act emerge from successive iterations of policy feedback, not from any single iteration.

Beyond StopCube, the viewer below shows how the learned memory is used at test time. For each task, it plays one test episode for every policy under Init and MEMI, with ✓ or ✗ marking whether the rollout succeeds. Each video shows the camera views together with the memory the policy writes at every decision. With the Plan action type, the policy reads text observations, so this memory holds only text. The Act policy instead reads images, so its memory can also hold keyframes.

3.4 Ablation Studies

Finally, we examine the importance of each proposed principle for MEMI. To this end, we rerun the optimization with the feedback of one principle withheld at a time. Each withheld principle leaves a different failure in place (Figure 7). Without sufficiency, the memory never gains the cube’s past positions, and StopCube stays at 36% instead of 94%. Without predictability, the policy never learns to write the route correctly, and RouteStick falls to 18% instead of 84%. Without necessity, a field the policy never reads stays in the memory, so success on VideoPlaceOrder stays about the same while each episode costs 24% more. This shows that each principle addresses a distinct failure, so all three are needed.

4. Future Work

In this blog, we explore robot memory learning as an inverse problem: reasoning backward from demonstrated decisions to infer what a robot should remember. MEMI turns this perspective into a practical framework for improving memory through feedback. We believe this approach has broad potential, particularly for embodied foundation models in recursive self-improvement, where agents need to continually learn what to retain from experience. We hope this work inspires further exploration of inverse memory learning and its role in building more capable embodied agents.

Acknowledgments

This work is supported by an NSF CAREER Award and Volkswagen, with additional API credits provided by OpenAI and Anthropic.

Citation

If you find this work useful, please cite it as follows.

@misc{ryu2026memi,
  title        = {MEMI: Reverse-Engineering What Robots Should Remember},
  author       = {Ryu, Sooyoung and Dai, Yinpei and Du, Yilun and Finn, Chelsea and Liu, Yuejiang},
  year         = {2026},
  url          = {https://memi-robot.github.io}
}