Optimal Dynamic Regret by Transformers for Non-Stationary Reinforcement Learning

We are grateful to the reviewer for their thorough review and valuable feedback on our paper. In the following, we hope to address each point raised in the review.

Why is the definition of the non-stationarity measure different from Wei and Luo?

We thank the reviewer for raising this important point. Our initial intention was to align with the learning setup in Lin et al. (2024), and thus defining the non-stationary measure as $\Delta(t) = \max_{s,a} | R_t (s,a) - R_{t+1} (s,a) |$. However, we inadvertently wrote it as $\Delta(t) = \max_{s,s',a,a'} | R_t (s,a) - R_{t+1} (s',a') |$. We have now corrected the definition to $\Delta(t) = \max_{s,a} | R_t (s,a) - R_{t+1} (s,a) |$, which is the upper bound of the one in Wei et al. It is worth noting that using $\Delta(t) = \max_{s,a} | R_t (s,a) - R_{t+1} (s,a) |$ instead of the definition in Wei et al. has minimal impact on the overall approximation, while offering better alignment with the setup in Lin et al.

Does using the “non-continuous” architecture complicate adapting the results of Lin et al (2024)?

No, using the "non-continuous" architecture does not complicate adapting the results of Lin et al. (2024). The non-continuous transformer is essentially equivalent to conventional transformers. In the proofs of Lin et al. (2024), a specific transformer was constructed to approximate each RL algorithm, demonstrating the existence of a transformer that can achieve the error bound. Similarly, we constructed a specific transformer to approximate the MASTER algorithm. While the non-continuous transformer may differ in terms of augmented inputs, the main impact of these inputs is simply increasing the model size. Since our work is focused on the theoretical approximation rather than practical considerations, this does not complicate adapting the results of Lin et al. (2024).

In fact, our initial proof used a conventional transformer without augmented inputs and relied on multiple heads to compensate. However, due to the complexity of the proof, we later switched to the augmented input strategy for simplification.

We appreciate the reviewer for highlighting this, and we will include a brief explanation of the non-continuous transformer in the revised version of the main text.

Should the transformer algorithm be trained to output both the auxiliary quantity and the actions? Or should it be given access to some other function computing auxiliary values? Or does it suffice for the transformer to be able to compute these values in some intermediate block?

Thank you for this insightful question. In conventional RL algorithms, auxiliary quantities (e.g., UCB scores, Q-values) are explicitly computed to inform action selection. In our transformer-based framework, we adopt a different but related strategy: the transformer is trained to output a probability distribution over actions, from which the action is sampled. This output serves a role of implicit encoding for the auxiliary information used to guide decision-making.

In other words, the auxiliary quantity is not a separate supervised target, but rather emerges from the transformer’s internal computation. Our theoretical construction (see Proposition 3 and Theorem 5) shows that the transformer can implement implicit mechanisms, such as window schedulers and restart tests, that rely on internal auxiliary statistics. Thus, it is sufficient for the transformer to internally represent and use such quantities, rather than outputting them explicitly or requiring access to an external function.

We appreciate the reviewer for raising this important point, and will add a brief explanation of this in the revised version to clarify how these auxiliary quantities are handled in the transformer framework.

The discussion of limitations seems inadequate. The authors should additionally discuss the discrepancies between the analyzed non-continuous transformer and the standard architecture. The limitations discussed in Lin et al (2024) also apply to this work and could be mentioned.

Thank you for your suggestions. We will expand the discussion of limitations in the revised version to address the important points you raised.

In particular, we will explicitly discuss the discrepancies between our analyzed non-continuous transformer and standard transformer architectures used in practice. Our theoretical construction relies on controlled masking and conditional restarts, which, while implementable in principle, do not directly correspond to standard autoregressive or encoder-decoder implementations. This mismatch may affect practical applicability and will be acknowledged clearly.

We also agree that the limitations discussed in Lin et al. (2024) such as brittleness under distribution shift, reliance on synthetic pretraining distributions, and sensitivity to context length—are relevant to our setting as well, and we will incorporate them in our limitations section.

