Meta-reinforcement learning aims to train agents that adapt quickly to unseen tasks from limited interaction, but this requires representing task evidence over long interaction histories. Classical recurrent agents compress experience into hidden states, while transformer-based agents retrieve from past trajectories at quadratic memory and computation cost. In this work, we propose Memory-Enhanced Test-Time Adaptation (METTA) for meta-reinforcement learning. It decomposes the fast adaptation phase into two temporal scales: a causal neural memory performs fine-grained within-episode adaptation with linear sequence complexity, while compact episode-level summaries support long-range cross-episode task inference without attending over the full interaction history. Moreover, equipping agents with neural memory provides sufficient expressive capacity and test-time training capability, enabling effective adaptation to unseen held-out tasks and task families at meta-test time. On unseen held-out Meta-World ML10/45 test environments, METTA achieves a success rate over 50% while maintaining a low training budget and short wall-clock time. On long-context sparse-reward PointMaze tasks, it outperforms transformer-based baselines by exhibiting more structured exploration behavior. These results show that METTA provides an efficient path toward more adaptive and scalable long-context meta-RL.
Problem Formulation and Preliminaries
Reinforcement learning studies how agents learn to make sequential decisions through interaction with an environment, using rewards as feedback to improve future behavior. Despite strong performance in fixed environments, standard RL methods often require large amounts of data to adapt to new tasks and generalize poorly under environment shifts. Meta-reinforcement learning addresses this limitation by training agents over a distribution of tasks so they can use past training experience to adapt quickly to unseen tasks from limited interaction at meta-test time. A fundamental problem in meta-RL is representing task evidence collected during interaction with the environment. Black-box sequence-based agents learn this representation directly from trajectories using internal dynamics, while latent-context methods make task inference more explicit by inferring a task variable or belief from collected experience. Yet despite their differences, most approaches still rely on compressed task representations, which may restrict their capacity in meta-RL settings where task evidence may be sparse, delayed, or unevenly distributed across episodes. Once experience is compressed into a latent task representation, the policy cannot easily recover specific past evidence that later becomes important.
METTA: Memory-Enhanced Test-Time Adaptation
In this work, we propose Memory-Enhanced Test-Time Adaptation (METTA) for meta-RL. We consider a trial-based meta-RL setting in which the agent interacts with multiple consecutive episodes from the same sampled task. Each trial is generated from a hidden task and contains E episodes, each with horizon H. At episode e and timestep t, the agent uses the interaction history to construct a context representation z_{e,t} that summarizes task-relevant evidence accumulated during adaptation. The policy conditions on both the current state and this context representation for action selection.
Besides the standard slow parameter updates during meta-training, METTA performs fast adaptation at two temporal scales. A within-episode neural memory supports fine-grained online adaptation, while a cross-episode memory aggregates information across episodes. Together, these modules produce the context representation z_{e,t}, which provides the policy with information accumulated through interaction. This temporal hierarchy provides an efficient mechanism for long-context adaptation with linear sequence complexity. Unlike standard recurrent approaches that rely on passive memory updates during inference, the TTT module performs gradient-based memory updates, providing an additional mechanism for fast adaptation to previously unseen task families.
Within-Episode Neural Memory
We use a Test-Time Training (TTT) architecture as the within-episode neural memory in METTA. Given an input sequence x_{e,t} = [s_{e,t}, a_{e,t-1}, r_{e,t-1}, d_{e,t-1}] processed in an autoregressive manner, sequence processing is performed by TTT layers with linear computational cost. We first embed the interaction input, and then pass this representation through L stacked TTT layers. For each layer, the model maintains a fast-weight memory that is initialized at the start of the episode and updated online as the episode unfolds.
The memory is updated by training the fast weights in a self-supervised manner to predict the label view from the training view, using a self-supervised loss. The output is then produced by applying the updated memory to the test view. Thus, the output is computed from a memory model that has been adapted on the causal prefix. The distinction between inner-loop and outer-loop parameters is central to TTT: the memory parameters are fast parameters initialized at the start of a sequence and updated while the sequence is processed, including at test time, while the view projections and initial memory parameters are slow parameters learned across training sequences.
Cross-Episode Memory
The within-episode TTT module adapts to the current episode, but meta-RL trials also contain useful information across episodes. We use the final adapted representation of each episode as an episode-level summary. At episode e, the agent causally has access to all previous summaries. We consider two mechanisms for aggregating these summaries to construct the context representation.
The first mechanism maintains an exponential moving average (EMA) of previous episode summaries. The current within-episode representation is combined with the recent cross-episode summary to form the current context representation, where a parameter controls the influence of previous episodes. For the first episode, no previous summary is available, so the context representation reduces to the within-episode representation, and the EMA summary for the first episode is set to the first episode's summary. The EMA summary for subsequent episodes is updated as a weighted combination of the previous summary and the current episode's summary. This variant stores only a single vector and updates cross-episode information in constant time.
The second mechanism, used as the default aggregation mechanism in METTA, employs an attention layer over stored episode summaries. A learnable attention layer forms a query from the current within-episode representation and uses the unordered set of episode summaries as keys and values for information retrieval. The retrieved cross-episode representation is fused with the current within-episode representation through a layer normalization operation. Compared with EMA, attention-based retrieval is more expressive because it can learn which previous episode summaries contain the most relevant task information.
Experimental Evaluation
We evaluate the effectiveness of METTA on standard meta-RL benchmarks covering locomotion, manipulation, and navigation tasks. On continuous-control adaptation, we evaluate on MuJoCo HalfCheetahVel, where the agent must infer the target velocity from interaction under a limited training budget. On robotic manipulation, we evaluate on ML1, ML10, and ML45 from the Meta-World benchmark. In ML1, goal coordinates are masked and must be inferred through interaction. ML10 and ML45 span multiple task families with disjoint train and test splits, allowing evaluation of generalization to families not observed during meta-training. During multi-class training, each rollout batch assigns environments to task classes in a balanced way and then samples a task variation within each class.
Finally, we use PointMaze with sparse rewards (only at the goal position) to evaluate long-context adaptation. In this environment, the agent may need to explore for many steps before receiving useful feedback, making it important to preserve information from earlier parts of the trajectory. This benchmark directly tests whether meta-RL methods can benefit from long context.
Main Results
On MuJoCo HalfCheetahVel with a limited training budget, METTA improves rapidly with low training variance. The episode-wise adaptation plot shows that METTA maintains strong returns from the first episode and remains stable across later episodes, suggesting that its neural memory can quickly infer the target velocity from interaction.
In long-horizon tasks with sparse rewards, METTA rapidly improves the success rate during both meta-training and meta-testing, while other baselines remain substantially lower. Transformer-based baselines struggle in this setting; TrMRL is difficult to scale to long effective horizons and tends to repeat similar exploration behavior, while ECET is limited by short-context cross-episode conditioning. On the more challenging ML10 and ML45 benchmarks, METTA surpasses other baselines by a large margin across diverse task families. The performance gap is notable not only on train environments, where evaluation uses task families observed during meta-training, but also on the Meta-World test environments, which consist of unseen held-out task families. This result indicates that the learned adaptation mechanism of METTA generalizes beyond the task families encountered during meta-training, providing evidence of task-family-level out-of-distribution generalization.
Exploration Behavior on PointMaze
METTA achieves the highest coverage and lowest overlap while requiring the fewest steps, suggesting that it maintains diverse exploration across episodes without sacrificing navigation efficiency. In contrast, ECET and TrMRL cover less of the maze while revisiting previously explored regions more frequently. The exploration patterns suggest that METTA can preserve useful evidence from earlier exploration and reuse it in later episodes, which is particularly important when useful evidence is sparse and must be accumulated through interaction.
Online Memory Updates and Training Efficiency
To isolate the contribution of test-time memory adaptation, we evaluate METTA with its online memory updates disabled while retaining cross-episode aggregation. On PointMaze, disabling the online update causes a large drop in test performance, showing that continual within-episode adaptation is particularly important when useful evidence is sparse. On ML10, the two variants perform similarly, where the goal is observed and adaptation mainly requires identifying the task family, which is easy to distinguish once learned. On ML1 Push, where the goal is hidden and must be inferred through interaction, METTA with online updates reaches substantially higher success in the first evaluation episode, before any previous-episode information is available for cross-episode aggregation. Without online updates, performance is much lower initially and improves as information accumulates across episodes, but remains below the full model even after many episodes.
METTA dominates transformer-based methods, achieving stronger performance with substantially lower GPU-hour cost. Its training time remains comparable to widely used baselines, while surpassing their performance by a large margin. This shows that neural memory provides a favorable trade-off between long-context adaptation and practical training cost. Over a complete trial, METTA requires O(E²T) computation with attention-based aggregation and O(ET) with EMA, compared with O(E²T²) for TrMRL and O(ET³+E³T) for ECET. Since typically E ≪ T, the dependence on timestep-level history dominates in long-context settings. TTT incurs a constant per-step cost from its inner update, but this cost does not grow with timestep-level history.
Conclusion
We presented METTA, a meta-RL framework based on neural memory and test-time training. The method separates fast adaptation into within-episode online memory updates and cross-episode evidence aggregation, allowing the agent to leverage long interaction histories without attending over the full history. This temporal hierarchy provides an efficient mechanism for long-context adaptation with linear sequence complexity. Unlike standard recurrent approaches that rely on passive memory updates during inference, the TTT module performs gradient-based memory updates, providing an additional mechanism for fast adaptation to previously unseen task families. Experiments showed that METTA improves adaptation across both standard and long-context meta-RL settings while maintaining favorable computational efficiency. The gains were especially strong on unseen held-out Meta-World task families, providing further evidence that neural memory with test-time adaptation can support task-family-level OOD generalization to previously unseen, novel task families. Our findings also suggest a promising direction for future work that extends METTA beyond the fixed-task trial setting considered here, for example by incorporating temporal or recency information into the cross-episode memory for continual or nonstationary environments that require fast online adaptation as task structure, dynamics, or rewards change over time.