PaperlayerView exact revision on arXiv

Source-linked AI summary

METTA: Memory-Enhanced Test-Time Adaptation for Meta-Reinforcement Learning

Ali Najar, Hosein Hasani, Amir Kooshan Fattah, Yasamin Rajabi, Mahdieh Soleymani Baghshah

arXiv:2610.31656v1cs.LGcs.AI

TL;DR

Meta-RL must infer task-relevant evidence from long interaction histories to adapt to unseen tasks, while compressed representations and full-history attention have limitations. METTA combines within-episode neural memory with cross-episode summaries; experiments report improved adaptation, especially on unseen task families, with favorable computational efficiency.

  • Problem

    Meta-RL requires agents to represent task evidence across long interaction histories, but compressed representations may lose details and full-history transformer attention becomes costly as context grows.

    Source: §1 · pp. 1–2

  • Method

    METTA uses causal neural memory for within-episode adaptation and episode-level summaries for cross-episode evidence aggregation without attending over the full interaction history.

    Source: §1 · p. 2, §3.1 · p. 5, §5 · p. 9

  • Results

    METTA improves adaptation across standard and long-context meta-RL benchmarks, with especially strong gains on unseen held-out Meta-World task families and lower training cost than transformer-based methods.

    Source: §1 · pp. 2–3, §5 · p. 9

  • Takeaways & Limitations

    The results suggest neural memory with test-time adaptation can support task-family-level generalization beyond families observed during meta-training.

    Source: §1 · p. 2, §5 · p. 9

    The evaluated setting uses fixed-task trials; extending METTA to continual or nonstationary environments is left for future work.

    Source: §5 · p. 10

Abstract

from arXiv · show

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.

1 INTRODUCTION

METTA addresses the challenge of using long interaction histories for meta-RL by separating within-episode adaptation from cross-episode evidence aggregation. Experiments report improved adaptation across benchmarks, with particularly notable gains on unseen held-out Meta-World task families.

1 INTRODUCTION

  • Compressed task representations may restrict meta-RL when evidence is sparse, delayed, or unevenly distributed, because later-relevant details can be difficult to recover.This concern applies to both sequence-based and latent-context approaches discussed in the introduction.Source: §1 · p. 1
  • On PointMaze, METTA explores new regions while reducing repeated trajectories across episodes; TrMRL struggles to scale to long multi-episode settings, while ECET has short within-episode context.The trajectories shown are from Episodes 1, 2, 5, and 10 of the same task.Source: Figure 1 · p. 2
  • Full-history transformer attention grows quadratically with sequence length, making long online adaptation histories increasingly costly in memory and computation.The effective context grows with both episode horizon and the number of adaptation episodes.Source: §1 · p. 2
  • METTA combines within-episode neural memory with cross-episode summaries, avoiding quadratic computation over interaction steps while using longer meta-test histories.This separates fine-grained within-episode adaptation from cross-episode evidence aggregation.Source: §1 · p. 2
  • METTA improves fast adaptation across locomotion, manipulation, and navigation benchmarks, with especially notable gains on unseen held-out Meta-World task families.The reported improvements cover both standard and long-context meta-RL settings.Source: §1 · p. 2

2 PRELIMINARIES

The preliminaries define meta-RL as adaptation across a distribution of tasks and describe neural memory as a state-update and readout mechanism. They introduce TTT as an online-trained memory model whose adapted state produces outputs from the causal prefix.

2.1 META-REINFORCEMENT LEARNING

  • An MDP specifies state and action spaces, transitions, rewards, an initial-state distribution, a discount factor, and a horizon.The agent selects actions from a policy, receives rewards, and transitions to subsequent states.Source: §2.1 · p. 3
  • Meta-RL samples tasks that share observation and action spaces but may differ in dynamics, rewards, or goals; agents train across the task distribution and adapt during multi-episode testing.At meta-test, behavior improves using internal memory and rollout interaction history.Source: §2.1 · p. 3