We’ll also address the gap between theory and experiments, as well as the generalizability of our results to broader RL settings, since we only did experiments on bandit settings.

In the experimental setup, the first paragraph states that the action set is changing over the rounds, but the subsequent text suggests that the actions are fixed and the range of the rewards is changing. In (a) it changes to [3, 4] for some rounds (how? do you just add 3 to the reward?) and in (b) reward is multiplied by a cosine? The Appendix does not clarify anything.

We apologize for the unclear explanation in the original text. To clarify:

• In (1), the rewards in the rounds $t\in[50,100]\cup[350,400]$ are within the range $[3,4]$, while for the rest of the rounds, the reward $r\in[0,1]$. To achieve this, we added 3 to the rewards in the specified rounds.
• In (2), the reward is indeed multiplied by a cosine function, as stated. We will revise the manuscript to make these details clearer and avoid any confusion in the future. Thank you for your understanding.

Clarity: In particular, the non-continuous transformer is only introduced and defined in the appendix, and the description is difficult to parse. This architecture and data augmentation should be discussed (or at least mentioned) in the main body.

Thank you for the suggestion. We have added descriptions and discussion of the relevant parts in the main text to provide greater clarity.

We thank the reviewer for the comment. It is indeed possible to use a regular transformer, and we apologize if this wasn't clear in our previous rebuttal. While approximating MASTER with a regular transformer is certainly feasible, it would introduce substantial mathematical complexity that is orthogonal to our main contributions. We chose the augmented input approach to keep the presentation focused on our core algorithmic innovations rather than the technical details of the approximation. Below, we clarify the motivation for using augmented inputs and provide an alternative approximation approach using a regular transformer.

The core challenge in using a regular transformer lies in approximating $\sigma_2$​, which must select the first non-zero entry:

$\sigma_2(\\{(s_t^{(i)}, a_t^{(i)}, r_t^{(i)})\\}_{i=0}^n) = [0; \ldots; 0; (s_t^{(k)}, a_t^{(k)}, r_t^{(k)}); 0; \ldots; 0]$

where $k = \min\\{j \in \\{0, \ldots, n\\} : (s_t^{(j)}, a_t^{(j)}, r_t^{(j)}) \neq 0\\}$. This operation requires comparing all $n+1$ trajectories simultaneously to identify $k$, and it is the main reason why the augmented inputs are used.

One alternative approach is to use a regular transformer with $n+1$ heads to approximate $\sigma_1$ and $\sigma_2$:

$\bullet$ Layer 1 (Compression): Each head uses masked multi-head attention to extract per-trajectory representations, i.e., $\mathbf{H}^{(i)}=\mathbf{W}_i^{\mathrm{comp}}\cdot \mathrm{MaskAttn}_m(\mathbf{H})\in\mathbb{R}^{\frac{D}{n+1}\times T}$. Here, $\mathrm{MaskAttn}$ restricts information flow across heads, and the combined output is $\mathbf{H}' = [\mathbf{H}^{(0)};...;\mathbf{H}^{(n)}]\in\mathbb{R}^{D\times T}$.

$\bullet$ Layer 2~K: Generate Bernoulli masks, block masks, and select the first non-zero trajectory to approximate $\sigma_1$ and $\sigma_2$ (same as the procedure in the paper).

$\bullet$ Final Layer (Decompression): Reconstruct the selected trajectory: $\mathbf{H}''=\sum_{m=0}^n\mathbf{W}_m^{\mathrm{decomp}}\cdot\mathbf{H}^{(m)}\in\mathbb{R}^{D\times T}$.

However, this method adds several analytical burdens: we must ensure that $\mathbf{W}_m^{\mathrm{decomp}}\cdot \mathbf{W}_m^{\mathrm{comp}}\approx{I_D}$ to preserve information, and the analysis becomes more cumbersome.

