# One for All: A Non-Linear Transformer can Enable Cross-Domain Generalization for In-Context Reinforcement Learning

arXiv: https://arxiv.org/abs/2605.09727 · HTML: https://mlsanity.com/paper/2605.09727/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Deep RL algorithms and representations, one of 15 directions in Deep reinforcement learning methods.
- Builds on 15 of its 25 references that are in the map.
- First posted on arXiv 10 May 2026.

## Placed in

- Deep reinforcement learning methods / [Deep RL algorithms and representations](https://mlsanity.com/deeprl/deep-rl-algorithms-and-representations/)

## Builds on (most cited first, 15 of 15)

- [An Explanation of In-context Learning as Implicit Bayesian Inference](https://mlsanity.com/paper/2111.02080/) (2021)
- [Transformers learn in-context by gradient descent](https://mlsanity.com/paper/2212.07677/) (2022)
- [What Can Transformers Learn In-Context? A Case Study of Simple Function Classes](https://mlsanity.com/paper/2208.01066/) (2022)
- [What learning algorithm is in-context learning? Investigations with linear models](https://mlsanity.com/paper/2211.15661/) (2022)
- [XCiT: Cross-Covariance Image Transformers](https://mlsanity.com/paper/2106.09681/) (2021)
- [Transformers learn to implement preconditioned gradient descent for in-context learning](https://mlsanity.com/paper/2306.00297/) (2023)
- [Transformer Dissection: A Unified Understanding of Transformer's Attention via the Lens of Kernel](https://mlsanity.com/paper/1908.11775/) (2019)
- [One Step of Gradient Descent is Provably the Optimal In-Context Learner with One Layer of Linear Self-Attention](https://mlsanity.com/paper/2307.03576/) (2023)
- [Large Language Models Are Latent Variable Models: Explaining and Finding Good Demonstrations for In-Context Learning](https://mlsanity.com/paper/2301.11916/) (2023)
- [Supervised Pretraining Can Learn In-Context Reinforcement Learning](https://mlsanity.com/paper/2306.14892/) (2023)
- [A Theory of Emergent In-Context Learning as Implicit Structure Induction](https://mlsanity.com/paper/2303.07971/) (2023)
- [Transformers as Decision Makers: Provable In-Context Reinforcement Learning via Supervised Pretraining](https://mlsanity.com/paper/2310.08566/) (2023)
- [Transformers Implement Functional Gradient Descent to Learn Non-Linear Functions In Context](https://mlsanity.com/paper/2312.06528/) (2023)
- [Emergent Agentic Transformer from Chain of Hindsight Experience](https://mlsanity.com/paper/2305.16554/) (2023)
- [Improving Transformers with Probabilistic Attention Keys](https://mlsanity.com/paper/2110.08678/) (2021)

## Abstract (arXiv)

A central challenge in reinforcement learning (RL) is to learn models that generalize beyond the tasks on which they are trained, a goal traditionally pursued through multi-task and meta RL. Recently, transformer architectures have emerged as a promising approach, enabling adaptation to new tasks via in-context learning without explicit parameter updates. From a functional perspective, a transformer can be viewed as a functional operator that maps a context to a task-specific function. It is thus fundamental to understand and design this operator to support stronger generalization in RL. In this work, we address this resulting question of generalization from a kernel-based perspective by establishing a connection between non-linear transformers and kernel-based temporal difference learning. By interpreting the transformer as performing regression in a Reproducing Kernel Hilbert Space (RKHS), we show that value functions from different domains can be represented using a shared set of weights, provided they lie within the same RKHS. Experiments on multiple MetaWorld domains support this interpretation, demonstrating convergence of the temporal-difference objective.