Memo: Training Memory-Efficient Embodied Agents with Reinforcement Learning
Authors: Gunshi Gupta
Karmesh Yadav
Zsolt Kira
Yarin Gal
Rahaf Aljundi
Venue: NeurIPS 2025 spotlight ()
Keywords: Reinforcement learning
transformers
memory
compression
Primary area: reinforcement_learning
TL;DR: We propose a way to create and access episodic memory when training transformer policies with RL on long horizon tasks that require remembering things from the past.

Abstract: To enable embodied agents to operate effectively over extended timeframes, it is crucial to develop models that form and access memories to stay contextualized in their environment. In the current paradigm of training transformer-based policies for embodied sequential decision-making tasks, visual inputs often overwhelm the context limits of transformers, while humans can maintain and utilize a lifetime of experience compressed as memories. Significant compression is possible in principle, as much of the input is irrelevant and can be abstracted. However, existing approaches predominantly focus on either recurrent models with fixed-size memory or transformers with full-context reliance. In this work, we propose Memo, a transformer-based architecture and training recipe for reinforcement learning (RL) on memory-intensive, long-horizon tasks. Memo incorporates the creation and retrieval of memory by interleaving periodic summarization tokens with the inputs of a model during training. We demonstrate Memo’s effectiveness on a grid-world meta-RL benchmark and a multi-object navigation task in photo-realistic indoor settings. Memo outperforms naive long-context transformer baselines while being more compute and storage efficient. Additionally, Memo generalizes better to longer contexts at inference time and remains robust in streaming settings, where historical context must be truncated to fit inference constraints.

Memo: Training Memory-Effciien Embodied Agents with Reinforcement Learning
Gunshi Gupta' University of Oxford
Karmesh Yadav Georgia Tech University
Zsolt Kira Georgia Tech University
Yarin Gal University of Oxford
Rahaf Aljundi Toyota Motor Europe
Abstract
To enable embodied agents operate effectively over extended timeframes, it is crucial to develop models that form and access memories to stay contextualized in their environment. In the current paradigm training transformer- based policies for embodied sequential decision-makinn tasks, visual inputs often overwhelm the context limits of transformers, while humans can maintain and utilize a lifetime of experience compressed as memories. Significant compression is possible in principle, as much of the input irrelevant and can be abstracted. However, existing approaches predominantly focus on either recurrent models with fixed-size memory or transformers with full-context reliance. In this work, we propose Memo, transformer-basee architecture and training recipe for reinforcement learning tasks. Memo incorporates the creation
(RL) on memory- cennoor-inteesvve,, and retrieval of memory by interleaving periodic summarization tokens with the inputs of a model during training. We demonstrate Memo's ' effectiveness on a grid- world meta- RL benchmark and multi-object navigation task in photo-realistic indoor settings. Memo outperforms naive long-context transformer haselines while being compute and storage efficient. Additionally, Memo generalizes better to longer contexts at inference time and remains robust in streaming settings, where historical context must be truncated to fit inference constraints.
Introduction
Humans intuitively prioritize and retain memories relevant to their current tasks, filtering out inrele- vant details such as the color house's walls along route unless it serves as a crucial navigation landmark. Similarly, reinforcement learning (RL) agents must learn to selectively retain and access task-specific memories guided by task objectives rather than predefined rules. This ability is espe- cially critical for long-horizon tasks, where delayed rewards make credit assignment challenging and necessitate efficient training and inference over extended timescales to effectively capture the relationship between past and future events.
Recently, Transformers have become the go-to framework for sequence modeling due to their flexible and expressive encoding, free from the fixed-size constraints of recurrent neural networks (RNNs) 14. However, it assumes access to all past information at each step rather than selectively retaining relevant memories, which lead to inefficiencies and make it harder to extract relevaelevnnt information for decision decision-making, Their quadratic attention complexity further limits scalability, making it challenging to process long contexts efficiently, especially when gradients must propagate over large sequences. At test time, Transformers face another practical challenge: storing an
"Correspondence to: mamshi.guptade ox. ox.ac.u.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

