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

• URL: http://arxiv.org/abs/2605.09727v1
• Date: Sun, 10 May 2026 19:52:29 GMT
• Title: One for All: A Non-Linear Transformer can Enable Cross-Domain Generalization for In-Context Reinforcement Learning
• Abstract summary: A central challenge in reinforcement learning (RL) is to learn models that generalize beyond the tasks on which they are trained.<n>Recently, transformer architectures have emerged as a promising approach, enabling adaptation to new tasks via in-context learning.<n>In this work, we establish a connection between non-linear transformers and kernel-based temporal difference learning.
• Score: 12.672720180188348
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: 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.

Related papers

• Ghost in the Kernel: In-Context Learning with Efficient Transformers via Domain Generalization [17.68725202024536]
  Transformer-based large models have demonstrated remarkable generalization abilities across different tasks.<n>With richer context, transformers adapt more effectively to the current use case without any parameter updates.<n>However, the quadratic computational and memory complexity with respect to context length significantly slows data processing in softmax transformers.
  arXiv  Detail & Related papers  (2026-07-01T06:03:14Z)
• Heuristic Transformer: Belief Augmented In-Context Reinforcement Learning [1.8791091507292152]
  Heuristic Transformer (HT) is an in-context reinforcement learning approach that augments the in-context dataset with a belief distribution over rewards to achieve better decision-making.<n>We show that HT consistently surpasses comparable baselines in terms of both effectiveness and generalization.<n>Our method presents a promising direction to bridge the gap between belief-based augmentations and transformer-based decision-making.
  arXiv  Detail & Related papers  (2025-11-13T12:32:36Z)
• Dita: Scaling Diffusion Transformer for Generalist Vision-Language-Action Policy [73.75271615101754]
  We present Dita, a scalable framework that leverages Transformer architectures to directly denoise continuous action sequences.<n>Dita employs in-context conditioning -- enabling fine-grained alignment between denoised actions and raw visual tokens from historical observations.<n>Dita effectively integrates cross-embodiment datasets across diverse camera perspectives, observation scenes, tasks, and action spaces.
  arXiv  Detail & Related papers  (2025-03-25T15:19:56Z)
• Interpreting Affine Recurrence Learning in GPT-style Transformers [54.01174470722201]
  In-context learning allows GPT-style transformers to generalize during inference without modifying their weights. This paper focuses specifically on their ability to learn and predict affine recurrences as an ICL task. We analyze the model's internal operations using both empirical and theoretical approaches.
  arXiv  Detail & Related papers  (2024-10-22T21:30:01Z)
• In-Context Learning with Representations: Contextual Generalization of Trained Transformers [66.78052387054593]
  In-context learning (ICL) refers to a capability of pretrained large language models, which can learn a new task given a few examples during inference. This paper investigates the training dynamics of transformers by gradient descent through the lens of non-linear regression tasks.
  arXiv  Detail & Related papers  (2024-08-19T16:47:46Z)
• Linear Transformers with Learnable Kernel Functions are Better In-Context Models [3.3865605512957453]
  We present an elegant alteration to the Based kernel that amplifies its In-Context Learning abilities. In our work, we present a singular, elegant alteration to the Based kernel that amplifies its In-Context Learning abilities evaluated with the Multi-Query Associative Recall task.
  arXiv  Detail & Related papers  (2024-02-16T12:44:15Z)
• HCVP: Leveraging Hierarchical Contrastive Visual Prompt for Domain Generalization [69.33162366130887]
  Domain Generalization (DG) endeavors to create machine learning models that excel in unseen scenarios by learning invariant features. We introduce a novel method designed to supplement the model with domain-level and task-specific characteristics. This approach aims to guide the model in more effectively separating invariant features from specific characteristics, thereby boosting the generalization.
  arXiv  Detail & Related papers  (2024-01-18T04:23:21Z)
• End-to-End Meta-Bayesian Optimisation with Transformer Neural Processes [52.818579746354665]
  This paper proposes the first end-to-end differentiable meta-BO framework that generalises neural processes to learn acquisition functions via transformer architectures. We enable this end-to-end framework with reinforcement learning (RL) to tackle the lack of labelled acquisition data.
  arXiv  Detail & Related papers  (2023-05-25T10:58:46Z)
• Transformers are Meta-Reinforcement Learners [0.060917028769172814]
  We present TrMRL, a meta-RL agent that mimics the memory reinstatement mechanism using the transformer architecture. We show that the self-attention computes a consensus representation that minimizes the Bayes Risk at each layer. Results show that TrMRL presents comparable or superior performance, sample efficiency, and out-of-distribution generalization.
  arXiv  Detail & Related papers  (2022-06-14T06:21:13Z)
• A New Representation of Successor Features for Transfer across Dissimilar Environments [60.813074750879615]
  Many real-world RL problems require transfer among environments with different dynamics. We propose an approach based on successor features in which we model successor feature functions with Gaussian Processes. Our theoretical analysis proves the convergence of this approach as well as the bounded error on modelling successor feature functions.
  arXiv  Detail & Related papers  (2021-07-18T12:37:05Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.