In contrast, our proposed method is a simple and transparent preprocessing step where augmented inputs are generated via a linear transformation $\mathfrak{h}_t = \mathbf{M}_t \cdot \mathbf{h}_t$. This preserves the core semantics while avoiding unnecessary architectural detours. Importantly, the resulting non-continuous transformer is structurally equivalent to a standard transformer with transformed inputs.

As for changes to the paper, we will include a brief description of the non-continuous transformer in the main text, while keeping it concise to remain within the length limits. The discussion of limitations will be moved to the appendix. We appreciate the reviewer’s consideration and helpful feedback.

We are grateful to the reviewer for their thorough review and valuable feedback on our paper. In the following, we hope to address each point raised in the review.

Q1/W1:

How strong Assumption 2 is?

We appreciate the reviewer raising the question. While Assumption 2 is indeed strong and difficult to verify empirically, it builds on solid theoretical foundations. For example, [1] shows that transformers are universal approximators over distributions, which supports the feasibility of such an assumption in principle.

It’s also worth emphasizing that the assumption doesn't require the transformer to exactly replicate the expert policy. Rather, it assumes the transformer can learn an amortized algorithm whose output distribution matches that of the expert, which is a more relaxed and realistic target in many cases.

A similar assumption is also used in [2], and several previous works (e.g., [3, 4, 5, 6]) have proved the expressivity of transformers from other aspects.

W2:

There are few minor typos and grammar mistakes.

Thank you for pointing this out. We have carefully revised the manuscript to correct all identified typos and grammatical issues.

[1] Furuya, T., de Hoop, M. V., & Peyré, G. Transformers are Universal In-context Learners. In The Thirteenth International Conference on Learning Representations.

[2] Lin, L., Bai, Y., & Mei, S. Transformers as Decision Makers: Provable In-Context Reinforcement Learning via Supervised Pretraining. In The Twelfth International Conference on Learning Representations.

[3] Yun, C., Bhojanapalli, S., Rawat, A. S., Reddi, S., & Kumar, S. Are Transformers universal approximators of sequence-to-sequence functions?. In International Conference on Learning Representations.

[4] Zhang, R., Frei, S., & Bartlett, P. L. (2024). Trained transformers learn linear models in-context. Journal of Machine Learning Research, 25(49), 1-55.

[5] Bai, Y., Chen, F., Wang, H., Xiong, C., & Mei, S. (2023). Transformers as statisticians: Provable in-context learning with in-context algorithm selection. Advances in neural information processing systems, 36, 57125-57211.

[6] Ahn, K., Cheng, X., Daneshmand, H., & Sra, S. (2023). Transformers learn to implement preconditioned gradient descent for in-context learning. Advances in Neural Information Processing Systems, 36, 45614-45650.

We are grateful to the reviewer for their thorough review and valuable feedback on our paper. In the following, we hope to address each point raised in the review.

Q1/W1:

While your theoretical results are strong for bandit settings, what are the challenges in extending this to more general RL algorithms and potentially more structured models of change?

Thank you for the insightful question. While our experiments are limited to bandit settings, the theoretical framework extends to more general reinforcement learning (RL) problems, as long as the expert policy can be expressed as a distribution over actions given trajectories—a property shared by many episodic or tabular RL algorithms.

A key challenge in general RL is the added complexity from temporal credit assignment and state transitions. Unlike bandits, RL requires modeling long-term dependencies, which may necessitate deeper architectures and longer contexts.

To address this, our approach does not require the transformer to replicate the expert’s full decision process. Instead, it approximates the expert’s output distribution (amortized inference), which is more tractable and better aligned with in-context learning paradigms. This relaxation is consistent with prior work (e.g., [1]) and backed by universal approximation results.

Extending our analysis to structured forms of non-stationarity (e.g., piecewise-stationary MDPs) is an important next step. We believe our decomposition, especially of restart and windowing mechanisms, provides a promising foundation, and we now discuss this in the conclusion.

W2:

Continuously generating or providing optimal input-output pairs for training the transformer to adapt to changing dynamics might be prohibitively expensive or impossible.

