What Must Generalist Agents Remember?

Paper Title: What Must Generalist Agents Remember?
Authors: Khurram Yamin, Namrata Deka, Maitreyi Swaroop, Albert Ting, Jeff Schneider, Bryan Wilder (Carnegie Mellon University, Georgia Institute of Technology)

Problem Statement

The paper addresses a fundamental question in the development of generalist Reinforcement Learning (RL) agents: what internal information must an agent store in memory to act near-optimally across multiple environments (domains) and diverse goals?

Unlike prior work that assumes the agent is given an explicit index or label for the current task or domain, this paper studies a setting where the agent receives no such index. The agent observes only its own trajectory (state-action history) and must infer the latent domain identity required for competent action selection. The authors frame this as a latent-domain Partially Observable Markov Decision Process (POMDP), where the hidden variable is the fixed domain index $\sigma$ which alters transition dynamics, while the reward mechanism remains shared across domains.

The core challenge is that two different domains may present the same observation (state) but require incompatible optimal actions. The paper asks: under what conditions is memory strictly necessary to resolve this ambiguity, and does a sufficiently competent memory allow the agent to reconstruct the local transition dynamics of the environment?

Methodology and Theoretical Framework

The authors formalize the interaction in discrete time with a family of goals $\Psi$ and a set of domains $\Sigma$. They introduce two primary theoretical results:

1. Representation Separation (Theorem 1)
The first result establishes a necessary condition for memory. The authors define a "shared bottleneck" scenario (Assumption 1) where two distinct domains require disjoint sets of near-optimal actions at the same state $s^\dagger$, separated by a significant value gap.

• Claim: If an agent is uniformly near-optimal across all goals and domains, its internal memory state $M_t$ must distinguish between these domains at the bottleneck state.
• Mechanism: The theorem proves that the conditional probability distributions of the memory state at the bottleneck, given the domain, must have a non-zero Total Variation (TV) distance. Specifically, if the agent's policy is $\epsilon$-near-optimal, the TV distance between the memory distributions of the two domains is bounded below by a function of the value gap, the probability of reaching the bottleneck, and the discount factor.
• Implication: Memory is not merely a sufficient statistic for optimal control (as in standard POMDP belief states) but is a necessary condition for learned representations to achieve uniform competence when hidden domain identity is action-relevant.

2. Decoder Existence (Theorem 2)
The second result investigates whether the memory required for separation is also sufficient to reconstruct the environment's dynamics.

• Setup: The authors introduce "probe goals" (an auxiliary family $\Psi_F$) that act as one-step predictive tests. They define "value-sufficient memory" as memory that can accurately predict the optimal action-values for these probe goals.
• Claim: If the memory is value-sufficient for a class of probe goals that can distinguish local transition kernels (via an integral probability metric), then there exists a decoder that can approximately reconstruct the local interventional transition kernel $P^\sigma(\cdot | s, a)$.
• Implication: A memory state that supports near-optimal planning for a rich set of goals inherently contains enough information to decode the local dynamics of the environment, enabling approximate planning even without an explicit world model provided by the environment.

Experimental Validation

The authors validate their theoretical claims using a custom environment called ForkWorld:

• Environment: A T-junction gridworld where a corridor leads to a fork (bottleneck). The domain is latent: in the "normal" domain, 'up' moves north and 'down' moves south; in the "swapped" domain, these mappings are inverted. The goal (reach top or bottom) is observable, but the domain is not.
• Bottleneck: At the fork, the optimal action depends entirely on the hidden domain. A memoryless agent cannot distinguish the domains and is forced to guess, capping performance at 50% success.
• Agents: The study compares memoryless agents (DQN), agents with explicit history windows (Stacked DQN), and recurrent agents (DRQN) against oracle baselines that know the domain.

Key Results:

1. Necessity of Memory: Memoryless agents failed to exceed the 50% baseline (achieving ~19%), confirming that without memory to separate domains, the policy cannot converge to a stable solution. Memory-equipped agents (Stacked DQN and DRQN) achieved success rates of 75–88%, closing the gap to the oracle.
2. Representation Separation: A linear probe trained to predict the domain from the agent's internal state at the bottleneck showed that memory-equipped agents achieved near-perfect domain classification (accuracy >97%) on subsequent visits to the fork, whereas memoryless agents remained at chance (50%). This empirically confirms Theorem 1.
3. Decodable Dynamics: The authors trained a decoder to predict the next state from the agent's memory representation. Memory-equipped agents significantly reduced transition prediction error compared to a baseline, whereas memoryless agents showed no improvement. This supports Theorem 2, demonstrating that the memory used for planning also encodes the local transition dynamics.

Key Contributions and Significance

The paper makes the following contributions:

1. Formal Necessity of Memory: It provides a control-theoretic proof that for generalist agents operating in latent domains, memory is not optional but a necessary condition for uniform near-optimality when domains share observations but require disjoint actions.
2. Memory as a Substrate for World Models: It demonstrates that the same memory representation required for action selection, if sufficiently rich to support value prediction for probe goals, inherently supports the decoding of local transition dynamics.
3. Shift in Analysis: Unlike previous "good regulator" or causal RL results that assume the agent is given a task index, this work analyzes the history-to-memory mapping, showing that the agent must implicitly learn to separate domains from trajectory data alone.

Significance:
The authors argue that these results characterize memory as the substrate that supports domain disambiguation, transition-model reconstruction, and planning for generalist agents. The work suggests that the ability to infer latent context and reconstruct local dynamics is not an architectural feature to be added, but a fundamental requirement for any agent that performs well across a family of shifting environments. The findings imply that successful generalist RL agents must internally encode a representation that separates latent domains, and that this encoding naturally facilitates the reconstruction of the environment's causal structure.