2.2 NEURAL MEMORY

  • Neural memory updates an internal state from sequence inputs and reads from that state to produce outputs, covering recurrent states and richer long-horizon memories.The preliminaries also describe key-value writing and query-based retrieval as an attention-style interface.Source: §2.2 · p. 3
  • TTT trains its memory online to predict a label view from a training view, then applies the updated memory to the test view from the causal prefix.The inner-loop memory parameters update during sequence processing; projections and downstream parameters are learned across training sequences.Source: §2.2 · p. 4

3 METTA

METTA divides fast adaptation between within-episode neural memory and cross-episode evidence aggregation. These modules produce context representations that condition the policy on information accumulated through interaction.

3 METTA

  • METTA handles trials of consecutive episodes from one sampled task, combining within-episode neural adaptation with cross-episode evidence aggregation into policy context.Source: §3 · pp. 4–5

3.1 WITHIN-EPISODE NEURAL MEMORY

  • Within each episode, stacked TTT layers update fast weights online, using causal sequence processing with linear computational cost.Source: §3.1 · p. 5

3.2 CROSS-EPISODE MEMORY

  • EMA combines the current episode representation with a prior summary, retaining one vector and updating cross-episode information in constant time.Source: §3.2 · p. 5
  • Default attention retrieves selectively from stored episode summaries and fuses retrieved information with current representation for actor and critic inputs.Source: §3.2 · pp. 5–6

4 EXPERIMENTS

METTA is evaluated on locomotion, manipulation, and navigation benchmarks, with tests of held-out task-family generalization, sparse-reward exploration, online memory updates, and training efficiency. The results report stronger adaptation across these settings, with the clearest held-out Meta-World gains on ML10 and ML45.

4.1 EXPERIMENTAL SETUP

  • The evaluation covers HalfCheetahVel, Meta-World ML1/ML10/ML45, and sparse-reward PointMaze, against recurrent, gradient-based, latent-context, and transformer-based baselines.The setup uses a limited training budget for HalfCheetahVel and disjoint train/test task-family splits for ML10 and ML45.Source: §4.1 · pp. 6–7

4.2.1 ROLE OF LONG-CONTEXT MEMORY

  • On held-out Meta-World ML10/ML45 test environments, METTA surpasses other baselines by up to a ∼40% success-rate margin.Figure 4 contrasts unseen instances from training families (top) with held-out test families (bottom).Source: Figure 4 · p. 7, §4.2.1 · p. 8
  • In sparse-reward PointMaze, METTA rapidly improves success during meta-training and meta-testing while baselines remain substantially lower.It achieves the highest coverage and lowest consecutive-episode overlap while requiring the fewest steps.Source: §4.2.1 · pp. 7–8
  • METTA’s Meta-World gains extend across ML1, ML10, and ML45, with large margins over baselines on the more challenging ML10 and ML45 benchmarks.Figure 4 includes both seen-family task instances and held-out task families for ML10 and ML45.Source: Figure 4 · p. 7, §4.2.1 · p. 8
  • Context representations cluster by task family in t-SNE, align test tasks with related training tasks, and show within-family similarity in the cosine matrix.Figure 5 displays t-SNE representations for all episodes and cosine similarities between final trial representations.Source: §4.2.1 · p. 8, Figure 5 · p. 8

4.3 ONLINE MEMORY UPDATES

  • With online updates disabled, PointMaze test performance drops substantially, whereas the two variants perform similarly on ML10.On hidden-goal ML1 Push, online updates yield higher success in episode 1; the disabled variant improves over 25 episodes but remains below the full model.Source: §4.3 · pp. 8–9, Figure 6 · p. 9

4.4 TRAINING EFFICIENCY

  • METTA achieves stronger performance than transformer-based methods at substantially lower GPU-hour cost, while training time remains comparable to standard baselines.Figure 7 compares success against GPU hours on ML10.Source: §4.4 · p. 9, Figure 7 · p. 9
  • Across a complete trial, METTA’s attention aggregation costs O(E2T) and EMA costs O(ET), versus O(E2T 2) for TrMRL and O(ET 3 + E3T) for ECET.TTT adds a constant per-step inner-update cost that does not grow with timestep-level history.Source: §4.4 · p. 9