We appreciate the reviewer’s observation. Indeed, in certain non-stationary reinforcement learning settings, conventional RL algorithms may offer greater reliability or data efficiency than transformer-based approaches. However, the primary goal of our work is to establish a theoretical guarantee for the robustness of transformers under non-stationarity, rather than to optimize for data efficiency or other empirical performance aspects.

Corollary 2 states that the regret bound holds when the number of training samples $N\geq CT^3 \log T$. If we take $T=10^3$, which already covers the horizon of many RL tasks, this gives a data requirement on the order of $10^9$. While this is certainly large compared to what's typically needed in traditional RL algorithms, it's still relatively modest compared to the scale of data used to train LLMs.

[1] Furuya, T., de Hoop, M. V., & Peyré, G. Transformers are Universal In-context Learners. In The Thirteenth International Conference on Learning Representations.

Thank you for the response! This answers my questions.

We are grateful to the reviewer for their thorough review and valuable feedback on our paper. In the following, we hope to address each point raised in the review.

Q2/W1:

Assumption 2 requires that the transformer can approximate the base algorithm arbitrarily well. How realistic is this in practice, especially for complex algorithms or environments with long horizons?

Thank you for raising this important point regarding the realizability assumption. We clarify below its motivation, scope, and practical implications.

Theoretical grounding: Assumption 2 is based on prior work showing that transformers can emulate standard reinforcement learning algorithms such as LinUCB and Thompson Sampling via in-context learning [1]. Furthermore, [2] shows that transformers are universal approximators over distributions, which supports the feasibility of such an assumption in principle. It’s also worth emphasizing that the assumption doesn't require the transformer to exactly replicate the expert policy. Rather, it assumes the transformer can learn an amortized algorithm that reproduces the expert’s output distribution, which is a more relaxed and realistic target in many cases.

