Reinforcement Learning Foundation Models Should Already Be A Thing

Abstract: Foundation models for language and vision are powered by internet-scale data, while structured domains such as tabular prediction are powered by synthetic data. This substitute shifts the challenge from collection to prior design. Such priors already exist for many structured tasks: TabPFN and its successors solve tabular classification with a transformer pretrained on a synthetic Bayesian prior. We make two points. \textbf{First}, reinforcement learning is the conspicuous gap: sampling a synthetic MDP is as feasible as sampling a synthetic tabular dataset, yet no in-context RL work treats prior design as a primary objective. \textbf{Second}, MDPs admit a fixed-size sufficient statistic, independent of the episodes observed and tabular in shape, which makes them directly amenable to the attention-based architectures used for tabular foundation models, with a policy head replacing the supervised target. Together these define the agenda for an RL foundation model. As a proof of concept, we train a Graph Attention Network entirely on synthetic MDPs and show that, with no task-specific tuning, it solves held-out tabular benchmarks in context, both online and offline: online, in far fewer episodes than UCB-VI and tabular Q-learning, and offline, competitively with VI-LCB.

• The paper introduces a novel RL foundation model that leverages a tabular sufficient-statistic matrix as a universal MDP prior for efficient policy estimation.
• It employs a Graph Attention Network with unrolled propagation layers to adaptively trade computation for improved policy inference on tabular benchmarks.
• Empirical results show the model outperforms conventional RL algorithms in both online and offline settings, underscoring the role of synthetic prior design.

Reinforcement Learning Foundation Models: The Case for Immediate Adoption

Motivation and Context

The foundational model paradigm—pretraining over vast, diverse distributions to enable strong in-context generalization—has transformed language and vision domains. However, reinforcement learning (RL) has not yet fully leveraged the foundation model approach, particularly with respect to explicit synthetic prior design. While modern architectures for structured data, such as TabPFN, have demonstrated that expressive priors over Bayesian classifiers can yield competitive tabular models, the analogous development for RL remains largely absent. This paper addresses this gap by proposing and experimentally validating a straightforward yet potent approach to RL foundation models, emphasizing that prior design for Markov Decision Processes (MDPs) should become a central concern in the field.

Critique of the Current RL Meta-Learning Landscape

Contemporary in-context RL and meta-RL predominantly rely on trajectory data sampled from handpicked or procedurally generated environments. This results in models heavily tailored to narrow task suites, without consideration for an explicit prior over MDPs. The paper argues that this focus misaligns with the principle of foundation models, wherein the distributional prior is the critical enabler of generalization. Notably, prior Bayesian RL research (Strens 2000, Dearden et al. 1998) did emphasize structured priors over transitions and rewards, but this thread has not been systematically incorporated in recent deep meta-RL, Decision Transformers, or sequence-modeling strategies.

A strong criticism is also levied against sequence-based models, which encounter practical limitations: vanishing/exploding gradients in RNNs and substantial memory and compute scaling in transformer-based models, where attention complexity is quadratic in context length. The result is that key exploratory experience may be lost during truncation or compressed ineffectively. The authors note that while latent compression has shown promise in the tabular supervised domain, scaling trajectory-based methods to challenging RL setups remains unresolved.

Tabularization of MDP Sufficient Statistics

The paper's central insight is the identification of the fixed-size, tabular sufficient-statistic matrix $Z$ for small finite MDPs. Each row corresponds to a state-action pair and encodes (1) empirical visit counts, (2) empirical mean rewards, and (3) empirical state transition probabilities. Notably, this matrix is invariant to the number of episodes observed and episode length, and for small MDPs, $Z$ can serve as a universal, information-complete representation.

This reframing allows the use of attention-based, permutation-equivariant set architectures (such as those developed for tabular data) to process RL data, replacing the supervised head with a policy output. The authors choose a Graph Attention Network (GAT) that explicitly utilizes the empirical dynamics as attention biases, enabling rapid, structure-aware planning for small MDPs.

Model Architecture and Training Protocol

The model takes, for each state-action pair, log-visit count, normalized empirical reward, and empirical transition probabilities. This is processed by a GAT with $K$ weight-tied propagation layers, each implementing a form of message-passing akin to value iteration. Importantly, these layers can be unrolled to greater depths at test time, trading computation for improved policy estimation, without retraining.

The model is trained offline. For each batch, multiple MDPs are sampled from a broad synthetic prior: number of states, actions, transition sparsity, latent geometry, stochasticity, and reward sparsity are all randomized. For each MDP, the optimal policy is computed; then, noisy empirical statistics are produced via simulated exploration, and the model is supervised to predict the Boltzmann-optimal policy using these statistics. The training objective is per-state cross-entropy between the predicted and optimal policy.

Empirical Results

Key empirical findings:

• On GridWorld and FrozenLake benchmarks, without any environment-specific tuning or data, the in-context RL foundation model rapidly achieves strong performance.
• In online RL, the model converges to the optimal policy in substantially fewer episodes than both tabular Q-learning and UCB-VI, a model-based algorithm with upper confidence bounds.
• In offline RL, when evaluated from fixed datasets collected with uniform random policies, the model is competitive with pessimistic value iteration (VI-LCB), achieving optimal or near-optimal policies with minimal data on deterministic benchmarks. On coverage-limited stochastic environments, the model’s reliance on the prior provides robustness where VI-LCB is brittle to data sparsity.

Limitations and Future Directions

The current prototype is limited to small tabular MDPs with discrete state and action spaces. While permutation equivariance confers invariance to arbitrary relabeling, it also removes any notion of state geometry, meaning the model cannot generalize across similar but unvisited states unless further state features are provided. This suggests that extensions integrating informative features or latent geometries into the prior and the input are a natural next step.

The tabular representation itself is not directly scalable to large or continuous state-action spaces. The authors propose future directions employing feature-indexed rows and continuous action heads (e.g., parameterizing policies with squashed Gaussians for continuous control tasks), but note that scale generalization will require more sophisticated priors and representations. Furthermore, extending beyond Markov environments is non-trivial, as partial observability invalidates the sufficiency of $Z$.

Amortizing computation is another open problem. While updating $Z$ is efficient, current evaluation necessitates full recomputation for each step—incremental updates would further enhance online applicability.

Theoretical and Practical Implications

This work argues strongly that RL foundation models should be built now, not post hoc as a consequence of advances in separate research threads. The central claim is that foundation models in RL can and should be trained using synthetic but well-designed priors over MDPs, analogously to the supervised tabular domain, because the sufficiency and invariance of the MDP statistic $Z$ endows them with broad generalization at moderate cost. This perspective fundamentally reframes the agenda for RL foundation models: the synthetic prior, not just architecture or data curation, is the primary axis for future progress.

If this approach is extended to large/continuous domains and robust feature priors are developed, it would enable RL agents to exhibit foundation-model-like properties—rapid adaptation across tasks, robustness to data sparsity, and efficient policy inference—long sought after in meta-RL and continual learning.

Conclusion

This paper provides a precise analysis and compelling prototype for RL foundation models based on synthetic MDP priors and permutation-equivariant attention architectures. It demonstrates state-of-the-art data efficiency on canonical tabular RL benchmarks, outperforming traditional algorithms with zero tuning. Theoretical considerations and empirical results underline the necessity of principled prior design in in-context RL, setting a clear agenda for future research: the systematic development and evaluation of synthetic-prior-fitted RL foundation models.