5 CONCLUSION

METTA separates fast adaptation into within-episode online memory updates and cross-episode evidence aggregation, with linear sequence complexity. Across the reported benchmarks it improves adaptation, with especially strong gains on unseen held-out Meta-World task families; the evaluated setting uses fixed-task trials.

5 CONCLUSION

  • Within-episode online memory updates and cross-episode evidence aggregation let METTA use long interaction histories without full-history attention.This temporal hierarchy provides long-context adaptation with linear sequence complexity.Source: §5 · p. 9
  • METTA improves adaptation across HalfCheetahVel, Meta-World, and sparse-reward PointMaze while maintaining favorable computational efficiency.Gains were especially strong on unseen held-out Meta-World task families.Source: §5 · p. 9
  • Evaluation considers fixed-task trials; extending METTA to continual or nonstationary settings is future work.The paper suggests adding temporal or recency information to cross-episode memory for environments where task structure, dynamics, or rewards change over time.Source: §5 · p. 10

AI USE STATEMENT

The statement describes the authors’ use of generative AI tools, while the related-work section situates meta-RL among approaches to adaptation and long-context memory. It contrasts compressed task representations and quadratic-cost transformer attention with neural memory’s linear-time updates.

AI USE STATEMENT

  • Generative AI tools assisted with language polishing and experimental-code debugging, but not core research or experimental design.The authors state that they verified the technical decisions, experiments, analyses, and interpretations.Source: AI USE STATEMENT · p. 10

A RELATED WORK

  • Meta-RL trains over task variation so agents can infer task structure and adapt behavior from limited interaction.The related-work discussion distinguishes this aim from conventional RL’s optimization of a policy for one environment.Source: A RELATED WORK · p. 14
  • Meta-learning approaches include optimization-, metric-, and memory-based methods; recurrent and latent context mechanisms are prominent in sequential, partially observable RL.Prior methods use recurrent hidden states, probabilistic task variables, context encoders, or gradient-based policy adaptation.Source: A RELATED WORK · p. 14
  • Transformer meta-RL supports long-range retrieval but incurs quadratic self-attention costs, whereas METTA uses neural memory with linear-time recurrent updates.Related work also describes test-time-training layers as online-updated learnable models and Titans as persistent neural memory for associative retrieval.Source: A RELATED WORK · pp. 14–15

B IMPLEMENTATION DETAILS

The implementation trains METTA with sequential PPO segments and reconstructs causal memory histories for optimization. It uses chunk-wise TTT updates for parallelism and separates learned slow parameters from episode-specific memories in training and meta-test evaluation.

B.1 SEQUENTIAL PPO OPTIMIZATION

  • Sequential PPO processes ordered trajectory segments chronologically, avoiding repeated reconstruction of arbitrary prefixes and roughly halving PPO update time in the implementation.In the experiments, each segment corresponds to one episode within a trial.Source: §B.1 · p. 15
  • For each PPO segment, preceding episodes and the current episode prefix are reprocessed with current slow parameters; gradients flow through the reconstructed history and TTT updates.The computation graph is discarded after each optimizer step, and history is recomputed for the next segment using updated parameters.Source: §B.1 · p. 15
  • PPO optimization uses GAE over the flattened trial trajectory, online observation normalization, running-return reward scaling, and a clipped policy objective.The behavior-policy log probability is held fixed across PPO epochs, while z_t is recomputed using current parameters and gradients propagate through the policy and memory components.Source: §B.1 · pp. 15–16

B.2 MORE ON TEST-TIME TRAINING

  • The TTT-MLP memory is instantiated as a residual two-layer MLP with GELU activation; its residual connection and normalization stabilize online updates.The MLP gives the memory more expressive capacity.Source: §B.2 · p. 16
  • Chunk-wise TTT evaluates gradients at the chunk’s initial memory, accumulates updates in sequence order, and passes the final memory to the next chunk.Within-chunk gradients can be computed in parallel, improving GPU utilization while preserving causal ordering across chunks.Source: §B.2 · p. 16