Practical feasibility: We emphasize that Assumption 2 does not imply that any complex algorithm or environment is immediately approximable by arbitrary transformers. Rather, it reflects a scaling law—namely, that approximation error $\varepsilon$ can be made small by increasing model size and training data, as reflected in the regret bound (Theorem 1) and its dependency on architecture parameters (e.g., $D' = O(\sqrt{T \log^2 T / \varepsilon})$). In this sense, the assumption is realistic to the extent that the base algorithm itself is implementable by a recurrent or sequential decision process. For many commonly studied algorithms, including sliding-window UCB and restart-based methods like MASTER, this has been shown to be the case in structured architectures [3].

Constructive verification: In Section 4.2 and Appendix D.2, we construct specific transformer blocks that replicate components of non-stationary algorithms—such as window schedulers and restart logic—using standard attention and MLP modules. This shows the assumption is not merely abstract but realizable within standard architectures.

W2:

The architectural implementation of non-stationary mechanisms (e.g., the restart mechanism via 2-layer MLPs) is elegant, but it may oversimplify real-world scenarios. The paper does not explore how robust these designs are to scaling or how sensitive performance is to architecture choices.

Theory: Instead of exhaustively exploring design choices, our algorithmic construction is designed to demonstrate the existence of a transformer that can approximate the expert algorithm across the entire input space. For the restart mechanism specifically, many RL algorithms adopt either deterministic strategies (e.g., periodic restarts) or rule-based ones (e.g., resets triggered by specific conditions). Deterministic approaches can be viewed as a form of window scheduling, where the algorithm either discards historical data or resets its internal state at fixed intervals. Rule-based strategies often rely on binary indicators, which can be naturally modeled by a 2-layer MLP. In this sense, our architectural implementation of non-stationary mechanisms is expressive enough to capture a broad class of restart behaviors, at least from a theoretical standpoint.

Experiment: While more in-depth experiments on scaling and performance sensitivity of the architecture are left to future work, we conducted experiments using a transformer with 4 heads and 4 layers (compared to the 16 heads and 16 layers described in the paper) under the Low Non-stationarity setting to address the reviewer’s concerns. The "Regret Over Time" results are summarized in the table below. It demonstrates that the transformer is capable of achieving performance comparable to that of the expert, even with 4 heads and 4 layers.

Q1/W3:

How do the authors anticipate their theoretical findings will extend to more general reinforcement learning settings, particularly those involving large or continuous state spaces and longer temporal dependencies?

Thank you for the question. We believe that the implications from our results are applicable to settings such as continuous state spaces and longer temporal dependencies. As this work is primarily theoretical, we initially focused less on experiments. However, to better address the reviewer’s concerns, we extended our evaluation to a continuous state space environment.

We ran additional experiments on the MuJoCo Hopper task, which features an 11D continuous state space (e.g., joint angles, velocities) and a 3D continuous action space corresponding to joint torques. The objective is to learn a hopping behavior that balances forward progress, energy efficiency, and stability. Rewards are shaped accordingly to encourage forward velocity while penalizing excessive control effort and contact forces.

We used Soft Actor-Critic (SAC) [4] as our base algorithm and SAC+MASTER as the expert. For comparison, we trained a 3-layer transformer (with 256D embeddings) on stationary data for 20,000 episodes. During training, the transformer was evaluated on non-stationary data every 100 episodes, by switching to model.eval() to perform evaluation for 10 episodes, and then switching back to model.train() for further training. Similarly, SAC was trained on stationary data, and every 10 episodes, it was evaluated in a non-stationary environment.

This setup follows prior work (e.g., [5, 6]). Although individual state components range over $\mathbb{R}$, typical ranges are more bounded: the z-axis lies roughly within [0, 2], joint angles within [–1, 1], and velocities within [–10, 10]. We introduce non-stationarity by adding a drift value to the states, defined as $\sin(0.01 t)$. (new_state = state + $\sin(0.01 t)$)

Results:

Transformer:

SAC:

SAC+MASTER:

** Typically, SAC requires over 1M steps to converge. Due to time constraints, we only trained SAC+MASTER for 1M steps, as its learning process is much faster than SAC's due to the restart mechanism, which maintains smaller active memory pools and results in faster computation.

** SAC+MASTER Avg Reward after 900K steps (converged at this phase): Among all rewards in the evaluation episodes, 21% of the rewards are over 50, 13% over 70, 7% over 100, and 2% over 150.

The results show that the transformer maintains more stable performance under observation drift and achieves rewards comparable to the expert. While the reported SAC and transformer scores may appear lower than those in earlier work (which often exceed 1,000), we validated our SAC and transformer implementation under standard stationary conditions and confirmed it aligns with previous findings before introducing non-stationarity.

We acknowledge that this single environment does not fully capture the complexity of general RL tasks, but we view this as a meaningful first step toward evaluating the theoretical framework in richer, more dynamic environments. We plan to explore additional tasks with longer horizons and more complex dynamics in future work.

[1] Lin, L., Bai, Y., & Mei, S. Transformers as Decision Makers: Provable In-Context Reinforcement Learning via Supervised Pretraining. In The Twelfth International Conference on Learning Representations.

[2] Furuya, T., de Hoop, M. V., & Peyré, G. Transformers are Universal In-context Learners. In The Thirteenth International Conference on Learning Representations.

[3] Wei & Luo (2021). Non-stationary reinforcement learning without prior knowledge: An optimal black-box approach. COLT. [4] Haarnoja, T., Zhou, A., Abbeel, P., & Levine, S. (2018, July). Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In International conference on machine learning (pp. 1861-1870). Pmlr.

[5] Chen, L., Lu, K., Rajeswaran, A., Lee, K., Grover, A., Laskin, M., ... & Mordatch, I. (2021). Decision transformer: Reinforcement learning via sequence modeling. Advances in neural information processing systems, 34, 15084-15097.

[6] Huang, S., Hu, J., Yang, Z., Yang, L., Luo, T., Chen, H., ... & Yang, B. (2024). Decision mamba: Reinforcement learning via hybrid selective sequence modeling. Advances in Neural Information Processing Systems, 37, 72688-72709.