Optimal Dynamic Regret by Transformers for Non-Stationary Reinforcement Learning

TL;DR

This paper addresses learning in non-stationary reinforcement learning settings by showing that transformers trained in an in-context manner can achieve near-optimal dynamic regret. It introduces a transformer architecture that can approximate non-stationary operation schemes and provides a regret bound decomposed into approximation, pretraining, and base algorithm components, with a corollary giving rate-optimal performance under sufficient data. The core contributions include a new proof technique that interprets the transformer's hidden state as maintaining multiple hypothesis policies, a demonstration that windowed scheduling and restart mechanisms can be effectively emulated by a transformer, and empirical results in a linear bandit setting showing competitive performance with established expert algorithms. The work highlights the practical significance of in-context learning for adaptive RL and lays groundwork for extending to broader non-stationary environments and richer benchmarks.

Abstract

Transformers have demonstrated exceptional performance across a wide range of domains. While their ability to perform reinforcement learning in-context has been established both theoretically and empirically, their behavior in non-stationary environments remains less understood. In this study, we address this gap by showing that transformers can achieve nearly optimal dynamic regret bounds in non-stationary settings. We prove that transformers are capable of approximating strategies used to handle non-stationary environments and can learn the approximator in the in-context learning setup. Our experiments further show that transformers can match or even outperform existing expert algorithms in such environments.

Figures (9)

• Figure 1: Overview of the proof structure for the approximation analysis (Theorem \ref{['thm-approx']}). The blue box represents the algorithmic operations for handling non-stationarity, while the green box represents the transformer approximating the algorithm. Each block or sub-architecture of the transformer corresponds to a specific operation of the algorithm.
• Figure 2: Illustration of WS when $n=3$. Purplish blocks represent instances scheduled by $\sigma_1$, while reddish blocks represent the active instances selected by $\sigma_2$. Reddish blocks connected by a dashed line are concatenated.
• Figure 3: Cumulative regret comparison for LinUCB, Thompson Sampling (TS), MASTER+LinUCB/TS, and transformer (TF) in linear bandits $d = 32$, $A = 10$. The first row corresponds to Low Non-Stationarity environments, while the second row shows High Non-Stationarity environments. Shading indicates the standard deviation of the regret estimates.
• Figure 4: Suboptimality Comparisons of LinUCB, Thompson Sampling (TS), MASTER+LinUCB/TS, and transformer (TF) in Low Non-stationary environments. Linear bandit with $d = 32$, $A = 10$. Shading indicates the standard deviation of the regret estimates.
• Figure 5: Suboptimality Comparisons of LinUCB, Thompson Sampling (TS), MASTER+LinUCB/TS, and transformer (TF) in High Non-stationary environments ($b=0.018$). Linear bandit with $d = 32$, $A = 10$. Shading indicates the standard deviation of the regret estimates.

Theorems & Definitions (15)

• Definition 1: Non-stationary measure
• Definition 2: Covering number
• Definition 3: Distribution ratio
• Theorem 1
• Corollary 2
• Proposition 3: Approximating WS
• Proposition 4: Approximating RM
• Theorem 5
• Theorem 6
• Lemma 7: Regret bound for multiple instances