B.3 METTA ALGORITHM

  • METTA training and evaluation maintain episode-summary histories and process episode prefixes with TTT before computing actions; PPO updates slow outer-loop parameters.Episode-specific memories are initialized from learned memories, updated online, and discarded after each episode; meta-test evaluation freezes the slow parameters.Source: §B.3 · pp. 16–17

C BENCHMARK DETAILS

The benchmark details cover robotic manipulation, velocity adaptation, and sparse-reward navigation. Their task structures differ in what is observed and how information can carry across episodes.

C.1 META-WORLD

  • Meta-World evaluates METTA on three environment families; ML1 masks goal coordinates, while ML10 and ML45 expose them and keep a task family fixed within each trial.ML10 and ML45 vary object and goal configurations across episodes, so persistent information primarily concerns task family and interaction structure.Source: §C.1 · p. 19

C.2 MUJOCO

  • HalfCheetahVel tests adaptation to a target running speed sampled uniformly from [0, 2], with episodes of 100 timesteps.The task runs in the MuJoCo simulator.Source: §C.2 · p. 19

C.3 POINTMAZE

  • PointMaze uses a 7 × 7 maze with sparse reward only upon reaching the goal, which terminates the episode.Start and goal are at least 10 shortest-path steps apart; tasks remain fixed across episodes, with 100 training tasks and 24 held-out test tasks.Source: §C.3 · p. 19

D EXPERIMENTAL PROTOCOL AND BASELINE DETAILS

Experiments use benchmark-recommended baseline configurations, matched interaction budgets and task splits, and five random seeds. The hyperparameters map rollout and TTT settings to episode, trial, and memory dimensions.

D EXPERIMENTAL PROTOCOL AND BASELINE DETAILS

  • Baselines use repository-recommended configurations where available, with matched interaction budgets, task splits, and evaluation protocols; results average five seeds with ±1 SEM.Learning curves show the mean across seeds, with shaded regions indicating ±1 standard error of the mean.Source: D EXPERIMENTAL PROTOCOL AND BASELINE… · p. 20
  • All methods share an NVIDIA RTX 3090 Ti server for experiments and wall-clock measurements, while retaining method-specific implementations.The common hardware and evaluation environment supports computational comparisons.Source: D EXPERIMENTAL PROTOCOL AND BASELINE… · p. 20

E HYPERPARAMETERS

  • The hyperparameters num envs, rollout steps, and trial length correspond to parallel environments B, episode horizon H, and episodes per trial E.These parameters define rollout collection and sequential PPO optimization.Source: E HYPERPARAMETERS · p. 20
  • TTT settings specify representation size, layer count, chunk length, and memory model; PPO and actor-critic settings specify update and hidden-layer dimensions.The experiments use a two-layer MLP as the TTT memory model.Source: E HYPERPARAMETERS · p. 21

F COMPUTATIONAL COMPLEXITY AND INFERENCE EFFICIENCY

The section compares full-trial online complexity and measured inference costs, then examines alternative within-episode sequence architectures. METTA avoids timestep-history-dependent computation in its fixed-size TTT memory, while architecture choice affects test adaptation.

F COMPUTATIONAL COMPLEXITY AND INFERENCE EFFICIENCY

  • TrMRL stores O(ET) timestep-level history, while ECET’s full-trial computation is O(E^2T^3 + E^3T), or O(T^3 + e^2T) per episode when episode computations are parallelized.The section distinguishes per-decision complexity from total computation accumulated across an online trial.Source: F COMPUTATIONAL COMPLEXITY AND… · p. 21
  • METTA’s fixed-size within-episode TTT memory yields O(1) history-dependent computation per timestep; attention over episode summaries gives O(E^2T) full-trial computation, while EMA gives O(ET) computation and O(1) history-dependent memory.The complexity analysis treats network dimensions and layer counts as constants.Source: F COMPUTATIONAL COMPLEXITY AND… · pp. 21–22, Table 3 · p. 22
  • At T = 500, TTT-MLP takes 2.16 ms per step and 5.03 MiB; TTT-Linear takes 1.62 ms and 2.40 MiB, with lower costs than ECET at E = 25.TTT-Linear has 65.2× lower latency and 42.2× lower memory than ECET with E = 25 at T = 500; TTT has higher constant cost at short contexts but scales better as history grows.Source: Table 4 · p. 22, F COMPUTATIONAL COMPLEXITY AND… · p. 22