increasingly large key-value (KV) cache is computationally expensive, and attending over ever- growing contexts requires the ability to generalize beyond the temporal patterns seen during training.
To address these challenges, we propose Memo, novel framework that enhances transformer-based agents by enabling them to summarize and index past experiences through specialized summary tokens. Memo introduces a simple yet effective training procedure where transformers learn to compress their experience by periodically generating summary tokens that encode relevant past information. These summary tokens are stored in a dedicated memory buffer, allowing the model to attend condensed representations of prior states instead of maintaining a full-context cache. This enables Memo to maintain compact memory footprint at inference, reducing computational costs while preserving long monizonn reasoning capabilities.
Inspired by advancements in extending context lengths in language models, Memo adapts these principles reinforcement learning (RI.), addressing the unique challenges presented by RL tasks. Specifically, Memo integrates with both on-policy and off-policy RL agents, leveraging learned compression mechanisms enhance sample efficiency and generalization in sequential decision- making settings. We demonstrate its effectiveness across diverse tasks, including a grid world and a 3D indoor navigation task in Habitat. Memo matches or outperforms the naive transformer baseline, which requires storing its entire context, leading to an 8-10x larger cache, Despite using significantly less memory, Memo achieves better in-context learning (ICL) behavior, higher success rates, and shorter path lengths in unseen environments. On targe-scale navigation tasks, we show that Memo also exhibits superior robustness in streaming settings where the cached context is truncated, demonstrating improved adaptability under limited inference budgets.
Our contributions are summarized follows:
We propose Memo, framework that enables transformers to learn to summarize and store task task-relevant past experiences, reducing computational costs while maintaining long-horizon reasoning capabilities.
We evaluate Memo on sequential decision decision-making tasks, demonstrating its efficiency and im- proved in-context learning behavior compared to transformers that require full-context storage.
We show that Memo is general and versatile memory augmentation technique applicable to both on-policy and off-policy RL agents for long-horizon tasks.
2
Related Work
In this section briefly review research on extending context limits of transformers in language modeling, as well as the state of the art in long sequence modeling and decision making in RL.
Extending Transformer Contexts in language: A significant limitation of current large language models is their restricted input context length, prompting research into methods for scaling and generalizing larger contexts [33 1 27 Key-value (KV) caching can reduce recomputation at training 99] or inference time [23), by storing and reusing self-attention outputs during decoding. 25 further reduce memory by maintaining compressed cache, though this compression is not completely task-guided.
Beyond compute efficiency, recent work explores dynamic context extension. Recurrent Memory Transformer (RMT) 4 introduces summarization tokens to periodically compress and propagate prior context while discarding older tokens. Autocompressors (AC) [ extend this by accumulating generated summaries across context windows, avoiding RMT's fixed-size memory but still truncating gradient propagation through summaries. We extend this context summarization approach to the more involved setting of reinforcement learning (RL), where summarization must support decision decision-making and handle credit assignment over long horizons. Unlike AC, which fine-tunes pretrained models on a supervised token prediction task with limited gradient propagation- aims to match but does not surpass full-context transformer baselines-ww train from scratch, integrating summarization directly into the RL optimization process and propagating gradients across all summaries. This enables more effective task-drived memory formation, surpassing full-context transformer baselines in both efficiency and scalability.
Transformers in RL: Transformers, while widely used in self-supervised language modeling 10 3 have seen slower adoption in reinforcement learning (RI.) due challenges like sparse rewards, optimization instability, and limited batch sizes in on-policy RL. However, move to more

---

Toher's
Trans Corner
Teans for 1e1
Trans former
O
sisuel Erooder
Infommtionn
Figure I: Architecture Diagram of Memo. The frames -1-3 depict the input sensor observations
to the agent at timesteps - 31seg- The figure depicts the information flow between three consecutive segments of a much longer context of inputs provided to a transformer during training. The summary tokens bottleneck the information passed from one chunk of inputs the next.
complex tasks that demand long-horizon planning and reasoning over past experiences, transformers have shown promising initial results, outperforming recurrent models in POMDPs when encoding longer temporal dependencies [22]
To stabilize transformer training in partially observable RL, Parisotto et al. [2 used KV caching
while restricting gradient propagation to small window of time steps, limiting the learning of long- term dependencies. More recently, ReLIC [1]] improved convergence of on-policy RL algorithms by unfreezing the historical context within a rollout and applying frequent updates, while AMAGO [12] adapted off-policy RL to train with shared transformer backbone. While these methods enhance long-horizon learning, ReLIC's high computational cost and slow training hinder scalability, and AMAGO is limited small state-space observations, reducing its applicability to more complex tasks. Our approach addresses these challenges by enabling efficient memory utilization and leveraging a more scalable transformer architecture for extreme long-houe hong-horizon decision decision-making.
Methodology
This section outlines our methodology for ceemory-augmentdd in-context reinforcement learning (RL). We first formalize the problem and describe the in-context RL algorithms and the sequence model consider, in Section 3.1 We then introduce context summarization mechanism in Section [3.2] and describe key implementation details including the attention masking, positional encoding, segment randomization and cache management scheme we use.
3.1 Problem Setting
In this work, we study tasks that require a model to perform sequential Decision-making over long executions while leveraging its experience history to improve efficiency over time. Many diverse reasoning problems fall into this category, ranging from embodied navigation to language modeling. We formalize these tasks within the fraaework of in-context reinforcement learning (RL), where models adapt and improve their performance using past inferences rather than pre-collected supervised datasets. By treating these tasks as partially observable Markov decision processes (POMDPs) and transforming them into fully observable MDPs through the integration historical context, is possible to train policies effectively using both on-policy and off- oof-po RL methods.
We use a sequence encoder model the agent's 's policy, processing past observations into a contextual representation for decision-making At each timestep I, the sequence encoder receives a sequence of observations, denoted as X, {o1..... ).),, and maps them to hidden representations he = SeqEnc(X, The policy distribution and value estimate are then computed using learnable actor and critic heads as T [aa ht), V (hr), respectively, where a₂ E AA is the action taken by the agent at timestep t, selected according to the policy over the action space A.
In-Context RL Algorithms: In in-context RL, an agent adapts over time by conditioning on its entire past within a trial. A trial consists of multiple episodes, where each episode resets upon termination, but the agent retains memory across episodes. These episodes share a common underlying as environment dynamics, task structure, or goal distribution eeriiibtiooneenalinngg enabling the model to leverage past experiences for more efficient adaptation. Sequence models like Transformers encode history by attending to all previous timesteps within trial, allowing the agent to infer temporal dependencies and adjust its behavior accordingly. However, as trial lengths increase, attending to all past timesteps becomes computationally infeasible.
3

---

To address this, we introduce Memo, a method for constructing and accessing memory representations in Transformers improve long-term context management in RL, as described in Section 3.2] To demonstrate its efficacy and broad applicability to long-conters RI., we integrate Memo with both on-policy and off off-policy RI. method:
On-policy RL: We incorporate our memory mechanism into RELIC M. an adaptation of DD- PPO [00,, which applies frequent partial updates during trajectory rollouts help transformer policies learn more effectively from in- in-context experiences.
Off-policy RL: We extend Memo to AMAGO [12), which trains transformer-basee policy using
shared sequence encoder and unified actor-critic loss to improve stability and learning efficiency.
3.2 Context Summarization in Memo
we introduce context summarization learned sub-task that enables the transformer encoder to compress and retain task-relevan information from long observation histories. This mechanism allows the model to efficiently process past experiences while optimizing for task rewards in RL. Our approach uses learnable summary embeddings to prompt the transformer to generate summary tokens at predefined intervals, ensuring efficient memory utilization in long-horizon tasks.
Following 77
Instead of attending to the full sequence history, the transformer periodically compresses past context into summary tokens, which are then fed back into the model in future timesteps. We partition long input sequences into segments of length (seg. and generate sam summary tokens at the end of each segment. These tokens act as compact memory representations, allowing the model to condition on past experiences without requiring access to raw observations. The encoding and integration of summary tokens into future segments are illustrated in Figure [I] and the pseudocoe is provided in Appendix A.4 The summarization mechanism is trained end-to-en through the RL objective, allowing the model to attend over and refine all previously generated summary tokens. Gradients propagate through the attention mechanism, ensuring that memory updates remain task-driven.
We use causal masking which includes all the previous summary tokens as well as the previous observations within the segment being processed at time step t. This excludes any previous observation that contributed to and that was processed before the most recent summarization. This creates a bottleneck for the flow of historical information to be solely through summary tokens.
Attention Masking:
To add positional awareness, assign indices an observation at time step t based on its relative position within a segment. We use the same positional encoding method (linear or ROPE) used by the baseline implementation which we will be comparing to. The position labeling scheme in a naive full-context transformer input has indices which run from 0 to the length of the context window. In Memo, the input context at time step t in the rollout consists of n Iyum summary tokens and the latest segment's 's observations, where n = [t/laegl- The position indices thus range from 0 to n * sum 1 for the summaries and then start from n (sam and go to / - n * (seg + + aawm for the most recent observations in the latest segment. The summary embeddings at the end of a segment are assigned indices following those of the last segment observation.
Positional Encoding:
Segment length randomization: Following Chevalier et al. 77. we randomize segment lengths during training by sampling uniformly within +20% of a fixed leeg. This approach, which previously showed small improvements in perplexity for token prediction, will be ablated in (.8 A fixed (eg is used during data collection and evaluation, while training uses randomized segment leggttteengths..
Maintaining the KV Cache: on-policy RL, the KV cache stored during action selection can become stale after a model update, as the new weights would encode different representations for the same past context. This divergence is problematic since policy decisions depend on context, making consistency crucial. ReLIC addresses this by refreshing and re-encoding the KV cache with the latest model weights after every on-policy update. We extend this approach to memory tokens, ensuring consistency by recomputing the summary vectors alongside the KV cache refresh. This guarantees that both cached representations and learned memory remain aligned with the current policy.
4 Experiments
In this section, we present experimental results highlighting different aspects of our proposed method. We first describe the benchmark tasks and baselines used in our experiments, followed by different experimental insights in the following subsections.

---

4.1 Benchmarks
The Extended Object Navigation (EXTOBINAV) task, first introduced in builds on the OBJECTNAV task commonly used in embodied AI research [ 20). While OBJECTNAV requires an agent to navigate to a single object goal in an episode, EXTOBINAV focuses on how performance evolves when the agent is repeatedly tasked with reaching different object goals within the same scene. The task evaluates the agent's ability to leverage memory-utilzzing information from prior exploration during navigation to reach subsequent goals.
EXTOBJNAV:
The EXTOBINAV task uses 37 training and 12 validation scenes from HSSD [16] and includes 20 object instances from the YCB dataset [6]. These objects can be randomly placed on receptacles throughout the scene, with each placement featuring an average of 30 objects. We use 11,100 novel placements during training and 108 during evaluation. A sample scene is shown in Figure 2a
Agents are trained and evaluated over trials, each consisting of sequence of episodes with the same object placement in a scene. Each episode mirrors OBJECTÑAV in structure: the agent is randomly spawned and tasked with locating an object from sampled category. Episodes capped 500 steps, while trials are limited to 4096 steps during training and 32k steps during evaluation.
The agent is Fetch robot with head-mounted camera (256x256 RGB) and odometry sensor that provides relative position and orientation. It operates in a discrete action space: Forward (0.25m), TurnLeft (0.3%), TurnRight, and Stop. An episode is successful if the agent stops within 2m of the goal with the object in view, occupying at least pixels in the image.
We report success rate (SR) and success weighted by path length (SPL), with SPL measuring navigation efficiency relative the shortest path. Results are computed over 32k interaction steps on the validation set and plotted at intervals of 500 steps. All experiments use 10 random seeds. Further details are in Appendix A We include results on another multi-goal 3D maze navigation task (Memory Maze) in the appendix.
18 is a Meta-RL [2] benchmark where an agent must find an invisible key to open an invisible door in a 9x9 2D gridworld. The agent receives +1 reward for finding the key and the door. It only observes its (x, y) position at each timestep and must remember key and door locations based on previous reward signals.
Park-Key-To-Door Derk-ooy-To-Door
Each episode lasts up 50 timesteps. with trial duration fixed at 500 steps. We evaluate performance by average reward across 960 trials and 3 validation seeds. An agent that leverages contextual information can complete multiple episodes per trial and achieve total reward well above 50.
4.2 Baselines
Below, we describe the main transformer-basee baselines that we compare Memo in this work.
The full context transformer baseline corresponding to ReLIC []] and AMAGO [2] We refer to it as FCT in the following sections for brevity.
Transformer without Inter-Episode Attention (no-IEA): We modify the attention masking such that the model can't 't attend over previous trials in the rollout. This sets a lower bound for what the performance would be if the model didn't 't do in-context learning and simply used its current trial's 's observations to explore and find objects.
Recurrent Memory Transformers (RMT): We also include a recurrent-only version of Memo that does not accumulate summary tokens. This bascline represents an implementation of RMT 5 adapted for the RL setting. The purpose of this comparison is to evaluate the benefits of Memo's is summary accumulation against fixed-size memory constraint. To ensure fair comparison, we set the fixed memory size in RMT equal to the total number of summary tokens Memo would accumulate over multiple segments (laum cuumeer-of-segmentt
Autocompressors (AC): Memo with pre-trained initialization and highly truncated backpropa- gation through time TBTT. This represents an implementation of AC [7] tailored to the RL setting.
We exclude RL2 and XL due to their poor performance in the results of [2] In the following subsections, we present experimental insights examining how Memo's memory creation mechanism and training recipe enhance long-horizon reasoning and in-context learning. Each subsection high- lights and focuses on specific finding and analyzes a subset of baselines drawn from figures that are shared across multiple sections, allowing to do more structured comparison.
5

---

Fallecorcat Traadooer
Sranstererr 183)
0.8
0.4
0.4
0.2
0.2
0.1
30k
Context (Steps)
Context (Steps)
(a) EXTOBINAV
(b) Results on EXTOBINAV
Figure 2: (Left) Overhead view of a training scene in Habitat simulator. The EXTOBINAV task involves placing multiple objects around the house, and then sampling an object category as the goal each time the previous goal is reached. The figure shows some sample objects that may be placed around the house, such Banana, Cracker Box, etc. (Right) Val success rate and SPL, curves 32k in in-context learning steps in novel scenes, for different methods trained with ReLIC. We compare Memo the the full-context transformer (FCT), the FCT variant which does not attend over previous episodes (no IEA), the recurrent memory transformer (RMT), and the Autocompressors (AC) variant.
4.3 Summarization Outperforms Full In-Context Access
We first show that Memo achieves higher performance and exhibits better ICI-ability compared to naive full context transformer (FCT) while using significantly less tokens in-context during evaluation.
In Figure 2b
we report the validation SR and SPL on the EXTOBINAV task over 32k environment steps. All methods are trained using on-policy RL technique proposed by ReLIC. Memo is trained with a segment length 256, 32 summarization tokens, and a trainable rollout length of 4096. We observe that Memo outperforms FCT achieving 7.5% higher SR and 2.5% higher SPL on average while using 8x fewer tokens (details in Table D We see that both Memo and FCT show in-context learning ability and continue to increase their performance up to 10k steps (~2.55x of the training context length). Beyond this point, both methods experience performance degradation due limitations in context length extrapolation. Notably, FCT degrades more than Memo in SR, while both show similar degradation in SPL. The single-episode variant (Transformer no IEA) performs significantly worse, highlighting the disadvantage of not leveraging historical context in this navigation task. We also compare the GPU memory and FLOPs for Memo and FCT in Appendix A.5 and their sample-efficiency in Appendix [A.I
In Figure 3a we present a similar comparison on the Dark- Kork-Key-To-Doww environment, reporting the average total return over trials as training progresses. All methods are trained using the off-policy RL algorithm AMAGO. We see that both Memo and RMT match the FCT baseline, achieving peak performance validation episodes by 40M training steps. Interestingly, FCT shows notable performance drop around the 35-40M step mark across all seeds, while Memo remains maintains stable convergence. This is likely due to training instability commonly observed in long-context RL 12 24,, highlighting further training challenges with full-context models.
The consistent performance gains observed across both on-policy (with ReLIC) and off-policy (with AMAGO) RL setups demonstrates that our proposal to augment transformer policies with summarization enhances long-context modeling capabilities while being independent of the specific RL algorithm used.
4.4 Advantage of Summary Accumulation Over Fixed Recurrent Memory
The summary accumulation mechanism in Memo addresses the limitations of both transformers and recurrent models. While transformers can model full context, they suffer from significant memory overhead with long sequences. In contrast, recurrent models are memory- propagate gradients effectively over long horizons due to vanishing exploding gradients. Summary accumulation mitigates context explosion by periodically summarizing experiences and propagating these summaries forward, allowing earlier memories to directly influence outcomes at later timesteps. Additionally, in purely recurrent models, gradients must pass through all intermediate memory update steps to reach the embedding of relevant input far back in time. In contrast, periodic summary
but struggle to

---

Nen
Futl Cont-Coreec Transforme
srt
4C (TOUT)
AC (all sogrents)
8.6
0.4
Avg. Return
50
Sungase
6.0
0.3
8.4
0. 0.2
20
8.2
0.1
5M
SOM 15M 28N 25M JOM SSH
JBM
0.0
5k
20k
30k
0.0
25k
Tinesteos
Cumulative Stops
Cumulative Steps
(a) Results on Dor--Key-To-Door
(b) Comparison between Memo and AC variants on EXTORINAV.
Figure 3: (Left) We plot the average return over evaluation trials of 500 steps, consisting of 10 or more episodes, plotted against the training progress. We compare Memo, RMT and FCT, when training with the AMAGO RL algorithm over random seeds. (Right) On EXTOBINAV. we find that restricting gradient propagation to a few consecutive windows (AC (TBTT)) performs much worse than allowing all past segments to remain trainable (AC (all segments)). While AC (all segments) initially performs slightly better than Memo, it suffers from severe performance degradation over time, converging to AC (TBTT) after 16k steps.
generation and accumulation (as in Memo) ensure that each summary vector contributes directly to the loss function at later timesteps without degradation. For example, at timestep all summaries created before are available as an undiminished input to the attention and subsequent optimization process.
We empirically validate this by comparing Memo's 's summary accumulation to a variant where only newly generated summaries (at the end of the current segment) are propagated to the next segment. This variant closely aligned to the Recurrent Memory Transformer (RMT) (5). To ensure fairness, we allow the recurrent model larger summary size. Even with this adjustment, Memo achieves faster convergence, improving training speed by over 10M steps on the Dark-Key-To- We further evaluate RMT on EXTOBJNAV (Figure 2b) using 1x and 2x Memo's summary size (32, 64). Although RMT outperforms a purely recurrent baseline such as Rho-showing that per-step memory updates are too limiting it still lags Memo by about 5% success rate. Larger summaries (RMT-128) destabilize training. This indicates that periodic summarization (like in Memo or Memo-RMT) improves expressivity and stability over purely recurrent per-step memory updates, while Memo's s accumulation mechanism further enhances both by introducing residual-like gradient shortcuts that enable efficient and stable long-horizon optimization (see Appendix A.7 for further analysis).
task (Figure Ba).
4.5 Importance of Long-Horizon Gradient Propagation: Comparison to Autocompressors
Since Memo's summarization approach is inspired by the Autocompressors (AC) method from the language modeling domain, we investigate how following the exact training setup from AC would perform in our setting. This involves pre-training a transformer policy on a shorter context length before fine-tuning it with summarization to extend its effective context length. We first train a baseline transformer with a context size of 1024 for 350M steps (half of the total training horizon for other baselines). This model serves as the weight initialization for our AC model, which is then fine-tuned for 350M steps with segment length leg 256, summary length Fum = 32, and trainable rollout length of 4096. Following AC, we employ Truncated Backpropagation Through Time (TBTT) during training, restricting gradient propagation to summary tokens from only the two preceding segments; we refer to this variant as AC (TBTT). For a fairer comparison to other methods, we also train version of AC where gradients are propagated across all segments within the rollout, referred to as AC (all segments). This setup is identical to Memo, except it's bootstrapped from pre-trained model with a shorter context length.
We discover that this limited propagation is insufficient for tasks requiring Long-horizon memory, as AC (TBTT) fail to capture dependencies spanning extended horizons. As illustrated in Figure 33b AC (TBTT) exhibits poorer in in-context improvement in object-finding tasks compared to AC (all segments). While AC (all segments) slightly surpasses Memo in SR and matches in SPL during the first 6-8k steps, its performance degrades significantly over time, eventually converging to AC (TBTT) after 16k steps. This suggests that while short context length pre-trained checkpoint can offer some initial advantage, it ultimately hinders effective generalization to longer contexts.
7

---

70m
Socaat Teatome
2181
1181
0.4
0.3
8
0..2
2.2
0.1
0.2
0.1
2.0
0.0
0.0
0.0
Cun www.lal..v
Slapp
Comulat.vvv
Domulat:ww
(a) Memo and Transformer streaming evaluations. (b) Finetuning for 16k steps on EXTOBINAV.
Figure 4: (Left) We show a comparison between the accumulating context (default setting) and the streaming version of Memo and the full-context transformer, where streaming starts after 6k eval steps. This comparison highlights the robustness of Streaming Memo, which maintains and even slightly improves performance, whereas Streaming Transformer suffers sharp decline. (Right) We see that finetuning over a longer number of steps during training serves to partially fix the degradation due to context length extrapolation.
4.6 Extrapolation and Experience Subsampling During Evaluation
Avoiding memory explosion due to KV cache accumulation is challenge for transformers at inference time. We evaluate Memo and the FCT baseline in streaming setting, where only the most recent 7 KV cache elements are retained for inference, enforcing fixed memory limit. To keep our setup simple, we do not update the KV cache to account for shifts in token position indices relative to the retained eache. We set T based on the context length at 6k env steps, as both Memo and FCT exhibit ICL. generalization up to at least 6k steps (see Figure 2b).. Thus, on EXTOBJNAV, we use T = 6k for Streaming Transformer and T 1024(= round (000002256) * 32 + 256) for Streaming Memo, since 6k steps correspond to 24 sets of summary tokens when feum 32 and (egg 256.
We present the results in Figure 4a Streaming Memo not only maintains performance through the trial but even outperforms Memo with full memory access in terms of ICL trend. In contrast, the Streaming Transformer undergoes a significant performance drop once streaming begins at 6k steps. Although methods like StreamingLLMM (31] address the issues faced by the Streaming Transformer, Memo achieves strong performance without requiring such modifications. Additionally, Streaming Memo provides a way to maintain peak ICL performance by identifying the highest-performing context length during evaluation and initiating streaming heyond that.
4.7 Finetuning
In Figures 2b and 4a we see a saturation in Memo's ICL performance ~60% SR. To investigate whether this limitation stems from suboptimal context length extrapolation, we take Memo's 's and FCT's checkpoints trained until 1B steps and fine-tune them on a 4x larger context size of 16,384 tokens for 500M steps. As shown in Figure 4b the ICL performance of both models improves when evaluated over 32k environment steps. Notably, Memo (16k) outperforms Memo (4k) in SPL, indicating that training with longer contexts enables the policy to find more efficient paths. While the 16k full-context transformer (FCT) improves significantly over its 4k variant, still achieves lower maximum SR and SPL compared to Memo (16k). Interestingly, Memo (16k) generalizes up to 1.5x its training context length after which it sees a slight degradation, whereas FCT (16k) maintains performance up to at least 2x. This necessitates further research into the factors limiting Memo's s long-context generalization, which we leave for future work.
4.8 Ablations.
Summary Length Ablation: There is a trade-off between creating or retaining more memory and the compute saved when generating fewer tokens at training or inference time. Additionally, creating more tokens pushes the model further into the "extrapolation" regime, where its positional encodings stop generalizing effectively. We vary the summary length between values 16, 32, and 64 (thereby varying the compression ratio, heg//sum: between 16x, 8x, and 4x) on the EXTOB/NAV AV task and plot the resulting performance in Figure Sa We find that Memo is highly sensitive to the choice of summary length. The main performance gap emerges after 6k environment steps, where a summary length of 32 achieves the best generalization across longer trajectories. In contrast, summary length of 64 performs the worst, likely due to summary tokens overfitting to the training context length (see Figure 10a), leading to redundant information storage and lower signal-to-noie ratio.
8

---

aegaent Langon
Fixed segentt Length
0.5
0.7
0.1
0.1
Comllttive 350ps
Dunutattive Sto90
cunulative
Stops
Ousulative Stops
(a) Summary length ablations
(b) Segment length randomization.
Figure 5: Memo ablations. (Left) Effect of varying the number memory tokens (16/32/64), where 32 outperforms 16, and 16 outperforms 64. (Right) Performance of Memo on EXTOBINAV with and without randomization of the segment lengths at the end of which summarization done, showing the latter is significantly less data-efficient.
Segment Randomization: In our main experiments, we adopted the strategy of randomizing seg- ment lengths in [7] for Memo, AC, and RMT make token generation more robust to specific environment time steps. We randomized segment lengths between [0.8, 1.2] x 256, resulting in a range of [205, 307]. In this ablation, we disable this randomization and evaluate the impact of using fixed segment length of 256, as shown in Figure 5b We observe that this variant performs significantly worse, not just during evaluation but in also training (see Figure TÖbj. We hypothesize that segment length randomization not only improves generalization by preventing overfitting to fixed boundaries but also serves as a form of curriculum. By exposing the model to segments of varying lengths, it naturally encounters both easier (shorter) and harder (longer) compression tasks, fostering a progressive learning effect. We present further results and ablations Appendices A.6 and A.8
5 Discussion
In this work, proposed simple yet effective approach for training transformers with memory in sequential decision-making tasks that require long-horizon reasoning. By introducing a context sum- marization mechanism, Memo enables transformers retain and retrieve task-relevant information efficiently, significantly reducing computational overhead while maintaining long-term dependencies. Our experiments demonstrated Memo's 's key advantages over full-context transformer baselines. In terms of performance vs. efficiency, Memo achieves comparable or superior performance while attending sequences over 8x shorter than the full-context transformer. This improves scalability, enhances extrapolation to longer contexts, and ensures robust performance in streaming settings, where past context must be discarded due to inference constraints.
Furthermore, our analysis highlighted two key design choices: propagating gradients across multiple summarization steps and accumulating summaries rather than using fixed memory size. We showed that allowing gradients to flow through sequential memory updates significantly enhances long- horizon reasoning, while memory accumulation improves retention of relevant past experiences, leading to more effective decision-making.
Limitations
Our experiments focus on object navigation in unseen scenes, where agents locate a fixed set of randomly placed objects. We do not evaluate semantic generalization-iia. navigating to entirely new object categories- sincee since this would conflate memorization with general object recognition capability. Future work could explore leveraging foundation models in more open-ended setting.
Additionally, while Memo effectively summarizes experience, we do not explore more flexible memory mechanisms, such as memory consolidation, where past summaries are progressively compressed. Investigating these strategies could further improve long-term memory efficiency. We also do not explicitly study length extrapolation. Our results indicate that compressed memory tokens improve robustness longer contexts and enable streaming inference, but extending generalization beyond training limits remains an open challenge. Lastly, our memory mechanism is trained end- to-end via the RI objective. Future work could explore self-supervised objectives, such as future prediction, to enhance data efficiency in training memory representations.

---

References
[1] Batra, D., Gokaslan, A., Kembhavi, A., Maksymets, O., Mottaghi, R., Savva, M., Toshev, A., and Wijmans, E. Objectnav revisited: On evaluation of embodied agents navigating to objects. arXiv preprint arXiv:2006. 13171, 2020.
[2] Beck, J., Vuorio, R., Liu, Z., Xiong, Z.. Zintgraf. L., Finn, C., and Whiteson, S. survey of meta-reinforrement learning. arXiv preprint 1rXiv:2301.8002 2023.
[3] Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A.. Shyam, P., Sastry, G., Askell, A., al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
[4] Bulatov, A., Kuratov, Y., Kapushev, Y., and Burtsev, M. S. Scaling transformer to 1m tokens and beyond with rmt. arXiv preprint 2 /062, 2023.
[5] Bulatov, A.. Kuratov, Y., and Burtsev, M. S. Scaling transformer to Im tokens and beyond with rnt. In Association for the Advancement of Artificial Intelligence, 2024.
[6] Calli, B., Singh, A., Walsman, A., Srinivasa, S., Abbeel, P., and Dollar, A. M. The ycb object and model set: Towards common benchmarks for manipulation research. In 2015 international conference on advanced robotics (ICAR), pp. 510-517. IEEE, 2015.
[7] Chevalier, A.. Wettig, A., Ajith, A., and Chen, D. Adapting language models to compress contexts. In Bouamor, H., Pino, J., and Bali, K. (eds.), Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, pp. 3829-3846, Singapore, December 2023. Association for Computational Linguistics. doi: 10. 111883//1/2233.emmm main.232. URL https:://acaaaanolllgg... org/ emnlp-main 232
[8] Cho, K., van Merrienboer, B., Gulcehre, C., Bahdanau, D., Bougares, F., Schwenk, H., and Bengio, Y. Learning phrase representations using mn encoder-decoder for statistical machine http://www..rrg////////////////
translation, 2014. URL URL http:://a......./g////////////// ats/////////// /ww.e/////...... 1078
[9] Dai, Z., Yang, Z., Yang, Y.. Carbonell, J., Le, Q., and Salakhutdinov, R. Transformer-Xi XL: Attentive language models beyond a fixed-length context. In Korhonen, A., Traum, D., and Mårquez, L. (eds.), Proceedings of the 57th Annual Meeting of the Association for Computa- tional Linguistics, pp. 1978-2988) Florence, Italy, July 2019. Association for Computational Linguistics. doi: 10. 10.18653/v1/P199 1285. hot.htttps//acaaatthti wwwphhoology.og/P99-111 1285
[10] Devlin, J. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv: 1810.04084 2018
[11] Elawady, A., Chhablani, G., Ramrakhya, R., Yadav, K., Batra, D., Kira, Z., and Szot, A. ReLIC:
A Recipe for 64k Steps of In-Context Reinforcement Learning for Embodied Al. arXiv e-prints, art. M:Xiv:2410.02755 October 2024. doi: 10.48550/axxx.2240007755
[12] Grigsby, J., Fan, L., and Zhu, AMAGO: Scalable in-context reinforcement learning for adaptive agents. In The Twelfth International Conference on Learning Representations, 2024. URR. wttprriieinnnnfffrreeeeeeee ddM
2024.
[13] He, Z., Qin, Z., Prakriya, N., Sun, Y., and Cong. J. Hmt: Hierarchical memory transformer for long context language processing. arXiv preprint 2eXiv:2405.06060 2024.
[14] Hochreiter, S. and Schmidhuber, J. Long short-term memory. Neural Comput., 98):1733- November 1997. ISSN 0899-7667. doi: 10. wwl.16//ncco.1997799 1735. URL https://doi
org/10. w162/necoo 1997 9.
[15] Huang, G., Sun, Y., Liu, Z., Sedra, D., and Weinberger, K. Q. Deep networks with stochas- tic depth. In Computer -isionEECCV 2016: 14th European Conference, Amsterdam, The Netherlands, October 11-14, 2016, Proceedings, Part IV 14, pp. 646-661. Springer, 2016.
[16] Khanna, M., Mao, Y., Jiang, H., Haresh, S., Shacklett, B., Batra, D., Clegg, A., Undersander, E., Chang, A. X., and Savva, M. Habitat synthetic scenes dataset (hssd-200): An analysis of 3d scene scale and realism tradeoffs for objectgoal navigation. In Proceedings of the IEEF/CVF Conference Computer Vision and Pattern Recognition, pp. (6384-16393 2024.
10

---

[17] Kingma, D. P. Adam: A method for stochastic optimization. arXiv preprint arXiv: arXiv:1412.6980. 2014.
[18] Laskin, M., Wang, L., Oh, J., Parisotto, E., Spencer, S., Steigerwald, R., Strouse, D., Hansen, S., Filos, A., Brooks, E., et al. In-context reinforcement learning with algorithm distillation. arXiv preprint arXiv:2210.142115 2022.
[19] Loshchilov, I. and Hutter, F. Sgdr: Stochastic gradient descent with warm restarts. arXiv preprint arXiv: 1608. 03983, 2016.
[20] Majumdar, A., Yadav, K., Arnaud, S., Ma, J., Chen, C., Silwal, S., Jain, A., Berges, V.-P., Wu, T., Vakil, J., et al. Where are we in the search for an artificial visual cortex for embodied intelligence? Advances Neural Information Processing Systems, 36:655-677, 2023,
[21] Munkhdalai, T., Faruqui, M., and Gopal, Leave no context behind: Efficient infinite context transformers with infini-attention. arXiv preprint arXiv:2404.07143, 2024.
[22] Ni, T., Ma, M., Eysenbach, B., and Bacon, P.-L. When do transformers shine in rl? decoupling memory from credit assignment. Advances in Neural Information Processing Systems, 36: 50429-50452, 2023.
[23] Ott, M., Edunov, S., Baevski. A., Fan, A., Gross, S., Ng, N., Grangier, D., and Auli, M. fairseq: fast, extensible toolkit for sequence modeling, 2019. URL https/ts:///ax http:/awwwwirraaaaaaaaaaa org/ h/sg//bb//// 1904 01038
[24] Parisotto, E., Song, H. Rae, J. W., Pascanu, R., Gulcehre, C., Jayakumar, M., Jaderberg, M., Kaufman, R. L., Clark, A., Noury. S., Botvinick, M. M., Heess, N., and Hadsell, R. Stabilizing transformers for reinforcement learning. In Proceedings of the 37th International Conference on Machine Learning, ICMI''0. 20. JMLR.org. 2020.
[25] Rae, J. W., Potapenko, A., Jayakumar, S. M., Hillier, C., and Lillicrap, T. P. Compressive transformers for long-range sequence modelling. International Conference on Learning Representations, 2020. URL httpps://pppereeiieee htt.t//see/pppeeeeeeeeee//////rs
[26] Shazeer, N. Glu variants improve transformer. arXiv preprint arXiv:2002.052022 2020.
[27] Shen, 7., Zhang, M., Zhao, H., Yi, S., and Li, H. Efficient attention: Attention with linear complexities, In Proceedings of the IEEE/CVF winter conference on applications of computer vision, pp. 3531-3539, 2021.
[28] Su, J., Ahmed, M., Lu, Pan, S., Bo. rotary position embedding. Neurocomputing, 568: 120663 2024.
Bo. W., and Liu, Y. Roformer: Enhanced transformer with
[29] Touvron, H., Lavril, T., Izacard, G., Martinet, X., Lachaux, M.-A., Lacroix, T., Rozière, B.,
Goyal, N., Hambro, E., Azhar, F., et al. Llama: Open and efficient foundation language models. arXiv preprint 2xXiv:2302.113777 2023.
[30] Wijmans, E., Kadian, A., Morcos, A., Lee, S., Essa, I,, Parikh, D., Savva, M., and Batra, D. Dd- ppo: Learning near-perfect pointgoal navigators from 2.5 billion frames. In The International Conference Leurning Representations, 2020.
[31] Xiao, G., Tian, Y., Chen, B., Han, S., and Lewis, M. Efficient streaming language models with attention sinks. arXiv preprint arXiv:2.309. 17533, 2023.
[32] Yenamandra, S., Ramachandran, A., Yadav, K., Wang, A., Khanna, M., Gervet, T., Yang, T.-Y., Jain, V., Clegg, A. W., Turner, J., al. Homerobot: arXiv preprint 1xXv:22306.11563 2023.
mobile manipulation.
[33] Zhang, P., Liu, Z., Xiao, S., Shao, N., Ye, Q., and Dou, Z. Long Context Compression with Activation Beacon. arXiv e-prints, art. Sii::2401.034622 January 2024. doi: 10.48550/rrXiv. 2401.02 103662...

---

A Appendix.
A.1
Task Setting and Model architecture: EXTOBJNAY
Our policy takes as imput an RGB image, the agent's relative position, the previous action, and the current goal object category. The RGB image is encoded using a ViT-B-sized VC-1 (0 visual encoder, finetumed by ReLIC to improve small object detection, a limitation of the original VC-1 model. The finetuned encoder remains frozen in all experiments, and downstream, applying an MLP for dimension reduction to size 256. The goal category one-hot vector and previous action are also embedded into 256 dimensions, while the agent's s position is converted into a 32-dimensional embedding. All these embeddings are concatenated together and passed to the sequence encoder.
VC-I's CLS token
All our sequence encoders (Memo, AC, etc) are based on a much smaller version of the causal auto-regressive transformer architecture used in LLaMA [29] (see details in Table The weights of the model are trained from scratch unless specified otherwise. To accelcrate data collection, the transformer's KV cache is maintained between rollout steps. Following ReLIC, we shuffle older episodes within each sequence and update the KV cache after each policy update, ensuring efficient memory reuse.
We use the same reward function for the navigation task as defined in Elawady et [.]. The reward consists of three components: (1) Geodesic distance reward: The agent receives reward
Id Ad
where d is the geodesic distance to the closest object, which change throughout the episode. (2) Slack penalty: A constant penalty of -0.003 per step encourages efficient navigation. (3) Success reward: The agent receives reward of 2 upon successfully reaching the target.
The SPL (Success weighted by Path Length) metric is computed as:
SPL 1 N ΣN S max(La,
where is the number of episodes, S; is a binary indicator of success in episode i, L is the length of the path taken by the agent, L, is the length of the shortest path the goal (computed using privileged simulator information). The SPI. values reported in the main graphs are computed at intervals 500 environment steps. Specifically, the SPL step calculated by averaging the SPL of all episodes that completed between steps t - 499 and t. So the SPL at step 5000 reflects the performance of all episodes that finished between step 4501 and 5000.
A.2 Task Setting and Model architecture: Kerk-Key-To-Door
We refer the reader to the AMAGO implementation for details on the policy architecture used on the Park-KeyTTo-Doon task. We keep all implementation and training details fixed.
http:////ggtuuuu.o////---ut CawuussinRRPP/aaann /ammaag
A.3 Training Scenes Dataset: EXTOBJNAV
Since the training dataset used in ReLIC only consisted of 333 episodes, we create our own training dataset while using the same validation dataset them. We use the same 37 train scenes [I] 32
However we generate 300 episodes per scene each having different object placement. This leads to total of 11 100 training episode. In each episode reduce the number of sampled objects make the navigation task harder since the agent has to traverse longer, more targeted paths to reach objects instead of getting away by guessing (i.e. navigating rooms with more clutter and many receptacles since the likelihood of finding an object would be higher there), We visualize the difference in sampled object distributions over episode, between the old and new datasets in Figure
A.4 Pseudocode for Memo
We present the algorithmic pseudocode for Memo in Algorithms I and 2
The modifications with respect to distributed PPO implementation (corresponding to ReLIC) are highlighted in blue. We plan to fully open-source our implementation upon publication and include an early release version at inttps://ee http://ttww...ooo//////oooooo//e icrl/mmmomm
12

---

Distribation of Objects Eplsods
Sel TC Datases
Detawet
Done
10 20 30
40
Munbor of Objects per Episode
Figure Distribution Objects in the training episodes: We visualize the changed distribution of objects sampled (and placed) randomly in the scenes corresponding to our dataset versus the original training dataset used [ our training scenes, avoid sampling more than 30 objects per scene to avoid cluttered rooms where multiple target objects could be placed together, to reduce the risk of models overfitting during training.
Algorithm 1 Memo
Initialize 0, o, buffer +- 0, summary tokens S
2: 01 reset()
3: --0
4: for tort= 1 to to N do
5:
Cit-1 t- context from B
6:
he + SeqEnca(S (Si 0z) a~ -= he)
7:
8:
(r,00+1, done) step(a) if done then
9:
O++1 reset()
11:
end
12:
i % legg = 0 then
leg is segment length
13:
St Seq Cat-1, itt
15:
end if
Append (or, a, r, 01+1, done)
17: if % num steps = 0 then
TRAIN(B, end if
19:
0: 600++1
21: end for
Algorithm 2 Train Memo
1: procedure TRAIN(B,
2: Initialize summary tokens Strain f- 0
3:
is num segments
4:
for = Otondo Cilnssiii)be f- context from
5:
6:
Strain + Seq Enco (Staann Cllcs
7:
end for
8:
Colteyyt f- context from
9:
C PPOOssss (aasssStrrani (Strains Calsegit)
10:
end procedure
13

---

A.5 Extended Result: FLOPs Comparison
Here, we include memory, FLOPs and latency comparison between Memo and Call-Context Trans- former during evaluation on EXTOBJNAV at 32k env steps. Memo requires 10x lesser memory and 4.2x lesser floating point operation, being 2x faster than the Full-Context Transformer
Table I: GPU memory usage comparison: Memo versus the Full-Context Transformer (FCT) at the end of 32k steps of evaluation on the EXTOB/NAV task. We observe 10x higher memory requirement for the KV cache of FCT, in line with the context compression ratio (~ 8) of Memo.
Memo
Full-Context Transformer
GPU Memory
51.8 MB
546.5 MB
Model FLOPs
17.61 MFLOPs
74.49 MFLOPs
Latency
5.3 mss
10.1 ms
A.6 Extended Result: Memory-Maze Benchmark
To further evaluate our method on long-horizon memory-intensive tasks, we include preliminary results on the Memory Maze benchmark. This environment poses challenging multi-goal navigation task procedurally generated Mujoco-based maze. At the start of each episode, up to six colored goal objects are randomly placed in the maze, and the agent is given randomly sampled sequence of goals to reach order. Upon reaching the current goal, the next target is revealed, continuing until the episode ends. We evaluate Memo and the Full-Context Transformer (FCT) the 9x9 version of the task, which includes two obstacles and has an episode length of 1000 steps. This setting demands both exploration and the ability recall spatial information about previously encountered goals and maze structure. Figure 8 show that Memo matches the performance of FCT on this task.
Memo
FCT
25
20
Reward
15
10
5
0
0
50M
100M
150M
200M
Context (Steps)
Figure 7: Returns of FCT vs AC on Memory Maze 9x9.
A.7 Extended Result: Benefits of Memory Accumulation on an Adversarially Long-Context RL Task (T-Maze)
To explore memory-reuse recall that we had included a recurrent-only summarisation (RMT) baseline in our experiments in Section 4.4 While RMT performance improved with larger memory sizes (e.g., RMT-64 outperformed RMT-32), Memo still achieved 5% higher success rate navigation. Runs with larger summary size (RMT-128) showed unstable convergence.
We observed an interesting trend across recurrent summarization baseline results on different bench- marks: while recurrent summarization can eventually achieve reasonable performance, they typically require significantly more training due to the need to propagate gradients through many sequential memory updates. In contrast, architectures like Memo-and transformers more generally- more efficient credit assignment by allowing earlier memories to receive gradient signals from later timesteps via direct attention mechanisms. To further validate the above intuition, we ran an exper- iment in the synthetic T-maze gridworld environment [22]. Here, the agent sees clue ("left" or
14

---

"right") at timestep O, which disappears at timestep 1. It then traverses a long corridor with only forward actions, and the end must choose the left right room based the initial clue receive reward. The corridor length can be made arbitrarily large (e.g. 10,000 steps). We hypothesized that purely recurrent models would especially struggle in this adversarial setting, as learning to retain the clue requires backpropagating gradients through all the segments corresponding to the 10,000 steps. Memo and FCT, however, can access that information directly at later timesteps-etie eether through full attention (FCT) or through small number of memory consolidation stages (Memo). Empirically, we observed the expected result: RMT required 10x longer to achieve an average reward of 1.0 for the first time during training, compared to FCT and Memo. In addition, it exhibited much greater instability in learning the correct policy, since it did not manage to maintain this performance (average reward 1.0) for three consecutive checkpoints any point during training.
RMT
Average Total Return
1.0
Meioo
0.5
-0.5
-1.0
500
1000
1500
2000
Environment Steps
Figure Returns of Memo-RM vs Memo on T-Maze (1 = 2000).
A.8 Extended Ablation: Direct KV Cache Construction from Summary Embeddings
To assess whether the improved performance of Memo stems from the additional computation it uses (due to summary embeddings being processed through the transformer twice), we perform an ablation that removes the re-encoding step.
In this variant, instead of generating summary tokens and re-encoding them through the transformer, we directly construct the memory (KV cache) from the internal embeddings produced at each layer while encoding the learnable summary embeddings. As in Memo, these summary embeddings are inserted at regular intervals between segments of tokens. However, unlike Memo, their layer-wise representations are extracted directly to serve as the KV cache, without further processing.
Meno
Pull-Contert transformer
Hero (directly generate KV coche)
0.
0
0.6
0.3 3
Success
0.4
SPL
0.2
0.1
0.0
0.0
5k
10k
15k
20k
25k
30k
5k
10k
15k
20k
25k
30k
Context (Steps)
Context (Steps)
Figure 9: Ablating the summary generation process in Memo
This approach is more compute-efficient, as each token including summaries-is passed through the transformer only once. However, as shown in Figure 9 this variant performs poorly and quickly degrades beyond the training context length. We hypothesize that this only due the lack of deeper summarization capacity, but also due to positional embedding generalization issues. Since summary embeddings are appended at the end of each segment, their positional indices lie in range [lweg /seg + summ] with offsets accumulating across segments. Consequently, these summary positions- and those of all future tokens progressively reach larger values, potentially leading to distribution shifts in the positional encodings that the model is unable to generalize over.
15

---

This result suggests that Memo's re-encoding step not only adds depth for more expressive sum- marization but also resets the positional index range for the summary tokens, both of which appear beneficial for stability and generalization.
A.9
Training-Vard Validation Gap
0.500
400
353
0..330
9.200
0.125
1.205
0.125
8.
0.000
0.000
0.600
Danulacivw
Ouaulativee St009
Dumulatioe Steps
Queulative Stace
(a) Summary length ablations.
(b) Segment length randomization.
Figure 10: Training-validation gap for Memo ablations on EXTOBJNAV. (Left) All summary tokens sizes (1 (16/32/64) perform equally well on the training set for ~ 7.5k environment steps after which 16 and 64 summary token policies start degrading faster. (Right) Fixed segment length's training performance is worse than the val performance of random segment length.
In this section, visualize the overlaid training and validation performance curves for Memo the EXTOBINAV task, which presents two key challenges. First, validation scenes are entirely novel, requiring models to generalize their in-context learning (ICL) solutions. Second, models are evaluated on 4x larger trial sizes, testing their context extrapolation capabilities. To assess generalization, we compare overfitting levels across models, distinguishing whether low validation performance stems from limited encoding expressivity (indicated by low training performance) or poor generalization despite strong training performance. We present train versus validation curves for Memo while ablating the summary lengths in Figure 10a and while ablating the segment length randomization in Figure TOb
Ckpt))
Moro (1B Ckot)
Full Context: Transforner (TREM Cipt)
Concext Transfocwor (18 Olpt)
0.8 8
0.4
4
0.6
0 3
Success
0.4
SPL
0.2
0.2
0.1
0.0
0.0
10k 15k
20k
25k
30k
10k 15k
20k 25k
30k
Cumulative Steps
Cumulative Steps
Figure 11: We compare the validation reward curves of checkpoints saved at 700M and 1000M steps of training for Memo and the full context transformer (FCT) variants on the EXTOBJNAV task trained with ReI.C.. Both checkpoints of Memo acheve higher success rate compared to the checkpoints of FCT, while its 1000M approaches the 700M Memo checkpoint's performance in terms of the SPL. metric. This highlights the sample efficiency of Memo, besides the compute efficiency which comes from using much smaller contexts in the transformer encoding and decoding process.
A.10 Sample Efficiency of Memo
We observed higher sample-effcienncy for Memo during training on the EXTORINAV and Dark- Key-To-Doo tasks compared to the full context transformer. For Derk-Key-To--Doo this is directly observable in Figure 3a since the performance is plotted against training progress. We present the this observation for the EXTOBJNA task in this section by visualizing the performance of earlier and later checkpoints (700M and IB training steps) of both FCT and Memo in Appendix A.9 We see that even earlier into training, Memo exhibits higher-ICL success rates compared to FCT which starts to get closer to Memo's ' SPL only after 300M more steps of training.
16

---

A.11 Hyperparameters
Hyperparameter
Ours/AC/RMT
Transformer
Model Architecture
Layers
4
Heads
8
Hidden dimensions
256
MLP Hidden dimensions
1024
Activation
GeLU [26
Sink-KV
0
Attention sink
N.A.
Sink K Vo
Episode index encoding
N.A. RoPE [28
RoPE [28
Within-episode position encoding
Learnable
Training Setup
Workers
320 (20 env workers per GPU x 16 GPUs)
Batch Size
160
RL Algorithm
DDPPO [30]
Discount Factor (?)
0.99
GAE Parameter (T)
0.95
Entropy Coefficient
0.1
Value Loss Coefficient
0.5
Optimization
Optimizer Regularization
Adam [7
Depth dropout [15] with value 0.1
In-context episodes shuffled after partial updates
Learning Rate Schedule
Warm-up for 100K env interactions
Initial Learning Rate Learning Rate at Warm-up End
4 10-7
x 10-1
Decay Schedule
Cosine decay [99 to after 1B steps
Computation
Precision
FP16 for visual encoder, FP32 for other model components
Rollout size
4096
Total # updates per rollout
16
partial updates
15
# full updates
1
Hardware and Performance
Environment Steps
700M Steps
Hardware
16x NVIDIA A40 GPUs
Training Time
2.5 days
Table 2: Comprehensive training setup and hyperparameters related to the EXTOBINAV task.
NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We only list the findings from experimental results in our claims.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
17

---

The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: Provided in Section
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers penalize honesty concerning limitations.
will be specifically instructed
3. Theory assumptions and proof's
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
Justification:
We do not include theorems and proofs in this work.
Guidelines:
The answer NA means that the paper does include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should complemented by formal proofs provided in appendix supplemental material.
18

---

Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: Provided in the Appendix. We also release our code.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g.. with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibilityy In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Answer: [Yes]
Justificatton: in Appendix A.4.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.cm cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new soun-s benchmark).
19

---

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
htt.ps.cc/pbblcc//uuieeeeeeeeeee isw.eeeeee/eooeeeeessiiiiisssss setieaummmsiiiiiiiii ssioniicccy for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if' applicable).
Providing much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Answer: [Yes]
Justification: Provided in the Appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: All results averaged over seeds and we indicate error intervals.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
is OK to report I-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not show in tables figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
20

---

For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: All information required reproduce the setup be found in the Appendix. Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics htttps:/neurips..e
Answer: [Yes]
Justification:
We have read through the ethics guidelines and conformed to them.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: Our work is related to foundational research on transformer architectures and does not immediately carry broader socictal impacts.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g.. disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deploymentt of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
21

---

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks. mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of MI.).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
We do not release any models or datasets.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have a high risk for misuse or dual- usal-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: All works that build cited.
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
assets are released, the license, copyright information, and terms use in package should be provided. For popular datasets, paperswi thcode. com datasete has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors encouraged to reach out to the asset's 's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification:
The released code is documented.
Guidelines:
The answer NA means that the paper does not release new assets.
22

---

Researchers should communicate the details of the dataset/code/mmodel as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
No crowd-sourcing is conducted.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: No research on human subjects is conducted.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if' applicable), such the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Answer: [NA]
Justification: Only used for editing.
Guidelines:
The answer NA means that the method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy (https /ttpa::///ee...eeeee cct.coneeeeeeccse/222555. for what should or should not be described.
23


Reviews:
[{"review_id": "mGPBlnVmiW", "reviewer": "Reviewer_SD9F", "summary": "The paper shows how to add memory tokens to a simple model-free Transformer-based RL pipeline, making it more compute and storage efficient, while outperforming baselines.", "questions": "* Did you consider using learnable positional encodings for the memory tokens?\n\n* If we reuse memory tokens rather than adding new ones each time, does the performance degrade significantly?", "limitations": "-", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "I found the paper straightforward and addressing an important problem, with well-designed experiments, strong results, and a thoughtful ablation study. The writing is clear. The results look good and largely as expected.\n\n**Weaknesses**\n\nThe most concerning aspect for me is the absence of comparisons with RNN-, SSM-, or Mamba-based baselines. Recurrent architectures eliminate the need to manage memory tokens or worry about context length, potentially making the proposed approach obsolete. So it would be valuable to include such comparisons. The authors mention that RL^2 doesn’t work well, which explains its exclusion. But what about SSM or Mamba layers? While I’m not an expert in applying SSM layers to RL, and I’m unaware of strong existing baselines using them, I understand that implementing and tuning them may be difficult and time-consuming. Nonetheless, I believe including such baselines would significantly strengthen the paper.\n\n**Comments**\n\n1) The approach may lack originality in the sense that it straightforwardly imports concepts from NLP into RL. However, I believe doing so is important, and I recognize that making such methods work reliably and perform well is non-trivial. This line of work is valuable for advancing in-context learning in RL.\n\n2) It seems the authors report the worst AC results in Figure 2(b). I would recommend showing the best results (i.e., without truncated backpropagation).", "was_revised": true, "final_justification": "I think the paper is technically strong and novel, and I am happy with the rebuttal and the additional experiments.", "rebuttal": "We thank the reviewer for their constructive feedback and positive reception of our work. We appreciate the recognition that our paper addresses a critical and interesting problem through clear writing and straightforward presentation.\n\nBelow, we respond to the remaining concerns and questions raised. We remain open and happy to address any further points during the discussion period.\n\n>“Absence of comparisons with RNN-, SSM-, or Mamba-based baselines. “ \n\nThanks for the suggestion! \n\nWe evaluated the **Mamba architecture on the Dark-Key-To-Door task**. Mamba achieved an average reward of approximately 43, compared to 55 for both FCT (Full Context Transformer) and Memo, and required about twice the training time to converge. This is consistent with a broader trend we observed across recurrent baselines: while **recurrent models can eventually achieve reasonable performance, they typically require significantly more training due to the need to propagate gradients through many sequential memory updates**. In contrast, architectures like Memo—and transformers more generally—enable more efficient credit assignment by allowing earlier memories to receive gradient signals from later timesteps via direct attention mechanisms. \n\nWe will add a similar experiment for Mamba on the extended object navigation task in the revised manuscript.\n\n\n**Synthetic Validation in T-Maze:**\n\n- To further validate the above intuition, we ran an experiment in the synthetic T-maze gridworld environment [1]. Here, the agent sees a clue (\"left\" or \"right\") at timestep 0, which disappears at timestep 1. It then traverses a long corridor with only forward actions, and at the end must choose the left or right room based on the initial clue to receive a reward. The corridor length can be made arbitrarily large (e.g., 10,000 steps).\n\n- We hypothesized that purely recurrent models would struggle in this setting, as learning to retain the clue requires backpropagating gradients through all 10,000 steps. Memo and FCT, however, can access that information directly at later timesteps—either through full attention (FCT) or through a small number of memory consolidation stages (Memo). Empirically, we observed the expected result: **RMT required ~10× longer to converge and exhibited greater instability in learning the correct policy compared to FCT and Memo**.\n\n[1] Ni et. al, When do transformers shine in rl? decoupling memory from credit assignment\n\n\n> It seems the authors report the worst AC results in Figure 2(b). I would recommend showing the best results (i.e., without truncated backpropagation).\n\nIn Fig 2(b), we tried to keep the setting close to the original AC implementation. The original AC method did not employ full backpropagation through all segments during fine-tuning, explicitly demonstrating that truncated backpropagation sufficed for their context-compression setting which had only 3 segments of summarisation on relatively short-range textual tasks. We are happy to replace the full-backpropagation variant in the main results figure if recommended by the reviewer.\n\n\n> Did you consider using learnable positional encodings for the memory tokens?\n\nWe opted against learnable positional encodings because fixed encodings (like RoPE) generalize better when extrapolating to longer contexts at inference. Although learnable encodings could potentially help differentiate summary from observation tokens, they typically struggle with extrapolation, which is critical in our setting since our evaluation steps are 8x the number of training steps.\n\n\n> If we reuse memory tokens rather than adding new ones each time, does the performance degrade significantly?\n\nWe appreciate this suggestion but would like to ensure we fully understand the intended modification. Could you please clarify what you mean by \"reusing memory tokens\" – would that mean repeatedly overwriting the same fixed set of summary tokens rather than introducing new tokens after each summarization segment (similar to RMT perhaps)? Additionally, could you elaborate on your motivation behind exploring this change—are you interested in memory efficiency, or another ablation?\n\n\nThanks again for your positive and encouraging feedback!"}, {"review_id": "V74suK2sm9", "reviewer": "Reviewer_Q5LN", "summary": "This paper introduces a transformer-based architecture and training method for reinforcement learning (RL) on memory-intensive, long-horizon tasks. The proposed model, Memo, summarizes and stores task-relevant past experiences using specialized summary tokens, which are fed back into the model at future timesteps. Crucially, Memo integrates this summarization mechanism directly into the RL optimization process, allowing gradients to propagate across all summaries by attending to and refining them. The method is evaluated on a grid-world meta-RL benchmark and a multi-object navigation (ExtObjNav) task.\nMemo addresses a key challenge in RL: unlike standard transformers that assume access to the full past context, RL agents must selectively retain and access task-relevant memories aligned with the reward signal. The proposed approach outperforms transformer-based baselines—despite having access to fewer tokens—in both scalability and efficiency.", "questions": "Please see the weaknesses.", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strengths:\n- The paper includes comprehensive experiments, strong baselines, and thorough ablations to demonstrate the effectiveness of Memo.\n- The paper demonstrates that Memo improves performance across both on-policy and off-policy RL settings.\n- It compares Memo to several strong baselines: full-context transformers, transformers without in-context learning, and recurrent memory transformers with fixed-size memory constraints.\n- On the ExtObjNav task, Memo achieves a 7.5% higher success rate and better in-context learning compared to the full-context transformer, while using 8x fewer in-context tokens and achieving 2x lower inference latency.\n- While both Memo and the full-context transformer show improved performance up to 10k steps in ExtObjNav, the full-context transformer’s performance degrades afterward—whereas Memo maintains higher performance—indicating more stable and effective in-context learning.\n\nWeaknesses:\n- Does the ability to propagate gradients to all past summary tokens introduce significant training overhead? How does this affect training time compared to baselines?\n- RMT appears unstable when using a higher number of summary tokens, especially in ExtObjNav. Could the authors clarify why this degradation occurs only in this task and not in the Dark-Key-to-Door task?\n- How does Memo differ architecturally and methodologically from Autocompressors (AC)? Is the novelty primarily in applying the summarization approach to RL and training from scratch?\n- Memo is highly sensitive to the number of memory tokens, with performance degrading significantly when increasing the token count to 64. Can the authors provide intuition for this sensitivity?", "was_revised": true, "final_justification": "I keep my positive score. Most of my concerns were addressed. The authors explained about the training overhead of propagating gradients to summary tokens, they mentioned an error in the RMT plot which will be fixed for the camera-ready with more expected results, and they explained the novelty of the method compared to Autocompressors (AC).", "rebuttal": "We thank the reviewer for their constructive feedback and appreciate their positive reception of our work, particularly for recognizing our experiment as comprehensive and well-designed with thoughtful ablations.\n\nWe address the remaining concerns and questions below, and we remain open to discussing any further points during the discussion period.\n\n> Does the ability to propagate gradients to all past summary tokens introduce significant training overhead? How does this affect training time compared to baselines?\n\n- In terms of wallclock time, it takes 60 vs 69 hours for our training runs of FCT (Full Context Transformer) v/s Memo (due to forward passes on sequence of split up context instead of one forward pass on a longer context). In contrast it takes 57 hrs for the same number of training steps for AC (TBTT).\n- In practice, we observe that there aren’t significant differences in the train time memory GPU usage of Memo v/s FCT since Flash Attention already optimises GPU memory to be below what the naive implementation would take, and so the savings in Memo from much shallower attention operations (because now tokens don’t attend over very long contexts and instead over memories and immediate context) is somewhat cancelled out by the memory needed due to the extra KV cache corresponding to memory tokens. \n- While Memo requires a modest increase in training time, this is **compensated by substantial efficiency and performance gains**. Specifically, Memo is 2× faster than FCT at inference (5.3 ms vs. 10.1 ms per step) and consumes 10× less GPU memory (51.8 MB vs. 546.5 MB). Additionally, Memo outperforms the AC baseline in downstream task performance, achieving 62% success compared to 50% for AC. \n\n> RMT appears unstable when using a higher number of summary tokens, especially in ExtObjNav. Could the authors clarify why this degradation occurs only in this task and not in the Dark-Key-to-Door task?\n\nWe had an important correction to make for the RMT plot: we reported the instability in the paper and after paper submission we decided to dig deeper into sources of bugs for that run and found an erroneous gradient accumulation config setting to be the cause of anomalous performance trends there. We reran and got much more expected results, the performance of the 64 summary len run is 4-5% higher than rmt-32 and 5% lower than the highest ICL performance achieved by Memo.  We will include this revised plot in the manuscript.\n\n> Memo is highly sensitive to the number of memory tokens, with performance degrading significantly when increasing the token count to 64. Can the authors provide intuition for this sensitivity?\n\nOn both the training (Fig. 9a) and validation sets (Fig. 5a), Memo's performance is comparable across all three summary lengths up to ~6k context length. However, beyond this point, the configuration with 64 summary tokens exhibits the weakest generalization. While we do not yet have a definitive answer, we offer the following context and hypotheses that we plan to investigate further for the camera-ready version:\n\n- Consistency with observation in AC: The AC paper reports a similar sensitivity to summary length but does not provide a concrete explanation for this behavior, suggesting this is a broader phenomenon not unique to our setup.\n- Context Explosion and Extrapolation: Increasing the number of summary tokens may lead to faster context growth across segments, pushing the model into a harder extrapolation regime more quickly. This might explain why the 64-token setup degrades earlier than 32 and 16.\n- Training Exposure Gap: However, this doesn't fully explain why the 16-token run underperforms relative to the 32-token run. One possibility is that with only 16 summary tokens per segment, the model sees a much smaller context (and therefore range of positional embeddings) during training. If we follow the findings in [1], which show that transformers trained on long sequences generalize disproportionately better to even longer sequences compared to transformers trained on short sequences and then evaluated on sequences with the same extrapolation factor. We believe Memo(16)’s quicker degradation here could be a consequence of both limited positional exposure and reduced effective context length during training.\n\n\nUnderstanding this interaction fully would require targeted ablations, which we plan to explore by the camera-ready deadline.\n\n\n[1] Zhou et. al, Transformers Can Achieve Length Generalization But Not Robustly\n\n\n> How does Memo differ architecturally and methodologically from Autocompressors (AC)? Is the novelty primarily in applying the summarization approach to RL and training from scratch?\n\nBesides the differences mentioned by the reviewer, we further note some subtle differences in our experimental setup and findings below:\n- AC compares models based on average perplexity on later chunks in a sequence, using a token-matched baseline (e.g., comparing 50 summary tokens against 50–150 raw tokens, rather than the full 750-token context that 50 tokens were meant to be summarising). Our baseline (Full Context Transformer) has access to the entire raw context that Memo is meant to be summarizing—making the comparison more stringent. Unlike the straightforward addition of summary tokens during supervised training in AC, Memo operates in an online in-context RL setting, which brings in additional challenges: needing to manage a KV cache of both the recent history and memory tokens during rollouts in the rollout buffer, recomputing all encodings to maintain on-policy-ness during update steps, followed by refreshing the cache for subsequent data collection. Getting memo to work in this multi-episode RL setting required non-trivial effort, which we believe will now be easier for others to adopt with our open-sourced implementation (Appendix.4).\n- We show that Memo matches or outperforms the full-context baseline, even when streaming (i.e., pruning older summaries’ KV cache), whereas the full-context model degrades.\n- Finally, we show that end-to-end trainability of the context is critical: Memo allows gradients to flow across the entire context, whereas TBTT (Truncated Backpropagation Through Time), as used in AC, fails to capture long-term dependencies effectively. This result has wider implications when trying to adopt AC for training LLMs to do long context tasks that impose a different reasoning problem from what a short context model has learnt to do.\n\nThanks again for your constructive and encouraging feedback!"}, {"review_id": "7ELtDAkTw0", "reviewer": "Reviewer_VTDX", "summary": "In this paper, the authors propose Memo, which is a transformer-based RL framework that uses periodic summary tokens to compress historical context, enabling efficient long-horizon decision-making. Memo interleaves learnable summarization tokens at fixed intervals during training, creating a condensed memory buffer that replaces the full KV cache. Experimental evaluations on ExtObjNav (3D navigation) and Dark-Key-To-Door, show that Memo outperforms full-context transformers (FCT) and recurrent baselines (RMT).", "questions": "1. Why does Memo generalize worse than FCT beyond 1.5× context (Fig. 4b)? Is this due to positional encoding saturation or information loss in summaries? Provide analysis beyond speculation.\n\n2. How does Memo ensure summaries capture task-critical information (e.g., object locations vs. irrelevant details)? The RL objective may not suffice for optimal compression.\n\n3. Memo requires 16×A40 GPUs for 2.5 days. Does the inference efficiency justify this? Compare total energy/compute vs. FCT.\n\n4. When does Memo perform poorly? E.g., in tasks requiring high-fidelity recall (e.g., \"wall color as landmark\" in Sec. 1), could summarization discard crucial details?", "limitations": "N/a.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "## Pros\n1. The topic of this paper is critical and very interesting. Memo adapts language-model context compression to RL’s unique challenges (credit assignment, partial observability), filling a gap in memory-efficient embodied AI.\n\n2. Extensive experiments across tasks (grid-world, photorealistic sim), RL algorithms (on/off-policy), and metrics (SR, SPL, FLOPs, latency) show the effectiveness of MeMo.\n\n3. Memo extrapolates better to longer contexts than FCT (Fig. 4b) and shows stable convergence where FCT fails (Fig. 3a).\n\n4. Well-designed attention masking, positional encoding, and KV-cache refresh mechanisms ensure gradient flow and memory consistency.\n\n## Cons\n\n1. In this paper, the experiments focus mainly on navigation with fixed object sets. No evaluation of semantic generalization (e.g., novel object categories) or dynamic environments, raising questions about scalability to real-world embodied tasks.\n\n2. Memo degrades beyond 1.5× training context (vs. FCT’s 2×). The paper attributes this to \"suboptimal extrapolation\" but lacks investigation or explanation into root causes (e.g., positional encoding drift).\n\n3. The paper omits some comparisons/ablations to state-of-the-art memory methods (e.g., Memorizing Transformers, StreamingLLM) or hybrid architectures (e.g., transformer-RNN).", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their positive feedback and recognition of our work, including the comprehensive experiments and the critical problem addressed. We also appreciate the acknowledgment of Memo’s superior stability in performance and convergence over longer contexts.\n\nWe address the remaining concerns and questions below, and we remain open to discussing any further points during the discussion period.\n\n> In this paper, the experiments focus mainly on navigation with fixed object sets. No evaluation of semantic generalization (e.g., novel object categories) or dynamic environments, raising questions about scalability to real-world embodied tasks.\n\nWe agree. However, at the time of our work, no established benchmarks offered both the diversity and long-range dependencies necessary to meaningfully train or evaluate such capabilities in embodied agents. That said, our method is designed with these challenges in mind and its general context summarization mechanism should be trainable given any long horizon task. We see the development of benchmarks targeting long-horizon memory in open-ended settings as a key enabler for future work, and we plan to pursue this in follow-up research.\n\n\n> Why does Memo generalize worse than FCT beyond 1.5× context (Fig. 4b)? Is this due to positional encoding saturation or information loss in summaries? Provide analysis beyond speculation.\n\nGreat question. We conducted several targeted evaluations to better understand this behavior.\n\nFirst, we note that Memo does generalize well under a streaming evaluation regime. In Fig. 4a (trained on 4k), Memo performs competitively with FCT even when past summaries are replaced continually, suggesting that **older summaries can be replaced by newer ones without degradation—as long as the model operates within the distribution of input lengths it was trained on**.\nTo further probe the 16k-finetuned Memo model used in Fig. 4b, we ran two streaming evaluations:\n\n- **Recent Memories Only:** We retained only memory summaries corresponding to the most recent 20k environment steps and discarded all earlier memories.\n- **Old + Recent:** We kept only the oldest 20k memory steps and appended a single recent memory chunk (assigning it position index N+1, where N is the number of old memories).\n\nIn both cases, Memo maintained its performance after streaming began. This suggests that information loss in summaries is not the primary cause of degradation when generalizing beyond 1.5× context.\n\nThe degradation likely stems from a distributional shift at inference time—specifically, Memo attending over far more input tokens than it was trained on. Due to summarization, Memo sees only a bounded number of memory tokens during training, whereas FCT is consistently exposed to full-length sequences. This also limits the range of positional embeddings that Memo experiences, further impacting its ability to generalize to longer contexts. These observations are consistent with findings in [1], which show that transformers trained on longer sequences generalize better to even longer ones. We believe Memo’s quicker degradation is a consequence of both limited positional exposure and reduced effective context length during training.\n\n[1] Zhou et. al, Transformers Can Achieve Length Generalization But Not Robustly\n\n> The paper omits some comparisons/ablations to state-of-the-art memory methods (e.g., Memorizing Transformers, StreamingLLM) or hybrid architectures (e.g., transformer-RNN).\n\n**StreamingLLM:**\n\nIn the paper, we already incorporate two aspects inspired by StreamingLLM. \n\n- The FCT baseline we used (from the RELIC paper) already includes an attention sink token, as recommended by StreamingLLM. RELIC had ablated this and found the sink token version to perform better. \n\n- We also included evaluations in a streaming setup where we do windowed inference over the KV cache, and found that FCT’s performance degraded significantly—even with a sink token (Figure 4a). We explored several mitigation strategies inspired by StreamingLLM (e.g., retaining a few early tokens along with the most recent window of tokens, preserving the tokens for the first episode in full to avoid mid-episode truncation discontinuities etc), but none resolved the degradation. In contrast, Memo used a simple FIFO strategy for selecting memory tokens and remained robust without any special design changes.\n\n**Mamba:**\n\nFollowing the suggestion from SD9F, we evaluated the Mamba architecture on the Dark-Key-To-Door task. Mamba achieved an average reward of approximately 43, compared to 55 for both FCT and Memo, and required about twice the training time to converge. This is consistent with a broader trend we observed across recurrent baselines: **while recurrent models can eventually achieve reasonable performance, they typically require significantly more training due to the need to propagate gradients through many sequential memory updates**. In contrast, architectures like Memo—and transformers more generally—enable more efficient credit assignment by allowing earlier memories to receive gradient signals from later timesteps via direct attention mechanisms. We will add a similar experiment for Mamba on the extended object navigation task in the revised manuscript.\n\n**Transformer-XL (Tr-XL):**\nWe also evaluated Tr-XL to test approaches where memory creation is based on fixed heuristics and does not benefit directly from learning signals (e.g., using a frozen kv cache or averaging the cache elements). As also reported in RELIC [3], **Tr-XL performed poorly in our setting**, achieving only ~15% success and showing no in-context learning capabilities.\n\n> How does Memo ensure that summaries capture task-critical information (e.g., object locations vs. irrelevant details)? The RL objective alone may not guarantee optimal compression.\n\nMemo is trained end-to-end using task-specific reinforcement learning objectives, such as navigation success. As a result, the summarization mechanism is implicitly optimized to retain information that directly impacts downstream performance (e.g., object locations, path-finding cues), while discarding irrelevant or redundant details (e.g., wall colors, if they are not task-relevant). We emphasize that this form of lossy compression is not arbitrary; rather, it is shaped by the learning signal provided through interaction. If certain features consistently influence task success, the model learns to preserve them through the summary mechanism. This stands in contrast to manually designed memory systems that may enforce fixed heuristics, which can be suboptimal across diverse tasks.\n\nThat said, it would be interesting to explore self supervised learning for pre-training the memorization, especially when considering large vision-language models and multi-modal dialogue datasets which implicitly capture what information humans tend to be interested in and talk about.\n\n> When does Memo perform poorly? Could summarization discard crucial details in tasks requiring high-fidelity recall (e.g., \"wall color as landmark\" in Sec. 1)?\n\nWe note that the failure modes we observe in Memo are qualitatively similar to those seen in baseline methods (e.g., agents getting stuck, or failing to detect small objects from a distance). This suggests that the summarization mechanism is likely not introducing new or unique failure modes, but rather inherits the broader limitations faced by embodied agents in complex visual environments. Addressing these foundational challenges in perception and exploration remains an important direction for future work.\n\n> Memo requires 16×A40 GPUs for 2.5 days. Does the inference efficiency justify this? Compare total energy/compute vs. FCT.\n\nMemo only takes 15% longer to train compared to FCT while being substantially more efficient during inference. \n- For a 4k sequence length training setup, FCT training takes ~60 hours with Memo taking 69.5 hours. The additional time in Memo stems from its segment-wise processing during PPO updates, which incurs extra forward passes on a split-up context. \n- Memory usage per GPU during training is also comparable: FCT uses ~48.3% of GPU memory per device, while Memo uses ~46.9% (we weren’t able to double the batch size as some GPUs exceeded 50% memory usage).\n- During inference at 32k environment steps, **Memo is 2× faster (5.3 ms vs 10.1 ms per step) than FCT** and **Memo uses 10× less GPU memory (51.8 MB vs 546.5 MB)**. \n\nThis significant reduction in inference cost is especially **valuable in embodied AI** scenarios, where models are deployed on edge devices with strict memory and latency constraints.  In such settings, multiple models—e.g., perception, mapping, and policy—often run concurrently on a single robot. A smaller and faster policy model, like Memo, frees up memory and compute, allowing other components to operate more efficiently or enabling the deployment of more capable models within the same hardware.\n\nThanks again for your constructive and encouraging feedback!"}]

Author rebuttal:

Meta-review: This paper introduces Memo, a transformer based model which incorporates memory tokens for long-horizon, memory intensive tasks. While the paper is an application of ideas from the NLP literature to RL, I think it is still a valuable contribution.

All the reviewers are positive about the paper. All the major concerns, like lack of comparison with Mamba, training efficiency, inference costs, and reasons for instability, are all properly addressed by the authors, and reviewers are satisfied with the answers. It is a good paper that is worth publishing at NeurIPS!
Decision: Accept (spotlight)