From Memories to Maps: Mechanisms of In-Context Reinforcement Learning in Transformers (2506.19686v2)

Abstract: Humans and animals show remarkable learning efficiency, adapting to new environments with minimal experience. This capability is not well captured by standard reinforcement learning algorithms that rely on incremental value updates. Rapid adaptation likely depends on episodic memory -- the ability to retrieve specific past experiences to guide decisions in novel contexts. Transformers provide a useful setting for studying these questions because of their ability to learn rapidly in-context and because their key-value architecture resembles episodic memory systems in the brain. We train a transformer to in-context reinforcement learn in a distribution of planning tasks inspired by rodent behavior. We then characterize the learning algorithms that emerge in the model. We first find that representation learning is supported by in-context structure learning and cross-context alignment, where representations are aligned across environments with different sensory stimuli. We next demonstrate that the reinforcement learning strategies developed by the model are not interpretable as standard model-free or model-based planning. Instead, we show that in-context reinforcement learning is supported by caching intermediate computations within the model's memory tokens, which are then accessed at decision time. Overall, we find that memory may serve as a computational resource, storing both raw experience and cached computations to support flexible behavior. Furthermore, the representations developed in the model resemble computations associated with the hippocampal-entorhinal system in the brain, suggesting that our findings may be relevant for natural cognition. Taken together, our work offers a mechanistic hypothesis for the rapid adaptation that underlies in-context learning in artificial and natural settings.

• The paper demonstrates that meta-trained transformers rapidly adapt to novel tasks, achieving near-optimal performance after a single reward exposure.
• The paper reveals that structured, geometry-aligned representations emerge from episodic memory mechanisms, enabling efficient planning in gridworlds and hierarchical mazes.
• The paper shows that transformers use memory tokens as dynamic computational workspaces, outperforming standard model-free and model-based RL methods.

Mechanisms of In-Context Reinforcement Learning in Transformers: From Memories to Maps

This paper investigates the mechanisms by which transformer architectures, when meta-trained for in-context reinforcement learning (RL), develop rapid adaptation strategies in novel environments. The paper draws inspiration from biological episodic memory systems, particularly the hippocampal-entorhinal circuit, and explores how transformers can leverage memory not only for storage but as an active computational substrate. The analysis is grounded in two classes of planning tasks—Euclidean gridworlds and hierarchical tree mazes—designed to probe the flexibility and generalization of learned RL strategies.

Experimental Framework

The authors employ a decision-pretrained transformer (DPT) architecture, following the framework of Lee et al. (2023), to perform in-context RL. The model is trained on a distribution of Markov Decision Processes (MDPs), each defined by unique state encodings, transition dynamics, and reward locations. For each task, the model receives a context dataset of exploratory trajectories and is required to predict the optimal action from a query state, relying solely on in-context information at test time. Notably, the context and query tokens are ordered such that the model processes past experiences before the current decision, aligning with the interpretation of episodic memory retrieval.

Two task distributions are used:

• Gridworlds: $5 \times 5$ spatially regular environments with fixed transition dynamics but varying state encodings and reward locations.
• Tree Mazes: Hierarchically structured, probabilistically branching binary trees with sparse rewards, designed to challenge compositional and hierarchical generalization.

Key Findings

1. Rapid In-Context Adaptation and Behavioral Efficiency

The meta-learned transformer demonstrates rapid adaptation in both gridworld and tree maze tasks, achieving near-optimal performance after minimal reward exposure. In gridworlds, the agent often navigates directly to the reward after a single exposure, paralleling one-shot learning observed in animal studies. In tree mazes, the agent similarly exhibits efficient learning, outperforming both tabular Q-learning and deep Q-network (DQN) baselines, even when these baselines are trained to convergence on the same in-context data.

A notable behavioral signature is the agent's ability to infer and execute shortcut paths to the reward, even when only circuitous trajectories are observed in context. The model selects optimal shortcuts in over 60% of test simulations, compared to 2% under a chance policy.

2. Emergence of Structured Representation Learning

Analysis of internal representations reveals that the transformer develops structured, geometry-aligned embeddings of the environment through in-context experience:

• In gridworlds, as context length increases, the model's representations organize to reflect the latent Euclidean geometry, as evidenced by principal component analysis and increased kernel alignment with the true environment structure.
• In tree mazes, representations form bifurcating structures that mirror the hierarchical layout, with kernel alignment increasing with context but remaining coarser than in gridworlds.

These findings indicate that the model meta-learns in-context structure learning strategies, facilitating efficient generalization and decision-making.

3. Cross-Context Alignment of Representations

The transformer aligns internal representations across environments with shared latent structure, despite differences in sensory input. States occupying the same position in different gridworlds or tree mazes are encoded similarly, as shown by increased cross-environment correlation of node representations with context length. This cross-context alignment is strongest for states at the periphery (edges or leaves) of the environment, consistent with theories of abstraction in the entorhinal cortex.

4. RL Strategies Beyond Model-Free and Model-Based Taxonomies

Mechanistic analysis demonstrates that the emergent RL strategies do not conform to standard model-free (value-based) or model-based (planning) frameworks:

• Model-Free RL: While value functions can be linearly decoded from representations, the decoded values primarily reflect spatial proximity to the goal rather than action-contingent reward prediction. In tree mazes, value gradients are too localized to support long-range planning.
• Model-Based RL: Attribution analyses (integrated gradients) and attention ablations show that the model's decisions depend primarily on memory tokens near the query and goal states, with little influence from intermediate path states. This is inconsistent with explicit path planning, which would require integrating transitions along the full route.

5. Memory as a Computational Workspace

The transformer leverages its memory tokens not only to store raw experiences but also to cache intermediate computations critical for decision-making. The specific strategies differ by task:

• Gridworlds: The model aligns state representations to Euclidean space, computes the angle between the query and goal states, and selects actions accordingly. Decoding analyses confirm that both position and goal-relative angle are encoded in memory tokens.
• Tree Mazes: The model tags context-memory tokens that are on the critical left-right (L-R) path from root to reward. At decision time, it checks if the query state is on the L-R path and, if so, extracts the optimal action from the tagged tokens; otherwise, it defaults to a parent-node transition. Decoding confirms that both L-R path membership and inverse actions are encoded in memory tokens.

Numerical Results and Contradictory Claims

• Performance: The meta-learned transformer achieves near-maximal return after a single reward exposure in both task domains, substantially outperforming Q-learning and DQN baselines.
• Shortcut Behavior: The model selects optimal shortcut paths in over 60% of gridworld test cases, a strong deviation from chance.
• Representation Alignment: Kernel alignment and cross-context correlation metrics quantitatively demonstrate the emergence of structured, generalizable representations.

A contradictory claim is that, despite the model's rapid adaptation and apparent planning behavior, its internal mechanisms do not align with either model-free or model-based RL, challenging the sufficiency of these traditional taxonomies for describing flexible, memory-based learning in artificial agents.

Implications and Future Directions

The findings have several implications for both AI and neuroscience:

• AI Systems: The results suggest that transformer architectures, when meta-trained for in-context RL, can develop flexible, memory-based strategies that support rapid adaptation in novel environments. This has practical relevance for designing agents capable of few-shot generalization and efficient exploration in complex, structured domains.
• Neuroscience: The parallels between the model's emergent representations and computations associated with the hippocampal-entorhinal system provide a computational hypothesis for episodic memory's role in rapid learning and decision-making.
• Theory: The work motivates a broader conceptualization of RL strategies, emphasizing the role of memory as an active computational resource rather than passive storage. It also highlights the need for new taxonomies and analytical tools to characterize the diversity of algorithms that can emerge in meta-learned systems.

Future research may explore:

• Scaling these findings to more complex, real-world environments and tasks.
• Investigating the interplay between episodic and semantic memory in in-context learning.
• Developing architectures or training regimes that further enhance the computational utility of memory.
• Applying similar mechanistic analyses to other domains, such as language or vision, to assess the generality of these strategies.

Conclusion

This paper provides a detailed mechanistic account of how transformers, when meta-trained for in-context RL, develop memory-based strategies that enable rapid adaptation and generalization. The emergent algorithms leverage memory as a computational workspace, supporting flexible behavior beyond the scope of standard RL frameworks. These insights have significant implications for the design of adaptive AI systems and for understanding the computational principles underlying natural cognition.