G.1 ALTERNATIVE NEURAL MEMORIES

  • On PointMaze, TTT achieves the strongest and most consistent test performance; on ML10, TTT and the linear Transformer generalize better than GRU and Mamba.METTA can use different within-episode sequence architectures, but their test-time adaptation and generalization differ substantially.Source: §G.1 · p. 22, Figure 9 · p. 23

G.2 CROSS-EPISODE AGGREGATION

The ablation compares cross-episode aggregation methods while holding the within-episode memory and the rest of METTA fixed. Aggregation matters more at test time, with softmax attention offering the strongest overall trade-off across the evaluated settings.

G.2 CROSS-EPISODE AGGREGATION

  • The ablation holds within-episode neural memory and the rest of METTA unchanged while comparing softmax attention, no aggregation, EMA (α = 1/2), and linear attention.Aggregation has relatively little effect on training performance but becomes substantially more important at test time.Source: §G.2 · p. 23
  • On ML10, no cross-episode aggregation remains around 0.3 test success; on PointMaze, it stays near zero for approximately the first 0.5M environment steps.The aggregation-based variants reuse information across episodes and begin improving earlier on PointMaze.Source: §G.2 · p. 24
  • Softmax attention provides the strongest overall trade-off across the four settings, while EMA and linear attention are competitive but less consistent on held-out tasks.EMA is less stable on PointMaze and somewhat weaker on ML10 training tasks; linear attention closely tracks softmax during training but is less consistent on held-out tasks.Source: §G.2 · p. 24

G.3 ADDITIONAL ONLINE MEMORY UPDATE RESULTS

The ablations show that online memory updates, memory capacity, and separating state from task context affect held-out adaptation. PointMaze trajectories and exploration measures characterize how agents reuse information across episodes.

G.3 ADDITIONAL ONLINE MEMORY UPDATE RESULTS

  • Online memory updates yield a substantially larger test-performance gap on PointMaze than during training; ML10 variants remain close, while updates improve ML1 Push test performance.Both ML1 Push variants learn the training task effectively.Source: §G.3 · p. 24

G.4 MEMORY CAPACITY

  • TTT-MLP outperforms TTT-Linear on held-out tasks, most clearly on PointMaze, despite similar training success; the gap is smaller on ML10 and ML45.TTT-MLP uses a two-layer MLP memory, whereas TTT-Linear uses a single linear layer; the latter remains a lower-cost alternative.Source: Figure 13 · p. 25, §G.4 · p. 25

G.5 SEPARATING STATE AND TASK CONTEXT

  • Separating the current state from task context improves held-out generalization, especially on PointMaze, while the variants are closer on ML45.The ablation removes the separate state input, requiring learned context to encode both state and task information.Source: §G.5 · p. 25, Figure 14 · p. 26

G.7 POINTMAZE EXPLORATION METRICS

  • PointMaze exploration is measured by maze-cell coverage, consecutive-episode Jaccard overlap, and environment steps before success or termination.Lower Jaccard values mean less overlap and less repetitive exploration; fewer steps indicate broader exploration without longer trajectories when considered with success and coverage.Source: §G.7 · p. 26

G.6 QUALITATIVE POINTMAZE TRAJECTORIES

  • METTA searches more broadly in Episode 1, then later episodes behave more directly toward the goal; ECET and TrMRL often revisit similar regions across episodes.The trajectory examples cover Episodes 1, 2, 5, and 10 at mid-training and end-training stages.Source: §G.6 · pp. 25–26, Figure 15 · p. 27

Related papers

Loading 2610.31656v1…