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

Problem Statement

A central challenge in reinforcement learning (RL) is developing models that generalize beyond the specific tasks on which they are trained. While traditional multi-task and meta-RL approaches aim to address this, recent work has explored using transformer architectures for in-context learning, where models adapt to new tasks via context without explicit parameter updates. However, a principled understanding of how and why transformers generalize across tasks—and specifically across different domains—remains limited. Existing theoretical frameworks, such as those by Wang et al. (2025), rely on strong linearity assumptions (linear Markov Reward Processes and linear transformers), which restrict their applicability to complex, realistic settings where value functions are non-linear. This paper addresses the gap in understanding the mechanisms of in-context RL, specifically focusing on policy evaluation, to determine if a single transformer with fixed weights can generalize across multiple domains.

Methodology

The authors approach the problem from a kernel-based perspective, establishing a connection between non-linear transformers and kernel-based temporal difference (TD) learning.

Theoretical Framework

1. Transformer as a Functional Operator: The paper interprets the transformer as a functional operator that maps a context (a sequence of transitions) to a task-specific function.
2. RKHS Interpretation: Building on the observation that non-linear attention mechanisms induce a kernel function (Cheng et al., 2024), the authors posit that non-linear transformers perform regression in a Reproducing Kernel Hilbert Space (RKHS) induced by this kernel.
3. Kernel TD Update Construction: The core theoretical contribution is an explicit construction showing that a depth-$L$ non-linear transformer with fixed weights can implement $L$ iterations of a kernelized semi-gradient TD update purely through its forward pass.
   • Input Representation: The input matrix $Z_0$ encodes context transitions (states, rewards) and a query state.
   • Mechanism: The transformer uses context states as kernel centers. Each transformer layer corresponds to one step of the TD update. Specifically, the attention mechanism computes weighted sums of TD residuals ($\delta_t = r_t + \gamma V(s_{t+1}) - V(s_t)$) weighted by the kernel function $\kappa(s_j, s_i)$.
   • Weight Construction: The authors provide explicit constructions for the Key ($K$), Query ($Q$), and Value ($V$) projection matrices for a two-head attention block. One head computes the subtraction of the current value estimate, and the other computes the bootstrapped target, effectively performing the update:
     $$v^{(i)}_{k+1} = v^{(i)}_k + \alpha_k \sum_{j=0}^{n-1} \delta^{(j)}_k \kappa(s^{(j)}, s^{(i)})$$
   • Output: The value of the query state is read out from the bottom-right entry of the output matrix after $L$ layers.

Limitations and Conditions

The paper identifies specific conditions and limitations for this construction:

• Shared RKHS: Cross-domain generalization is theoretically guaranteed only if the value functions of the different domains lie within the same RKHS induced by the transformer's fixed kernel. If domains require fundamentally different kernel characteristics (e.g., different bandwidths for smooth vs. sharp value landscapes), a single fixed kernel cannot represent both well.
• Constant Offset: The construction introduces a state-independent constant bias in absolute value estimates. However, the authors argue this does not affect policy improvement (e.g., policy gradients) or TD error convergence, as the offset cancels out in relative comparisons and error calculations.
• Fixed Kernel Parameters: The kernel parameters (e.g., temperature $\delta$) are fixed and not adapted in-context. This restricts the model's ability to handle tasks requiring different kernel properties.

Key Contributions

1. Kernel View of In-Context Policy Evaluation: The paper frames in-context policy evaluation through the lens of kernel-based RL, relating transformer inference to kernelized TD updates.
2. Explicit Weight Construction: It provides a theoretical construction demonstrating that a fixed-weight non-linear transformer can implement iterative kernel TD updates, using context states as kernel centers.
3. Sufficient Condition for Cross-Domain Reuse: It establishes that cross-domain generalization is possible if the value functions of the tasks lie within a shared RKHS.
4. Empirical Validation: The authors instantiate this construction on real-world tasks, validating the theoretical predictions.

Experimental Results

The authors evaluated their method on MetaWorld (Yu et al., 2021), a benchmark comprising diverse manipulation domains (e.g., Pick-Place, Shelf-Place, Button-Press).

• One-Domain Generalization: A transformer with fixed weights (trained only on high-noise tasks within a domain) successfully converged on TD loss for medium and low-noise tasks within the same domain and across different domains (Pick-Place, Pick-Place-Wall, Shelf-Place, Plate-Slide).
• Cross-Domain Generalization: Models trained on one domain (e.g., Pick-Place-v3) were applied to other domains without fine-tuning. The model successfully generalized across Pick-Place, Pick-Place-Wall, Shelf-Place, and Plate-Slide. However, it failed to transfer to Button-Press-v3.
  • Analysis: This failure aligns with the theoretical limitation regarding the shared RKHS. Hyperparameter analysis revealed that Button-Press requires a significantly different kernel temperature (activation parameter) compared to the other four domains, implying its value function lies in a different RKHS.
• Linear vs. Non-Linear: Linear transformers (as proposed by Wang et al., 2025) failed to converge on MetaWorld tasks, confirming that non-linearity is necessary to capture the complex value function structures.
• Synthetic Environment: In a controlled synthetic MRP where the ground-truth value function was explicitly constructed to lie within the induced RKHS, the transformer accurately recovered the shape of the value function, validating the kernel regression mechanism.

Significance and Claims

The paper claims to provide the first theoretical construction demonstrating that in-context RL with non-linear transformers can generalize across domains. By interpreting the transformer as performing kernel regression, the authors offer a principled explanation for why a single set of weights can learn across multiple domains: it is possible when the tasks share a common RKHS.

The work serves as a foundational step toward developing more universal policies. It clarifies that while fixed-weight transformers can generalize, this capability is bounded by the expressiveness of the induced kernel. The authors explicitly state that their focus is on policy evaluation as a necessary foundational step, leaving policy improvement and adaptive kernel structures for future work. They do not claim to have solved the generalization problem for all RL tasks, but rather to have established the theoretical conditions under which such generalization is feasible.