Balancing Context Length and Mixing Times for Reinforcement Learning at Scale
Authors: Matthew Riemer
Khimya Khetarpal
Janarthanan Rajendran
Sarath Chandar
Venue: NeurIPS 2024 poster (2024)
Keywords: Mixing Times
Non-Markovian
Average Reward
Long Context
Partial Observability
Causal Structure
Primary area: reinforcement_learning

Abstract: Due to the recent remarkable advances in artificial intelligence, researchers have begun to consider challenging learning problems such as learning to generalize behavior from large offline datasets or learning online in non-Markovian environments. Meanwhile, recent advances in both of these areas have increasingly relied on conditioning policies on large context lengths. A natural question is if there is a limit to the performance benefits of increasing the context length if the computation needed is available. In this work, we establish a novel theoretical result that links the context length of a policy to the time needed to reliably evaluate its performance (i.e., its mixing time) in large scale partially observable reinforcement learning environments that exhibit latent sub-task structure. This analysis underscores a key tradeoff: when we extend the context length, our policy can more effectively model non-Markovian dependencies, but this comes at the cost of potentially slower policy evaluation and as a result slower downstream learning. Moreover, our empirical results highlight the relevance of this analysis when leveraging Transformer based neural networks. This perspective will become increasingly pertinent as the field scales towards larger and more realistic environments, opening up a number of potential future directions for improving the way we design learning agents.

Balancing Context Length and Mixing Times for Reinforcement Learning at Scale
Matthew Riemer"
Khimya Khetarpal Mila
IBM Research, Mila, Université de Montréal
Janarthanan Rajendran
Sarath Chandar
Dalhousie University
Mila, École Polytechnique de Montréal
Abstract
Due the recent remarkable advances in artificial intelligence, researchers have begun to consider challenging learning problems such as learning to generalize behavior from large offline datasets or learning online in non- non-Markovian environ- ments. Meanwhile, recent advances both of these areas have increasingly relied conditioning policies on large context lengths. A natural question is if there limit to the performance benefits of increasing the context length if the computation needed is available. In this work, the context length of a policy to the time needed to reliably evaluate its performance (i.e., its mixing time) in large scale partially observable reinforcement learning environments that exhibit latent sub-task structure. This analysis underscores a key tradeoff: when we extend the context length, our policy can more effectively model non-Markovian dependencies, but this comes at the cost of potentially slower policy evaluation and result slower downstream learning. Moreover, our empirical results highlight the relevance of this analysis when leveraging Transformer based neural networks. This perspective will become increasingly pertinent as the field scales towards larger and more realistic environments, opening up number of potential future directions for improving the way we design learning agents.
establish a novel theoretical result that links
Introduction
Scaling the Context Length: As the field of AI moves towards harder problems where agents must model higher order non-Markovian dependencies, it is only natural to consider conditioning models on larger context lengths of the interaction history. When considering only small context lengths or single observations, models become restricted to some lesser notion of the best possible performance comparison what is achievable with an infinite context length [1]. Meanwhile, there are two obvious downsides to increasing the context length agent considers. The first downside is the increase in necessary computation. However, there is also ongoing research related to minimizing the needed quantity of computation [2; 3], increasing the parallelism of computation [4; 5; 6). and increasing the throughput of computation on modern hardware [7]. The second downside the high dimensionality of the input representation we must process and the difficulty of learning and generalizing from this kind of input. But it is unclear how worried researchers and practitioners should be about this downside given the recent improvements in processing long sequences that came with attention models [8; 9] and particularly the Transformer architecture [10]. Additionally, there has been trend towards training internet-scale datasets [11: 12; 13] that far exceed the quantity
Please direct correspondences to mdri carrimmer@us...m ibm. com. Work done during Postdoc at Mila, Université de Montréal,
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

of any single human's 's experience in domains of interest, which significantly alleviates requirements for generalization. Our paper aims to add to this discourse by being the first to establish another downside: policies that are conditioned on longer context lengths take longer to reliably evaluate because they impose looser bounds on the mixing time. This may make less safe to deploy these policies in the real-world because we are able to extrapolate less about their true behavior from our limited evaluations. Downstream learning may also be slower for the same reason as Monte Carlo algorithms can only perform reliable updates after being rolled out for the mixing time number of steps and the convergence time TD algorithms also depends on the mixing time.'
Addressing High Mixing Times: Recent work has highlighted the challenges presented by high mixing times as the field moves towards large scale and continual reinforcement learning problems [14]. Unfortunately, as highlighted by Riemer et al. [14] there are not many known solutions that address this problem. As a result, it is important for researchers to consider how changes to the policy architecture and learning algorithms may have an impact on mixing behavior. The effect that the input variables sent to a policy could have on bounds on the mixing time was previously considered by the pioneering work of Kearns and Koller [15]. In Theorem I present a tighter version of the bound presented in their paper. Moreover, our paper is the first to highlight the connection between this result and the context length in partially observable settings in Theorem 2. Our paper focuses solely on how design decisions related to the policy's s architecture can impact the mixing time and is complementary to recent insights about how changes the learning algorithm, such as incorporation of replay [16] or multi-level critics [17], can lead to faster learning when mixing times are high.
The core insight of this paper is that there is tradeoff to consider between increasing the context length and increasing the mixing time. When we increase the context length sent to agent, we increase its ability to model non-Markovian dependencies. However, this increase in context length may also cause an increase in the mixing time, which makes it take longer to reliably evaluate agent's 's policy. Dong et al. [18] propose the solution of restricting the class of policies they optimize over such that the mixing times are bounded, which is also the perspective we embrace in this work. However, Dong et al. [18) restrict this class of policies by limiting the planning horizon and thus making optimization more myopic. In our work, we highlight a different mechanism for restricting the mixing times of the policies we optimize over by simply limiting the context length. To this end, our key contributions are as follows:
Balancing Context Length and Mixing Times:
1. We present theoretical analysis with supporting toy examples shedding light on how the input sent to policy can influence its mixing time in Section 2.
2. In Section 3 we formally establish the connection between the context length considered by policy in partially observable environments and the mixing time. We also empirically verify the relevance of this theory to Transformer based agents in partially observable settings.
3. Finally, in Section 4 we highlight the relevance of this theory when building foundation models that imitate diverse set of behavior policies for RL. Specifically, we show that Decision Transformers [19] must use much larger context lengths than the policies that generated their dataset, and result require more interaction for reliable evaluation.
To reproduce our experiments see tttpps//giite httpr//nessteaatego////oo//ooooa cowm///tttriiieee//Coeotttooontt
2 Understanding How the Policy's Input Impacts the Mixing Time
Before discussing the connection between mixing times and context length in the next section, we first provide the technical grounding for readers to understand how the problem structure and the policy combine to jointly affect bounds on the mixing time. We will begin by detailing the problem formulation and present a novel upper bound for the mixing time in this context in Theorem 1. We then highlight key implications with illustrative example problems.
2.1 Preliminaries: the Average Reward RL Setting and Definitions of the Mixing Time
Problem Formulation: We adopt the average reward setting to enable long-term analysis of agents. Please note that we only use this perspective to evaluate agent's 's behavior and place no restriction
See Riemer al. [14] for a discussion of the difficulties presented for learning when mixing times are high.

---

on using discounting within the learning process (as we do in our experiments). 4 We consider a finite, discrete-tines timee, infinite horizon Markov Decision Process (MDP) [20; 21]: M (S, T, R). where S is the set of states, A is the set of actions, R : S x A - [0, R is the reward function, and T S A -> [0, 1] is the environment transition probability function. At each time step, an agent perceives state E S and takes an action E A drawn from a policy S x A -> [0. 1]. The agent then receives a scalar reward drawn from the function R(s, a) and with probability T [ss a) enters next state s'. Markov chains may be periodic and have multiple recurrent classes, but optimality is difficult to define in such cases [22], making the following assumption necessary:
Assumption All stationary policies are aperiodic and unichain, resulting in Markov chain with a single recurrent class that is recurrent in the Markov chain of every policy."
Any RL problem may be modified such that Assumption holds by adding arbitrarily small positive constant to all transition probabilities in ((''|,, a) [23]. An important corollary is that the steady-state distribution nº induced by the policy is independent of the initial state:
Corollary All policies π induce a unique steady-state distribution (T) lim Pa (s+- 8(80) that is independent of the initial state EvES 1eth (s) EaEa (sals) = (s') Vs' E.
such that
Corollary implies that long-term rewards and thus the average reward per step objective p(TT) can be defined independently of the current state
[21]:
m
p(π)= = lim m=000 M
Ex
R(s ai)
= lim E R(s), ai)
11000
Σ (s) SES
Σ π (clss)
AEA
Computing the average reward (i.e., the reward rate) with the last expression is limited by the amount of time the Markov chain induced by the policy T (s'ls) ΣeEA л (a]s))(ss||ss a) needs to be run for before reaching the distribution (s). This amount of time is called the mixing time of the induced Markov chain. We denote tmix (c) as the e-mixing time of the chain induced by π:
E = min {m max drv (P* (so). (-) S {
(-))
where drv is the total variation distance between the two distributions. The conventional mixing time defined as " = 1. (1/4). This only gives insight about distributional mismatch with respect to the steady-state distribution, which led Kearns and Singh [24] to introduce the notion of mismatch with respect to the reward rate. The e-return mixing time is a measure of the time it takes to formulate an accurate estimate of the true reward rate. Formally, if we denote the m-step average return starting from state so as p(TT, So, m) == E, [ Em relsp]. then we can define the e-return mixing time as:
1e (e) min { (m |I((T, so, m') p(T)) Se
Vso E S and Vm' 2 m
An Interpretable Metric: We will use this definition of the mixing time in our experiments as it directly measures the amount of time needed to evaluate policy at given threshold of precision c.
2.2 Mixing in MDPs with Multidimensional States
Multidimensional State Space: We will also assume that problems have a state spaces comprised of multiple dimensions. Formally, each state is separated into n 2 1 variables i.e., [st s") with s² E St for all variables i E {1, n such that S C S x x S". It is also important to note that each state variable can be modeled independently based on the previous state across all variables such that T(s'|s (s, a) = (I........ |s, a). In Appendix A. we discuss why this setting is strictly more general than the typical discrete MDP setting and does not limit the scope of our results.
'See Appendix C for additional explanation about why this is most general possible problem formulation. 'This corresponds to what is called an ergodicity assumption for stationary policies in Sutton and Barto [21]. See Appendix A for more detailed discussion of this assumption and its implications.
3

---

Dynamic Bayesian Network (DBN) Structure: A Bayesian Network (BN) is probabilistic graph- ical model that represents a set of variables and their conditional dependencies via a directed acyelic graph (DAG) and a dynamic Bayesian network (DBN) is BN that relates variables to each other over adjacent time steps. In general, any Markov chain induced by a policy can be seen as a DBN over the state variables 8² ..... 8 that is also influenced by an action variable dictated by the policy [15]. The dependency graph D for the DBN is a directed cyclic graph whose nodes are s ."... s" and where there is then a directed path from s' to s in DF iff sj influences sy for some t' >> t under policy л. While can ignore actions nodes in the graph, it is still important to consider their effect in forming causal connections among the state variables. Note that every node in 2" should influence itself and thus have a self-loop. If there is direct path from s to s and from s to s', the variables/ and j are said to be in the same strongly connected component. In contrast, if there is either a directed path from s to s or s to 8', but not both, the variables and j are weakly connected such that the one that has a causal influence on the other can be considered as ordered before the one it causes. Let IT, ..., 17 be the maximal strongly connected components in D when an agent behaves with an arbitrary policy π that is sorted such that if i < j there are no directed edges from I to rT. We can then define gd max; [[]] as the maximum number of variables in any strongly connected component of "D'T and g = max, 9 as the maximum possible within a given policy class or parameterization æ. In Appendix A.I we provide examples to help readers better understand how this works and demonstrate that this formulation does not limit the scope of our results.
Mixing Times and Coupling Times: We can now leverage the well-studied connection between the mixing time and the so called coupling time to provide an upper bound on the mixing time of the Markov chain induced by [25]. Let T be the random variable that represents the coupling time defined as the smallest m for which two Markov chains starting from different initial states are in the same state. Following the analysis provided in Lemma 5.2 of [15], for any e, let be the number of steps such that for any two starting states in S if we can say that P(T > m) S e then the Markov chain is e-mixed time m. For this analysis we also must introduce the parameter BI. which defines the minimum amount of common probability mass between any two state configurations for any state variable i so that 3 ming min, 3x,x E [0, 1] further corresponds to the minimum over all state variables i and possible parameterizations of the policy n.
Σ min (o'1))))) (1(1)))))))11)))))) S'ES'
))
B1, = min #1.82ES
We can now introduce a mixing time bound as a function of the strongly connected component structure. While our result has similar motivation to Theorem 5.4 of Kearns and Koller [15]. we note that the proof of Theorem I outlined in Appendix A is entirely our novel contribution building off Lemma 5.2 and Definition 5.3 of Kearns and Koller [15]. Additionally, there is no version of Kearns and Koller [15] online that includes the proof of Theorem 5.4.
Theorem 1 (Strongly Connected State Variables Bound): If the Markov chain T (d/1s) induced by policy π has l maximal strongly connected components in D with maximum size of state variables and a minimum of 3 common probability mass between any two state configurations for any state variable, the mixing times (T)(() and (E) can be upper bounded:
Eye (c) E (E) O
39 log (o9(1
Proof Sketch: Due to the sorting of the strongly connected components, our analysis is based on coupling each of the IT''s in succession. Because it is possible that multiple IT's couple at the same step, every step where the Markov chain does not fully couple must be a step where some I'T does not couple. Our proof proceeds in the following high-level steps: 1) The probability of 1 coupling at given step once IT. VT.....I? Ih_ have all already coupled is 2 89. 2) Thus the probability of IT not coupling at step when IT. ...... r have all already coupled is < (1 - 39). 3) So the joint probability of IT not coupling for m, 2 0 steps when IT ........? have all already coupled is
< (1- - 3 mm. 4) If 7 > m then D i=1 m. = m and the joint probability that m--teps have been spent not coupling in some IT has a probability bound independent of the particular allocation m into individual m,, Thus we can conclude that P(T > m) S (1 f39mm 5) Leveraging the identity
that 1 - S - for > 0, we find that P(T > m) < (e- 9a ym 6) The Markov chain is e-mixed
4

---

if P(T > m) S c, so it must be e-mixed if (e- (e- I e, which implies that m 2 3o log(1/c). 7)
Finally, we note the relationship between fre (c) and (())() following Lemma 1 of [24].
Bound Tightness: Theorem is a tighter version of Thhhorem 5.4 presented by Kearns and Koller [15] by a factor of 8 from their result of O
Rd. 13
(lo((1/e) This tighter bound underscores definitively
that having more strongly connected components of a smaller size leads to tighter bounds on the mixing time than having fewer strongly connected components of bigger size. In fact, without this contribution beyond the result by Kearns and Koller [15] it would not be possible to show our main result in Theorem 2 as would thus be unclear if the mixing time bound gets tighter with the shrinking context length. An intuition we could provide for how this bound is achieved is that if the Markov chain has not coupled for m steps, each time step must be spent not coupling in least one of the strongly connected components, which has a bounded conditional probability for any component. Our detailed proof is in Appendix A. Moreover, while it may seem that only a restricted class of Markov chains have B > O, in Corollary 2 of Appendix A, we demonstrate the extension of the definition for the probability mass in common over e > steps, which must be > for some value of e due to Assumption 1. Our approach could also yield an even tighter yet more cumbersome bound in terms of the maximum probability that any strongly connected component does not couple rather than using 3 and g. which is important when we see significant variation in 3; across variables.
Conditioning on Variable Subsets: One of the primary implications of Theorem 1 is that the mixing time of a policy is closely connected to the number of state variables that influence it's actions. For example, a policy in the class T--(s). which takes full states = s ..... s" S as input can be associated with 3n, and In. Mcanwhile, a policy in an alternative class that only considers some potentially time dependent arbitrary subset of n (t) S n state variables as input at any moment in time such that s² ..., sx((t) C "'.,.., S can be associated with B, and gnc. Lemma I in Appendix A then shows that the mixing time bound gets strictly tighter because 1//3a < 1//3g.. We illustrate this phenomenon more concretely through two toy MDP examples detailed in Figure 1.
2.3 Building an Intuition with Examples
We will consider the mixing times for the two examples highlighted in Figure over the space of deterministic policies. As in Riemer et al. [14], we relax the definition of the e-return mixing time using an average over start states (rather than maximum) in order to emphasize mixing times encountered in practice. See Appendix B for details on our methodology for estimating mixing times.
MDPs with Irrelevant Variables (Figure 1): In the first example we consider, there are two independent state variables and y that are both influenced by the agent's actions. As a result, æ and are in the same strongly connected component if both variables also influence the agent's actions. Because the reward only depends on x, the agent can achieve the optimal policy without considering the variable in its decision about actions. If the policy is only conditioned on the variable x, influences y through its actions, but y no longer has an influence on æ. As a result, the Markov chain induced by the policy that only conditions on the relevant variable x has two strongly connected components of size one such that g = 1 whereas the policy that conditions on both relevant and irrelevant variables potentially has a single strongly connected component of size two such that g = 2 with 3 held constant. Based on Theorem I, we expect the policy conditioned on both variables to experience higher mixing times even despite the small scale of this toy problem. Indeed, this is the case when we compute the e-return mixing time at reward rate precision of e = 0.05, which must reflect a minimum of 10% relative estimation error because no policy has a reward rate of 0.5. The optimal policy has a mixing time of 17.6 and only depends on the variable x. In fact, no policy only depending on æ has mixing time above 44.9. Mcanwhile, there is policy dependent on both variables that is within the reward rate precision threshold of the optimal policy with a mixing time of 45.2, and policy conditioned on both variables with a mixing time as high as 340.3.
Impact of Reward Density: To test a tighter result is possible with respect to the ratio between the maximum mixing times predicted by Theorem 1, we have tried a setting where the reward is still 1.0 when x = z0, but edited to -0.25 when x = x1 and -1.5 when x x2 rather than This setting was chosen to have impact on the optimal policy while densifying rewards without making a substantial change to the magnitude of best reward rates achieved. The ratio between the maximum mixing times is now improved to 8.8 times smaller (233.0 vs. 2049.7) when only focusing on the
5

---

y2
al
80
a0
al
a0
ya
ab
a0
XO
yo
0.1
z0
0.1
Figure I: Illustrative Toy Examples. The above figure details the three relevant variables y, and that we will consider for our toy examples. Note the action a0 is interpreted as al or vice versa with a 10% failure probability, which is the same as the rate at which = switches regardless of the agent's actions. Irrelevant Variables Example: In this case, we consider only the variables a and (z is not needed in this example) where the reward is +1 if I = z0 and 0 otherwise. The result is that variable x is relevant to the task whereas variable y is irrelevant. Independent Subtasks Example: In this case, we consider that the variable z evolves (according to the diagram) with y remaining constant when 2 z0 and the variables y evolves with x remaining constant when 2 = z1. The reward is +1 if =0 and x z0 or if z1 and y = y0, 0 otherwise.
relevant variables rather 7.6 times smaller (44.9 vs. 340.3) with the previous sparse rewards. Note that the reward still has no dependence on y and is thus invariant to large part of the change in the state. So, mixing is still substantially faster with respect to the reward (tign (c)) than the state (fmix (c)).
MDPs with Independent Subtasks (Figure 1): We can now consider an example with three vari- ables x, y, and z where the agent's 's actions influence both and y. A policy conditioned on all three variables has a strongly connected component of size one including (because z is not influenced by the agents actions or x or and a strongly connected component of size two including and such that g = 2. On the other hand, if the policy were able to condition only on x when 2 = =0 and only condition on y when
21, then the actions that still influence both and y no longer serve to make their values dependent on each other. This results in three strongly connected components of size one such that g 1 with 33 held constant. One for each variable y. and z. Based on the analysis of Theorem 1, we would again expect the policy conditioned on all variables always to experience higher mixing times than the policy conditioned specific subtask even despite the small scale this toy problem. As for the previous example, we compute the e-return mixing time reward rate precision of 0.05, reflecting minimum of 10% relative estimation error. The optimal policy has a mixing time of 16.0 and only depends on the variable when = =0 and variable y when a mixing time above 267.1. Mcanwhile, there is a policy dependent on all variables that is within the reward rate precision threshold of the optimal policy with mixing time of 341.2. Moreover, there is a policy conditioned on all variables with a mixing time as high as 968.1. This example helps illustrate the important fact that a policy can limit its mixing time while still conditioning on all variables if the subset it conditions on has some time dependence.
only the relevant variables for the
=1. In fact, policy structured like this has
3 Understanding How the Context Length Impacts the Mixing Time
In the previous section we were concerned only with fully observable MDPs, but real-world and large scale problems tend to be partially observable. We will begin by highlighting new notation that allows us to talk about notions of optimality as function of the interaction history context length and present mixing time bounds that are dependent on the context length considered by a policy. We go on to conduct experiments that verify the relevance of this analysis during online RL in partially observable domains and demonstrate that Transformer neural network models actually experience higher mixing times than alternative function approximators and tabular models.

---

3.1 Partially Observable Environments with Local Observation Structure
Partially Observable Environments: We can extend the notion of an MDP from the previous section to consider an unknowable but ultimately stationary Partially Observable Markov Decision Process (POMDP) [26], which is comprised of an augmented tuple P (S, A. O, R. N.,TOO, O). This adds to the definition of an MDP M an observation space O and an observation function O(oss) that maps states s E to observations o E O. The state 8 E of such a POMDP ultimately serves as a theoretical quantity that is never actually observed from the agent's perspective and the observations o E O that does receive may have an arbitrarily non- Markovian relationship.
Interaction Histories: At time t the union of all things the agent has observed about an envi- romment can be called its interaction history ht == {on, do, To, 01 "I,,,,,,,. 0t). We can also consider the case where the agent only maintains finite window of history with size k i.e.,
h,()) = {oL-++1, at-k+1, -+- Ft 0t-k+2; Pt Tt= 0z where our notation is chosen so that memoryless policy that only processes the current observation corresponds to k 1. We will henceforth call the context length of the interaction history. A given POMDP p is non- Markovian to order kp which implies that for every combination of state s, action and history window (*), T(+ss, a) = ((-) a) and R(-[s, a) = R(- a) Ak 2 kp.
Optimality in POMDPs: An advantage of using policy that processes the full interaction history is that it is possible to handle arbitrarily non-Markovian environments, but this is not computationally scalable in the long-run. As 00, no (-(e)) processes infinite length sequence hg. In practice,
must limit the context length sent to the policy. The best policy no ( (-1h% (*))))) over a finite context length has the same optimal reward rate the best policy no To (-1s)) over true state inputs when k 2 kp such that limy-+s0 m =.(-In (1.,.))..0 [ra] = lim,,-so maxøre@r E (-1+,) [ra] However, this implies
the context length must scale with kp, which could become quite large in highly non-Markovian real-world env avviroments where achieving optimality is difficult.
The Effeet of Context Length: As described in the previous section, are interested in multi- dimensional state spaces consisting of n variables and in particular problems where n is large. A given observation generated from O(ols) may be then caused by only subset of the state variables Par(o) S {1..., n where (ar(o) denotes the causal parent state variables of o. As such, a finite history window (k) may also be caused by a subset of the state variables because the same can be said for rewards and actions that are potentially included in this input. In general, finite context length is reflective of a maximum n& (f) sized subset the state variables at a time where n 2 ner(t) 2 n2 (t(()) if context length k' 2 k for all L. This is because (t) C (a)()())) C n... Theorem then follows from the combined results of Lemma in Appendix A and Theorem
Theorem 2 (Limiting Mixing Times with Context Length): If the Markov chain induced by policy conditioned on a finite interaction history window π ()) with a context length of k has lu maximal strongly connected components in Dn with maximum size of gk variables and minimum of By common probability mass between any two state configurations, the mixing times the (e) and (e) can be bounded for any ke 2 k:
1
t (c) E tHT (c)
O log(1/c)
8
log(1/e)
where the dependence on the context length implies that S BR S B. S 1 and gk. 2 gk 2
Proof Sketch; Lemma in the appendix considers the mixing time relationship of policy classes conditioned on subsets of the state variables that other policy classes conditioned on. Our
'Note that any non-stationary MDP can alternatively be viewed a POMDP (see Proposition 2 of Khetarpal [27]) and by the same logic any non-stationary POMDP can alternatively be viewed as stationary POMDP, so there is no loss of generality in making the assumption that the POMDP is stationary.
7 Maany exclude reward from the history representation in POMDPs, but we include it to be general as possible. That said, theoretical results do not require this particular choice of the history representation.
"Not the distinction between O', the possible policies over st, and θ, the possible policies over (!())
7

---

proof includes the following high-level steps by applying our notation in which k' 2 k for all to the results of Theorem and Lemma 1: 1) We consider the causal parent state variables of each observation, action, and reward to conclude that Par (h(+)) S Par (h()) S n], which implies that n 2 n& (t) Tek (t). 2) Through Lemma 1 we show that by rule of Cartesian products over subsets that S BK S 3x S 1. 3) Through Lemma we also demonstrate that g 2 9k 2 1 because causal connections in D are only added and not removed when the context length is increased. 4)
This then implies that // 2 1/5gg. which is sufficient to prove Theorem using Theorem 1
because is independent of k.
Theorem 2 formally estab- lishes novel connection between the context length k of a policy and bounds on its mixing time. Growing k increases the performance of the best achievable policy when hp is high. However, any increase in k leads to a looser mixing time bound.
Formalizing the Tradeoff:
0OO0 ST
Se+1
S7+1
Figure 2: POMDPs with Local Obser- vation Structure. An example of an en- vironment with a multidimensional state space where near term observations only causally dependent on a subset of the variables. At time T >> I see different dimensions of the state space influencing observations and thus these separate sub- sets do not need to be in the same strongly connected component.
When it Really Matters: While this bound is always true, there may not be a meaningful dependence on k for some problems. This includes very simple environ- ments where kp + 1 or environments with no local structure where all state variables influence all observa- tions. However, if changing the context length l does actually lead a change in the strongly connected com- ponent structure, the impact the mixing time will be exponential. In Figure 2 we present an illustration of the kinds of problems where the chosen value of may make a large difference the value of (4). We can see that at each time only a subset of the total state variables contributes to the local observations of the agent and that which variables they are has a degree of local consistency across time-steps of the Markov chain induced by the policy in the environment. When observations are only caused by a subset of the state variables at each step, there is more potential to break the problem up into independent subtasks in our example in the previous section.
3.2 Empirical Verification During Online RI.
Average Mixing Time
1250-
1000
Context Length and Encountered Mixing Times:
750
While Theorem 2 draws a clear connection between context length and mixing times, there still remains question about if this analysis is too conservative and it unclear if these policies will actually be encountered during learning. To test this question empirically, we consider the simple RGB world environment Figure 6a, which is modeled after the classic T-maze [28] en- vironment. We conducted comprehensive experiments leveraging tabular Q-lcarning for 1 billion steps over 100 seeds (see Appendix B). In Figure 3 we highlight that longer context lengths experience larger average mixing times at the same approximate reward rate as policies conditioned on lower context lengths during learning. When interpreting this kind of plot, it is im- portant to consider that there is a difference between the input required to express a policy and the input available to that policy. Any policy that receives sufficient statistic of the environment state as input i.c. k > 2 will arrive at the same optimal policy by the end of training, so what is more interesting is the points encountered along the way.
500-
250
0.0
0.1
0.2
0.3
0.4
Reward Rate
Figure 3: Mixing Times Encountered vs. Context Length. We plot the average mix- ing time and 95% confidence intervals en- countered during billion steps of learning over 100 seeds at each context length k. We bin the average mixing time compu- tation by the nearest 0.1 increment of the reward rate of the policy. k = 1 is not visible because the reward rate is always O.
Effect of Neural Network Architecture: We extend these experiments to prominent neural network architectures on the same problem leveraging Q-learning (see Appendix B). We plot the result when each architecture has about IM parameters in Figure 4. All models have the same mixing time when they arrive at the same policy i.e. the optimal policy. However, at intermediate reward rates we see a higher average mixing time for Transformer models. It is important to note that sometimes an architecture will learn to achieve reward rate that others never learn to achieve. In this case, there
8

---

Foundation models may struggle to model the training data when the context window length A used for training is less than the maximum context window of any policy kn := model is not clear basis for comparison. We also ran experiments evaluating the role of the number of parameters. The extra capacity did not have statistically significant effect on the mixing times for MI.P and LSTM models. Meanwhile, extra capacity resulted in higher encountered mixing times throughout learning and a larger effect when increasing for Transformer models (see Figure 8). appears that attention mechanisms make it easier to focus on the full context rather than i.e. only the recent parts and that this capability is predictably enhanced when the model capacity is increased. In Appendix B we take closer look at the Transformer attention maps in the decoder and present some evidence that the larger models are paying attention more uniformly to the entire context.
3000
4 Understanding Growing Context Lengths in Foundation Models for RL
Averaee Mixing Time
Foundation Models for RL: Large scale founda- tion models trained to recreate behaviors from large and diverse distribution have recently had a disrupting effect across the field of AI. Concretely, a foundation model leverages an offline dataset D, containing data from N different now unknown behavior policies T. for i E N}. Deci- sion Transformers [19] is a popular approach that achieved state of the art performance in offline RI. by treating learning as a sequence modeling prob- lem akin to language modeling. The objective is to imitate the behavior of all policies π,, in the dataset using& window of their interaction history (R). See Appendix B for additional details.
1600
500
MLP
18TM
Traskerer
-0.1
0.0
0.1
0.2 Reward Rate
0.3
0.5
Figure 4: Policy Architecture vs. Encoun- tered Mixing Times. We plot the average mix- ing time and 95% confidence intervals for each choice of policy architecture between tabular, MLP, LSTM, and Transformer models with av- erages binned by the reward rate. We provide a representative example with k = 5 and include comprehensive plots in Appendix B. Mixing times goes down as the absolute value of the reward rate gets close to 0 as a consequence of the sparse and local reward structure.
Fitting Training Data:
max k2 The reason is because must not only produce the behavior of each agent n, but also, must store additional context to disambiguate R, from all ny Vj E [1. NY i. Thus it will be necessary to consider large contexts, well beyond that of any context length used by the behavior policies, especially as the number and diversity of behavior policies considered grows.
Comparison to Behavior Policy Context Lengths: To understand the connection between the behavior policy context length and the context length needed to train Decision Transformers, we consider a similar setup to the Key-Door experiments in the original paper [19], but with the publicly available Minigrid Crossing environment [29] (Figure 6b). We randomly initialized 1,000 deterministic behavior policies, each rolled out for one episode in the environment to collect dataset DD leveraging a CNN architecture following past work on the Minigrid domain [30; 31]. The context length of each behavior policy is set to 1, implying kn 1. In Figure 5 we plot the training accuracy achieved by Decision Transformer model as function of its context length k. We report that the training accuracy is only optimal in general for > 50, which is significantly higher than ko with k = 25 reaching optimal performance for some random seeds. We also considered 1,000 episodes generated by different stochastic behavior
100%
Random Policies
Training Accuracy
Learning Policies
90%
80%
70%
60%
10
10' Context Length
10
Figure 5: Context Length vs. Training Accu- racy. We plot the achieved training accuracy and 95% confidence intervals across 5 random seeds as a function of the Decision Transformer context length in the crossing environment with 1,000 episodes of data generated by either ran- dom behavior policies with context length k = 1 or a REINFORCE based learning agent using context length k

---

policies of REINFORCE based learning agents [32] that learn following every episode. Figure 5 demonstrates that is even harder to model these policies. This makes sense both because the policy distinctions are more subtle and less diverse between episodes and because the policies are stochastic, leaving an irreducible source of uncertainty.
Evaluating Learned Models: It could be argued that, although a larger context length may be required capture the complete range of behavior policies compared to each policy individually, this merely promotes overfitting of the training data and does not effectively enhance downstream performance. So, to see if Decision Transformers need near optimal training accuracy for good downstream performance, we evaluated each Decision Transformer model for particular random policy seed where 25 is able to achieve optimal training performance across a variety of return to go prompts stepping by 0.01 from O to 1.0 and report the average performance across 1,000 episodes. We plot our results in Figure which validate the importance of fitting the training data. Decision Transformer with A = 25 can achieve performance as good as any behavior policy used to generate the data. Meanwhile, models with smaller fail to achieve the same performance. Moreover, increasing context length k, while important for optimization, raises the mixing time. The average mixing time with precision e = 0.01 across policies and start states for the Decision Transformer with k = 25 is 298.1 episodes or 86,834.3 steps while the average mixing time is merely 11.2 episodes or 3,271.2 steps for the behavior policy across policies and start states. Figure 9 also demonstrates the not very surprising conclusion that unnecessarily large context lengths more prone to overfitting than the minimal context length that achieves 100% training accuracy.
Discussion and Future Work
In this work, we have highlighted the potential limitations of training models conditioned on ever increasing context lengths and particularly the effect that these growing context lengths have on mixing times. This motivates a number of interesting research questions to explore moving forward.
New Architectures and Algorithms: Most work on RL that even acknowledges the challenges associated with high mixing times does so with a defcatist mentality, assuming that problems with high mixing times are unavoidably harder and that there is basically nothing that could be done about it. Our work highlights that this isn't actually true and that the policy class we choose to optimize over itself can have big impact on mixing properties. What our paper shows in Theorem 2 is that what leads to potentially high mixing times is when our model leverages a monolithic representation that is highly sensitive to a large part of the interaction history at all times. This is particularly descriptive of how vanilla transformers work, but there are multiple already existing research directions that seem well suited to scaling to high context lengths while providing less history sensitivity at each step. We refer interested readers to Appendix D for an in depth discussion of related directions.
We believe the settings of greatest relevance to our work are those related to continual or multi-task environments where agents are evaluated as generalists over a number of skills rather than just solving single narrow task. As such, we believe that focus on the difficulties presented by high mixing times is timely in the age of foundation models. As mentioned at the end of Section 3.1, our analysis will not have relevance in problems where there are few state variables that each impact every observation. However, composite tasks that test a number of sub-skills naturally tend to have many total state variables with relatively few impacting each observation. So, for example, simple Atari domains will not suffer from high mixing times, but i.c. continual learning over multiple Atari games will [14] result of the sparsity of causal impact of variables across games. Broadly speaking. Al assistant tasks that include providing help on a number of topics rather than just one should also suffer from issues with high mixing times.
Scaling to Complex Environments:
Our work additionally highlights how the way we pretrain foundation models may make high confidence evaluation of these models more difficult. Trusted evaluation of foundation models remains an important challenge for the research community to grapple with. Towards this end, our work is the first to establish that more interaction is needed to reliably evaluate models that have larger context lengths. This novel perspective very important for researchers to consider as these models are increasingly being deployed in the real-world.
Evaluation of Foundation Models:
10

---

Acknowledgments and Disclosure of Funding
We would like to thank Murray Campbell, Miao Liu, and Payel Das for valuable conversations over the course of this project. This project was supported by the IBM-Mila collaboration grant. We would also like to acknowledge our support from the Canada CIFAR AI Chair Program and from the Canada Excellence Research Chairs (CERC) Program. We thank the IBM Cognitive Compute Cluster and the Mila cluster for providing computational resources. Finally, we really appreciate the feedback we received from the NeurlPS reviewers, which helped to improve the presentation of our key findings.
References
[1] Satinder P Singh, Tommi Jaakkola, and Michael Jordan. Learning without state-estimation in partially observable markovian decision processes. In Machine Learning Proceedings 1994, pages 284-292. Elsevier, 1994.
[2] Sarath Chandar, Sungjin Ahn, Hugo Larochelle, Pascal Vincent, Gerald Tesauro, and Yoshua Bengio. Hierarchical memory networks. arXiv preprint arXiv: 1605.07427, 2016.
[3] Wenhan Xiong, Jingyu Liu, Igor Molybog, Hejia Zhang, Prajjwal Bhargava, Rui Hou, Louis Martin, Rashi Rungta, Karthik Abinav Sankararaman, Barlas Oguz, et al. Effective long-context scaling of foundation models. arXiv preprint 2rXiv:2309.100 16039, 2023.
[4] Julian Richard Medina and Jugal Kalita. Parallel attention mechanisms in neural machine translation. In 2018 17th IEEE international conference on machine learning and applications (ICMLA). pages 547-552. IEEE, 2018
[5] Zhihao Jia, Matei Zaharia, and Alex Aiken. Beyond data and model parallelism for deep neural networks. Proceedings of Machine Learning and Systems, 1:1-13, 2019.
[6] Deepak Narayanan, Mohammad Shoeybi, Jared Casper, Patrick LeGresley, Mostofa Patwary, Vijay Korthikanti, Dmitri Vainbrand, Prethvi Kashinkunti, Julie Bernauer, Bryan Catanzaro, al. Efficient large-scale language model training gpu clusters using megatron-Im. In Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, pages 1-15, 2021.
[7] Norman P Jouppi, Cliff Young, Nishant Patil, David Patterson, Gaurav Agrawal, Raminder Bajwa, Sarah Bates, Suresh Bhatia, Nan Boden, Al Borchers, et al. In-datacenter performance analysis of tensor processing unit. In Proceedings of the 44th annual international symposium on computer architecture, pages 1-12, 2017.
[8] Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate. arXiv preprint arXiv:1409.0k4 0000, 2014.
[9] Kyunghyun Cho, Aaron Courville, and Yoshua Bengio. Describing multimedia content using attention-base based encoder-decoder networks. IEEE Transactions on Multimedia, 17(11):1875- 1886, 2015.
[10] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Eukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[11] David Silver, Julian Schritwieserr, Karen Simonyan, Ioannis Antonoglou, Aja Huang, Arthur Guez, Thomas Hubert, Lucas Baker, Matthew Lai, Adrian Bolton, et al. Mastering the game of go without human knowledge. nature, 550(7676):354-359 2017.
[12] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
[13] John Schulman, Barret Zoph, Christina Kim, Jacob Hilton, Jacob Menick, Jiayi Weng, Juan Felipe Ceron Uribe, Liam Fedus, Luke Metz, Michael Pokorny, et al. Chatgpt: Optimizing language models for dialogue. OpenAl blog, 2022.

---

[14] Matthew Riemer, Sharath Chandra Raparthy, Ignacio Cases, Gopeshh Raaj Subbaraj, Maximil- ian Puelma Touzel, and Irina Rish. Continual learning in environments with polynomial mixing times. In Advances in Neural Information Processing Systems, 2022.
[15] Michael Kearns and Daphne Koller. Efficient reinforcement learning in factored mdps. In IJCAI, volume 16, pages 740-747, 1999.
[16] Dheeraj Nagaraj, Xian Wu, Guy Bresler, Prateek Jain, and Praneeth Netrapalli. Least squares regression with markovian data: Fundamental limits and algorithms. Advances in neural information processing systems, 33:16666-1666 16676, 2020.
[17] Wesley A Suttle, Amrit Bedi, Bhrij Patel, Brian M Sadler, Alec Koppel, and Dinesh Manocha. Beyond exponentially fast mixing average-reward reinforcement learning via multi-level monte carlo actor-critic. In International Conference on Machine Learning, pages 33240-33267. PMLR, 2023.
[18] Shi Dong, Benjamin Van Roy, and Zhengyuan Zhou Simple agent, complex environment: Efficient reinforcement learning with agent states. Journal of Machine Learning Research, 23 (255):1-54, 2022.
[19] Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling, Advances in neural information processing systems, 34: 15815554-55097,
2021.
[20] ML Puterman. Markov decision processes. 1994. Jhon Wiley & Sons, New Jersey, 1994.
[21] Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. 2018.
[22] Yi Wan, Abhishek Naik, and Richard S Sutton. Learning and planning in average-reward markov decision processes. arXiv preprint arXiv:2006. 16318, 2020.
[23] Dimitri Bertsekas. A new value iteration method for the average cost dynamic programming problem. SIAM journal on control and optimization, 36(2):742-75, 1998.
[24] Michael Kearns and Satinder Singh. Near-optimal reinforcement learning in polynomial time. Machine learning, 49(2):209-233 2002.
[25]
Torgny Lindvall. Lectures the coupling method. Courier Corporation, 2002.
[26] Edward Sondik. The optimal control of partially observable markov processes over the infinite horizon: Discounted costs. Operations research, 26(2):282-304, 1978.
[27] Khimya Khetarpal, Matthew Riemer, Irina Rish, and Doina Precup. Towards continual rein- forcement learning: A review and perspectives. arXiv preprint arXiv: 13490, 2020.
[28] Bram Bakker. Reinforcement learning with long short-term memory. Advances in neural information processing systems, 14, 2001.
[29] Maxime Bhoveeier-Boisvert, Bolun Dai, Mark Towers, Rodrigo de Lazcano, Lucas Willems, Salem Lahlou, Suman Pal, Pablo Samuel Castro, and Jordan Terry. Minigrid & miniworld: Modular & customizable reinforcement learning environments for goal-oriented tasks. CORR, abs/2306. 13831, 2023.
[30] Anirudh Goyal, Alex Lamb, Jordan Hoffmaan, Shagun Sodhani, Sergey Levine, Yoshua Bengio, and Bernhard Schölkopf. Recurrent independent mechanisms. arXiv preprint wxXiv:199..1899 2019.
[31] Marwa Abdulhai, Dong-K Ki Kim. Matthew Riemer, Miao Liu, Gerald Tesauro, and Jonathan How. -optext-specific representation abstraction for deep option learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 36, pages 5959-5967, 2022.
[32] Ronald J Williams. Simple statistical gradient-following algorithms for connectionist reinforce- ment learning, Machine learning, 8:229-256, 1992.
12

---

[33] Thomas Dean and Keiji Kanazawa. A model for reasoning about persistence and causation. Computational intelligence. 5(2):142-150, 1989.
[34] Felipe M Santos, Leliane N Barros, and Felipe W Trevizan. Reachability- based model reduction for markov decision process. Journal of the Brazilian Computer Sociery, 21:1-16, 2015.
[35] Peng Dai and Judy Goldsmith. Topological value iteration algorithm for markov decision processes, In IJCAI, pages 1860-1865, 2007.
[36] Charles R. Harris, K. Jarrod Millman, Stéfan
der Walt, Ralf Gommers, Pauli Vir- tanen, David Cournapeau, Eric Wieser, Julian Taylor, Sebastian Berg, Nathaniel J. Smith, Robert Kern, Matti Picus, Stephan Hoyer, Marten H. van Kerkwijk, Matthew Brett, Allan Hal- dane, Jaime Fernández del Río, Mark Wiebe, Pearu Peterson, Pierre Gérard-Marchant, Kevin Sheppard, Tyler Reddy, Warren Weckesser, Hameer Abbasi, Christoph Gohlke, and Travis E. Oliphant. Array programming with NumPy. Nature, 585(7825) September 2020. doi: 00038/44158-000-2220
URL https: hww.ddi..oggg/g/ooo 10 1038 0000////2215560002222240000
[37] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high- performance deep learning library. In Advances in Neural Information Processing Systems 32, pages 8024-8035. Curran Associates, Inc., 2019. URL http: http://papers neurips cc/paper/ 1015-pytorch-- an- www-iiipratieete/eeee high- wwtpe--eerooooeeeeeeeeee pdf.
htgpapaeeeeeettteeeeeeeererrorre -earninnn------------------nn-ne library.
[38] Volodymyr Mnih, Koray Kavukeuoglu, David Silver, Andrei Rusu, Joel Veness, Marc G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. nature, 518(7540):529-533) 2015.
[39] Abhishek Naik, Roshan Shariff, Niko Yasui, Hengshuai Yao, and Richard $ Sutton. Discounted reinforcement learning is not optimization problem. arXiv preprint arXiv:1910.0. 140, 2019.
[40] Chris Nota and Philip S. Thomas. Is the policy gradient gradient? CoRR, abs/1906.07073.
2019. URL http:://rxiiii org/ abs/a//s666 07073.
[41] Saurabh Kumar, Henrik Marklund, Ashish Rao, Yifan Zhu, Hong Jun Jeon, Yueyang Liu, and Benjamin Van Roy. Continual learning as computationally constrained reinforcement learning. arXiv preprint arXiv:2307.04334 043444 2023.
[42] Anton Schwartz. reinforcement learning method for maximizing undiscounted rewards. In Proceedings of the tenth international conference machine learning, volume 298, pages 298-300 305, 1993.
[43] Huang Bojun. Steady state analysis of episodic reinforcement learning. Advances in Neural Information Processing Systems, 33:9335-9345, 2020.
[44] Richard S Sutton, Doina Precup, and Satinder Singh. Between mdps and semi-mdps: framework for temporal abstraction in reinforcement learning. Artificial intelligence, 112(1-2): 181-211, 1999.
[45] Pierre-Luc Bacon, Jean Harb, and Doina Precup. The option-critic architecture. In Thirty-First AAAI Conference on Artificial Intelligence, 2017.
[46] Matthew Riemer, Miao Liu, and Gerald Tesauro. Learning abstract options. NIPS, 2018.
[47] Matthew Riemer, Ignacio Cases, Clemens Rosenbaum, Miao Liu, and Gerald Tesauro. On the role of weight sharing during deep option learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, pages 5519-5526, 2020.
[48] Clemens Rosenbaum, Tim Klinger, and Matthew Riemer. Routing networks: Adaptive selection of non-linear functions for multi-task learning. International Conference on Learning Representations, 2018.
13

---

[49] Clemens Rosenbaum, Ignacio Cases, Matthew Riemer, and Tim Klinger. Routing networks and the challenges of modular and compositional computation. arXiv preprint arXiv:1904.12774 2019.
[50] Modjtaba Shokrian Zini, Mohammad Pedramfar, Matthew Riemer, Ahmadreza Moradipari, and Miao Liu. Coagent networks revisited. arXiv preprint arXiv:2001.10474 2020.
[51] David Abel, Nate Umbanhowar, Khimya Khetarpal, Dilip Arumugam, Doina Precup, and Michael Littman. Value preserving state-action abstractions. In International Conference on Artificial Intelligence and Statistics, pages 1639-1650. PMLR, 2020.
[52] Khimya Khetarpal, Martin Klissarov, Maxime Shevalier-Boisverr Pierre-Luc Bacon, and Doina Precup. Options of interest: Temporal abstraction with interest functions. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, pages 4444-445 2020.
[53] Khimya Khetarpal, Zafarali Ahmed, Gheorghe Comanici, and Doina Precup. Temporally abstract partial models. Advances Neural Information Processing Systems. 34:1979-1999 2021.
[54] Andrei Cristian Nica, Khimya Khetarpal, and Doina Precup. role of attention in hierarchical reinforcement learning. In NeurlPS'22 Workshop on All Things Atlention: Bridging Different Perspectives on Attention.
The paradox of choice: On the
[55] Kha Pham, Hung Le, Man Ngo, Truyen Tran, Bao Ho, and Svetha Venkatesh. Generative pseudo-inverse memory. In International Conference Learning Representations, 2022.
[56] Payel Das, Subhajit Chaudhury, Elliot Nelson, Igor Melnyk, Sarath Swaminathan, Sihui Dai, Aurélie Lozano, Georgios Kollias, Vijil Chenthamarakshan, Soham Dan, et al. Larimar: Large language models with episodic memory control. arXiv preprint arXiv:2403. 1901,, 2024.
[57] Tsendsuren Munkhdalai, Manaal Faruqui, and Siddharth Gopal. Leave no context behind: Efficient infinite context transformers with infini-attention. arXiv preprint arXiv:2404.07143, 2024.
[58] Bingbin Liu, Jordan Ash, Surbhi Goel, Akshay Krishnamurthy, and Cyril Zhang. Exposing attention glitches with flip-flop language modeling. Advances in Neural Information Processing Systems, 36, 2024.
[59] Kiran Voderhobli Holla, Chaithanya Kumar, and Aryan Singh. Large language models aren't all that you need. arXiv preprint arXiv:2401.00698 2024.
[60] Ziwei Xu, Sanjay Jain, and Mohan Kankanhalli. Hallucination inevitable: An innate limitation of large language models. arXiv preprint arXiv:2401.11817 2024.
[61] Yuan Tian and Tianyi Zhang. Selective prompt anchoring for code generation. arXiv preprint arXiv:2408.0911 09121, 2024.
[62] Zizhong Li. Haopeng Zhang, and Jiawei Zhang. Unveiling the magic: Investigating attention distillation in eetrieval-auggennee generation. arXiv preprint arXiv:2402. 11794, 2024.
[63] Thomas Schmied, Fabian Paischer, Vihang Patil, Markus Hofmarcher, Razvan Pascanu, and Sepp Hochreiter. Retrieval-augmenter decision transformer: External memory for in-context arXiv preprint arXiv:2410.07071 2024.
[64] Tianzhu Ye, Li Dong, Yuqing Xia, Yutao Sun, Yi Zhu, Gao Huang, and Furu Wei. Differential transformer. arXiv preprint arXiv:2410.05255, 2024.
[65] Jason Weston and Sainbayar Sukhbaatar. System 2 attention (is something you might need too). arXiv preprint arXiv:231 1829, 2023.
[66] Richard S Sutton, Anna Koop, and David Silver. On the role of tracking in stationary envi- ronments. Proceedings of the 24th international conference on Machine learning, pages 871-878, 2007.
14

---

[67] Suhas Kowshik, Dheeraj Nagaraj, Prateek Jain, and Praneeth Netrapalli. Streaming linear system identification with reverse experience replay. Advances Neural Information Processing Systems, 34:30140-30152, 2021.
[68] Matthew Riemer, Ignacio Cases, Robert Ajemian, Miao Liu, Irina Rish, Yuhai Tu, and Ger- ald Tesauro. Learning to learn without forgetting by maximizing transfer and minimizing interference. arXiv preprint arXiv:1810.11911 2018.
[69] Matthew Riemer, Tim Klinger, Djallel Bouneffouf, and Michele Franceschini. Scalable recol- lections for continual lifelong learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pages 1352-1359, 2019.
15

---

Appendix Overview: In this Appendix we provide important details that there was not enough space include in the main text. First Appendix A, we provide detailed proofs for the key theoretical results of our paper. Next in Appendix B, we provide additional details about our experiments discussed in the main text. In Appendix we further justify why we chose the particular setting of average reward RL for exploration in this paper. Finally, in Appendix D we discuss follow up research directions that could be considered in order to impact bounds on the mixing time.
A Deriving Upper Bounds for the Mixing Time
Section Overview: In this section, we provide proofs for Theorem and Theorem in the main text. We also highlight Corollary 2, which is key extension of Theorem I that relaxes the definition of 3 to be over c-steps rather than single step. Moreover, we highlight Lemma 1, which intermediate result to bridge the gap between the results in Theorem I and Theorem 2. For each theoretical result, we first state the result and then provide detailed proof. Our proofs build off the definitions and notation established in the main text.
key
For clarity, we will briefly recap the key assumptions discussed in the main text and why each is needed:
1. Unichain Policies: All stationary policies are aperiodic and unichain, meaning they give rise Markov chain with a single recurrent class that is recurrent in the Markov chain of every policy.
Needed for Theorem 1, Corollary 2, Lemma 1, and Theorem 2.
This assumption is necessary given definitions of the mixing time, the average reward, and coupling time the main text. For example, it is otherwise possible that two copies of the Markov chain originating from different start states may never couple (even in the infinite limit).
2. Multidimensional State Space: Each state is separated into n variables i.c., 8 = [st ..... 8"] with s² E S¹ for all variables i {1...., n) such that S C S' x x S" and T T(s' |s
Il (1 Tss" a).
Needed for Theorem 1, Corollary 2, Lemma 1, and Theorem
This assumption is necessary to make it meaningful to talk about the Markov chain induced by policy DBN over state variables and the strongly connected component structure of such DBN.
3. History Windows Depend on Subsets of State Variables: Observation o, action a, and reward r are caused by of a subset of the state variables such that Par(o) S {I, n). Par(a) {1, (1, n.),
and Par(r) S [I,,..,,. n... As result, any interaction history window (k) of size k is also caused by a subset of the state variables such that (h) (L,.,, n).
Needed for Theorem 2 only because it is the only result in this section concerned with
policies over (k).
This assumption is not strictly necessary state in the sense that it also covers the comne case where (k) depends all variables and thus must always hold in practice. However, it is worth emphasizing because if we do not actually see a change in the causal subset of variables as k decreases, then our bounds on the mixing time do not actually get tighter but just stay the same. As a result, this helps us build an intuition for the kinds of problems where the bound tightening showcased in Theorem 2 is meaningful (even though the result of Theorem 2 is still valid as stated even when it is not).
A.1 Understanding the Generality of our Assumptions
Multidimensional State Space: The assumption that there are n state variables s [s] s subsumes the corner case where n = 1 leaving only one variable. Moreover, the assumption that T(s'|s, a) [I ((<e.....),(('(| a)) also is not restrictive because it is only required that this structure hold for some representation of the state variables. For example, if two variables s S and d S do not involve independently of each other, new variable can be made to represent the
16

---

combination of these variables S E S 5" x S In the worst case, the system reduces to n = 1 where there is a single discrete variable evolving on its own. If this is the case, it does not invalidate the theoretical results of Theorem and Theorem 2. However, as explained in the paragraph on When Really Matters in Section 3, when n = the bound on the mixing time does not change with & as there is only one strongly connected component and the inequality from both theorems is simply an equality. As such, our results are most meaningful in environments with local observation structure, where are actually quite common for real-world and large-scale applications.
DBN Structure: Multidimensional transition functions can be represented by set of DBNs including a separate DBN for each action [33]. The DBN for any action a E A is an acyclic directed graph with two layers. The first layer the set of state variables in the current state and the second layer is the set of state variables in the next state [34]. The Markov chain induced by given policy can be seen as DBN formed as mixture of the DBNs for each action according to the policy and with the actions produced by the policy also critically being causally dependent on certain state variables. The assumption that there are strongly connected components TÝ. ..... l is also not restrictive because {= 1 is corner case of this formulation. Indeed, the generality of this phenomena underlies the generality of planning algorithms such as Topological Value Iteration (TVI) [35], which efficiently applies Value Iteration (VI) on each strongly connected component in reversed topological order. TVI is indeed known to cover the case when = 1 [34]. Again, the case where l = 1, which also happens by definition whenever n 1, is not very interesting for our theory, but does not invalidate our results.
Motivating Examples: To motivate the generality of these assumptions in real-world settings, we will walk through how both of the domains used our experiments naturally possess the desired structure to make our bounds in Theorem 1 and Theorem 2 meaningful despite their simplicity. See Figure for visual depictions of the two domains.
RGB World (Figure 6a): This domain can be represented as a multidimensional state space with the RGB color of each cell in the line representing its own state variable. The colors of each state variable is only possibly influenced by the variable on either side of it (depending on the chosen action of the agent). As result, it is clear that the variables in the transition graph can be considered to evolve independently of each other given the previous state and action while displaying sparse interaction structure among variables. As such, the number of strongly connected components is highly dependent on the policy parameterization. Moreover, the observations are local to the agent, and are only the causal result of the agent's 's surrounding state variables in a single observation with multiple observations needing to be combined to reflect the full state space.
Minigrid Crossing (Figure 6b): This domain can be represented as a multidimensional state space with each grid cell representing its own state variable. Each variable can either be the location of the agent, wall, an empty space or the location of the goal. Each variable can then only be influenced by the variables adjacent the the variables in the transition graph can be considered to evolve independently of each other given the previous state and action while displaying sparse interaction structure among variables. As a result, the number of strongly connected components is yet again highly dependent on the policy parameterization. The observations are also yet again local to the agent, and only are causally dependent on the state variables within limited radius of the agent's location. To reflect all state variables, it is indeed necessary that multiple observations are combined.
across a single time step. Thus
A.2 Proof of Theorem
Theorem (Strongly Connected State Variables Bound): If the Markov chain T (s'ls) induced by policy has E maximal strongly connected components in DE with maximum size of g state variables and a minimum of 3 common probability mass between any two state configurations for any state variable, the mixing times (e)(e) and th
(c) can be upper bounded:
Au (e) E t (c) E (
1 39
1og(1//))
17

---

Proof Sketch: Due to the sorting of the strongly connected components, our analysis is based on coupling each of the I in succession. Because it is possible that multiple I couple at the same step, every step where the Markov chain does not fully couple must be a step where some ry does not couple. Our proof proceeds in the following high-level steps:
1. The probability of 1 coupling at a given step once 1 "T,...,""" 17-1 have all already coupled is 2 ,
2. Thus the probability of T not coupling at step when [T 1 TT_1 have all already coupled
is S (1 3.)
3. So the joint probability of I not coupling for m, 2 0 steps when II....,T_. have all already coupled is S (1 - mm
4. If > m then Ei= m, = m and the probability that m-steps have been spent not coupling in some 1T has probability bound independent of the particular allocation of m into m.;
Thus we can conclude that P(r > m) S (1 - 39mm.
5. Leveraging the identity that 1 - x Se e for x 2 0, we find that P(T > m) S (c ymm
6. The Markov chain is c-mixed if P(T > m) S c, so it must be e-mixed if (c-00) ym S c. which implies that m 2 ga [og(c/)).
7. Finally, we note that 1 (t) E E following Lemma of Kearns and Singh [24].
Additional Details: Due to the sorting of the l strongly connected components, our analysis will be based on stabilizing the 17 's in succession as in Kearns and Koller [15]. For example, if we assume that TT. .... r 1 have all stabilized by time I, all the variables in I must then couple at the same time for i to stabilize. This event then happens at time with probability 2 ,3. As soon as i stabilizes, we can move on to stabilizing i + 1. When all é strongly connected components have stabilized, we are done and have surpassed the coupling time. An important subtlety to note is that this process need not take at least É steps and will in fact complete in just a single step with probability 2 (39)². It is not that strongly connected component can't 't couple until i 1 has coupled at the previous step, but rather that if i has coupled and i 1 has not yet stabilized, this coupling over a subset of variables is not necessarily meaningful for i stabilizing.
To formalize this idea, let us consider that each strongly connected component was the focus of the successive stabilization for m; 2 O steps without stabilizing yet. If the entire Markov chain across all & components has been running for m steps and has not yet fully stabilized, we note that Ei= m, = m. This is because every step where a component did stabilize, either the next component i + did not stabilize at the same step, each stabilizes immediately i.e. m, = 0 until the next component j + 1 does not stabilize. Note that components have stabilized after m steps that implies that my
consider series of components j > where
if only still describes the setting where not all components actually have the opportunity to stabilize despite my being defined for all E {1, (L.,.. f). The probability that the system has not stabilized by m steps P(T > m) is then equal to the probability across all allocations of m. that each component i did not couple for m, 2 steps when all components j < had already stabilized. The probability of i also stabilized at or before that step is 2 33, so the probability of not stabilizing at a given step is I (1 (1-33 3, and the probability of not stabilizing for each of my steps is < (1 33)m; Then we can consider the joint probability of this across all é components to upper
for all j >>> i, this
stabilizing at given step if
bound P(T > m)) noting that the total probability of all possible allocations of m, must be < 1:
P(r > m) S II (1 - 39)m = (1-39)Eis Eo m. = 31- S (0- 1=1
The dependence between the strongly connected components in the joint probability is incorporated through interdependence between the m; allocations in the product. However, we see that this interdependence and dependence on the number of strongly connected components goes away when noting that Σi m, In the final inequality of the above equation, we note that - z S -=
18

---

for x 2 O. The system is e-mixed after m steps if P(T > m) - therefore by the transitive property it is then also e-mixed if:
(e (m c.
(2)
We then proceed by taking the logarithm of both sides of this equation:
m(-39) S log(e).
(3)
We can then multiply each side by negative one, which also reverses the inequality:
m39)) 2 (og()/e).
(4)
Finally, we divide both sides by 39 and yield the result that must e-mixed if:
m 2 33
log(1/c).
(5)
This then implies that:
mix (e) o (
(6)
The final step of the proof is simply that the result from Lemma of Kearns and Singh [24] that mix(e) E 12 (((r(((e)) for any policy æ.
A.3 Proof of Corollary 2
Corollary (Strongly Connected State Variables Bound with c-Step Transitions): If the Markov chain T (s'(s) induced by policy has & maximal strongly connected components in DD with maximum size of state variables and minimum of Be common probability mass between any two state configurations for any state variable after 2 1 steps, the mixing times the (c) and t ( can be upper bounded:
tn (c) t (c) E of c log(1/e)
32
potential criticism of 33 from [15] is the potential for 3 = 0 when there are deterministic elements of the transition dynamics. However, this concept could be relaxed by considering the c-step transition dynamics instead of just step. For policies following Assumption 1, there must be a common recurrent of states that all policies experience regardless of the starting state, so values must become greater than 0 for appropriately large values of c. As such, we define Be a parameter that defines the minimum amount of common probability mass for any two state configurations for any state variable after c-steps so that 3r = Bi.c E [0, 1]:
Bi. min u.w'ES
Σ S'ES,
min (s (Tl,
is easy to show this because the probability of a strongly connect component coupling at each step the same in the proof of Theorem 1, but we must simply account for the difference in definitions of m' = mc and Be 3 because the probabilities are function of steps now rather than step before. So we can simply plug these substitutions into our equations from Theorem I show that m 2 de (og(e)). The result is then a simple multiplication of the bound from Theorem by given the new definition of B..
19

---

A.4 Proof of Lemma
Lemma (Conditioning on Subsets of State Variables): Consider the Markov chains induced by two policies drawn from different policy classes. The first policy is <(-|xx), which takes n(t) states variables s² .... s*(()) input time t. has l maximal strongly connected components with maximum size of g state variables and a minimum of 3 common probability mass between any two state configurations for any state variable. The second policy π (-|x'). only takes some arbitrary subset n' (t) S nı(t) state variables as input at any moment in time i such that x' =s²,, ss,,.., s" (t) C s² :.... st((t) has ll maximal strongly connected components with maximum size of state variables and a minimum of 3 common probability mass between any two state configurations for any state variable. The mixing times the policy classes of π. and T' can then be upper bounded such that the bound on nO is tighter:
(e)(e) (m) E ( (0gg(0)())) EO ( [g
The policy class over x' removes edges in the DBN that were previously present in the class over and cannot add any. If these edges were within pre-cxisting maximal strongly connected component, that component may be split into multiple smaller strongly connected components, otherwise the number of strongly connected components stays the same. This implies that 2 f. If removing edges resulted in splitting the largest strongly connected component, we say that g' < g and = g' otherwise. This implies that 2 g'.
To understand the relation between 33 and 3' we must make our notation bit more detailed than what was presented in the main text. 3r can then be stated as follows while making the parameterization of explicit:
3: X(A)x((A)
min
min $in
min
Σ x (()))
Σ (18) agEA
(8)
Likewise, 3 can be defined in a similar fashion:
a2) }
8 = min ((())(((A)
min
min Σ
Σ n' (a2)<66)) (6 agEA
#1.82ES
R.EA
(n(() and X' St ... S (t) such that X' C X(t) for all while á and z'j are analogously obtained from the full state sj for j E [1,2}. The main difference between these formulations is then minimizing over (() A(A) in the case of B; and x'( x A(A) in case of 8. By rule of Cartesian products over subsets, we then know that where (t)) St x
(1) ((A) C x(1) (AA) and because the minimization over subset must yield a result larger or the same size as a minimization of the full set, 8, I 3 for all state variables i, This also implics that B < 8 it also must hold for the smallest value over state variables.
Bringing it all together, 0 < 3 < B' < 1, and 2 o' > 1, so 1/3'd S 1/39. So, when combined with the results from Theorem this result has been proven.
A.5 Proof of Theorem
Theorem 2 (Limiting Mixing Times with Context Length): If the Markov chuin induced by policy conditioned on a finite interaction history window x ()) with a context length of k has Eg maximal strongly connected components in D" with maximum size of ge variables and minimum of Bx
20

---

common probability mass between any two state configurations, the mixing times t (c) and mix (E) can be bounded for any k' k:
(E) E mix
E o
1
log (og())))
O
log(/e)
300
where the dependence on the context length implies that 0 S Br S Bx S and gk 2 gk 2 1.
Proof Sketch: Lemma considers the mixing time relationship of policy classes conditioned on subsets of the state variables that other policy classes are conditioned on. Our proof includes the following high-level steps by applying our notation from Section 3 in which k' 2 k for all t to the results of Theorem and Lemma
1. We consider the causal parent state variables of each observation, action, and reward to conclude that (ar(h(k) S (a)(h())) S (1,..,n), which implies that n 2 (t) 2 (t))).
2. Through Lemma we show that by rule of Cartesian products over subsets that S Bi S
3x S 1.
3. Through Lemma I we also demonstrate that gk 2 gs 2 because causal connections in DD are only added and not removed when the context length is increased.
4. This then implies that 1/ 2 1/3ge, which is sufficient to prove Theorem 2 using Theorem because is independent 'f
Additional Details: As mentioned in the section overview, we assume that observation o, action
a, and reward r are caused by of a subset of the state variables such that Par(o) S {1, n). Par(a) S {1,.., n), and Par(r) S [1....n) n). As result, any interaction history window (R) of size is also caused by a subset of the state variables such that Par(h S {L, n). This is always true in practice as this includes as corner case the situation where the interaction history is dependent on all variables. However, our bound is not meaningful in domains where (a)( {I,.., n for all values of & because then it never gets tighter as gets smaller. In general, finite context length is reflective of a maximum ny (t) sized subset of the state variables at a time where n 2 ny (t) 2 n*( if context length k' 2 k for all t. This is because Par(h C Par(n(()) S {1,... n). Theorem 2 then follows directly from the combined results of Lemma and Theorem 1.
Specifically, 0 < 8 < 32 < Bx < 1 because they are the result of minimizing the same function over progressively smaller subsets as Lemma 1. Moreover, g 2 9 2 Ib 2 1 because causal connections in the DBN are only removed as in Lemma 1. This then implies that 1/ 2 1/5 > 1/ which demonstrates the relative tightness of the bounds. Finally, Theorem 1 does not depend on the parameterization of the policy, so it can equally be applied to policies of history length & to bound the mixing time.
B Additional Details for Experiments
Section Overview: Here we elaborate on experimental details that we did not have room to include in the main text. each subsection, we fill in missing details from Sections 2, 3, and respectively.
Compute Infrastructure: Our experiments were deployed on a cluster of Intel x86 machines. Each of our toy experiments and tabular experiments were run on single CPU. Meanwhile, the online function approximation and Decision Transformers experiments were each run with one V100 GPU.
Software Libraries:
Simple tabular models were coded from scratch primarily using Numpy [36], which is publicly available following a BSD license. Neural network models were developed using Pytorch [37], which is publicly available following modified BSD license.
Hyperparameter Tuning Protocol: Our models train to convergence in all experiments. For our simple RGB world experiments, we did search over Adam learning rates with gradient clipping
21

---

until finding one that consistently converged at every context length. For our Decision Transformers experiments, we followed the code from the original paper by Chen et al. [19], which included gradient clipping, Adam optimization, and learning rate schedule that consists of linear warm-up following by cosine learning rate decay. These details were copied from the original paper [19] and we also did grid search over the initial learning rate.
B.1 Toy Examples (Section 2)
Mixing Time Calculation: Following Riemer et al. [14], we compute a relaxed version of the c- return mixing time that is averaged over start states to be more reflective of mixing times encountered rather than the worst case. This computed by first rolling out the policy for 100,000 steps in the environment to approximate the reward rate from the first state in our list. This state is arbitrarily chosen such that all variables are set to index 0. Then for each environment state we rollout the policy for 10,000 steps, recording the final step at which the environment is not within of this reward rate. The reported mixing times in the main text are an average across all possible starting states. For the smaller irrelevant variables example, we also compute the mixing time for each starting state 5 times and take the average over trials in order to lower the variance of our estimates.
Irrelevant Variables Example Details: For the irrelevant variables example, there are 9 total states across variables z and y and only total states when ignoring y. As a result, there are 2 512 deterministic policies across both variables and 2 = 8 deterministic policies over the relevant variable. 0.474 the highest reward rate estimate for any policy, so if the reward rate has not mixed to precision of e = 0.05, there must be significant bias and greater than 10% relative error in the reward rate estimation regardless of the policy.
Independent Subtasks Example Details: For the independent subtasks example, there are 18 total states across all variables æ, y, and and only 6 total states when ignoring the irrelevant variable for the subtask. As result, there are 2 = 262,, 144 deterministic policies across all variables and 2 = 64 deterministic policies over the contextually relevant variables. 0.474 is again the highest reward rate estimate for any policy, so if the reward rate has not mixed to precision of c 0.05, there must be significant bias and greater than 10% relative error in the reward rate estimation regardless of the policy. As the search space over deterministic policies considering all variables is very large for this domain, we randomly sort the policies and only evaluate policies for 24 hours on a single cpu. 1 1,597 total policies were considered. Notice that this means the numbers reported in the main text are conservative as less than an arbitrary 1% of the search space was actually explored to find these values.
B.2 Online Learning in a POMDP (Section 3)
Statistical Significance: The lightly shaded regions around each data point in Figures 4, 5, and represent the 95% confidence interval of the average mixing time of that point given the number of data points found at that particular binning with respect to the reward rate across runs and seeds. This is computed as 1.96. /ni where is the standard deviation and n is the number of data points at particular binning of the reward rate.
For each setting considered, we conducted a version of the experiment with 100 different seeds ranging from 0 to 99, Our tabular experiments were run for I billion steps (corresponding to about 7 days for the highest context length k 10). Meanwhile, our function
Experimental Protocol:
(b) Offfine Learning: Minigrid Crossing Figure 6: Domains. Figure 6a depicts the environment used for online learning with left and right actions. The reward is +1 if the agent is the blue square and -1 the red square. If either square reached, the agent returns to the center while the red and blue squares are randomly assigned to the edges. Figure 6b depicts random episodic configuration of the environment used for offline learning. Observations are denoted by yellow dotted lines or light shading.
(a) Online Learning: RGB World
22

---

approximation experiments were run for 6 hours each, corresponding to about 1 million steps for the Transformer models, million steps for the MLP models, and 600 thousand steps for the LSTM models. We set the number of green squares to 3, which implies that the k = 1 policies cannot achieve the optimal policy, but all k 2 2 can. Based on the results of preliminary testing, the SGD learning rate was set to 0.01 in our tabular experiments and the Adam learning rate was set to 0.000001 in our function approximation experiments. Online Q-learning with an exploration rate of 0.05 and a discount factor of 0.99 was used in all experiments. During our function approximation experiments, the target network was updated every 10,000 steps following best practices [38].
For our MLP models, the history representation is flattened at each step and sent to a two layer MLP model with 724 hidden units in each layer and linear layer deriving an output for each action's 's value. ReLU activations are used at each hidden layer. Our LSTM model sent a sequence of dimensional inputs for each step where dimensions represent the observation, one dimension represents the action, and one dimension represents the reward. The current observation then sent the LSTM with an action of 1 and reward of in order to preserve the dimensionality. The LSTM, implemented with the Pytorch [37] LSTM library, is unidirectional and has two hidden layers of size 256 followed by a linear layer to produce a value for each action. The input sent to the Transformer model has 6 dimensions at each step, adding to the LSTM input an encoding of the ordering of the input ranging from 1 to k. Our model is implemented using the Pytorch [37] Transformer library leveraging a hidden dimension of 256, one encoder layer and one decoder layer. The output of the Transformer is once again sent to a linear layer producing a value for each action. The number of heads was set to 8 based on preliminary experiments.
Neural Network Architecture Details:
We adapt our mixing time estimation approach from the last section as we cannot easily generate the full set of possible histories corresponding to each state. Instead we rollout single chain from the current state for 10,000 steps using the greedy exploitation policy, At each step along the way, we calculate the estimated reward rate and report the mixing time as the last time where the estimated reward rate is not within = 0.01 of the full estimate over all of the steps. This is similar to the approximation used by Riemer et al. [14] for larger scale domains.
Mixing Time Calculation:
Mixing Times Across Context Lengths: Our results in the main text only plot the mixing time for the context length & = 5. We provide more comprehensive results across values of k in Figure 7. These additional settings follow a similar overall trend. Reward rates are binned by increments of 0.0333.
Parameter Scaling Experiments: We varied the number of hidden units in each architecture in order to test the effect of model size on encountered mixing times over 100 random seeds and the same learning configuration.
MLP Scaling, For the MLP model we considered: 256 hidden units resulting in 134,402
total parameters (at & 2). 404 hidden units resulting in 331,686 total parameters (at h = 2), 724 hidden units resulting in 1,057,766 total parameters (at k = 2). 1280 hidden units resulting in 3,293,442 total parameters (at k- = 2), and 2250 hidden units resulting in 10, total parameters (at h = 2). While the MLP model does grow with the context length, most parameters are not in the first layer and this only results in 8.7% relative parameter growth from context length of I to 10. We considered both k = 2 and k = 10 for each number of hidden units and found no statistically significant effect varying the number of parameters while keeping k fixed.
LSTM Scaling. For the LSTM model we considered: 92 hidden units resulting in 137,634 total parameters, 143 hidden units resulting in 330,618 total parameters, 256 hidden units resulting in 1, 1,054,722 total parameters, 454 hidden units resulting in 3,308,754 total param- eters, and 792 hidden units resulting in 10,055,234 total parameters. We again considered both k: = 2 and & = 10 for each number of hidden units and found no statistically significant effect varying the number of parameters while keeping l fixed.
Transformer Scaling, For the Transformer model we considered: 88 hidden units resulting in 127,338 total parameters, 136 hidden units resulting in 301. 20,2242 total parameters, 256 hidden units resulting in 1. 1,058,562 total parameters, 448 hidden units resulting in 3,228,738 total parameters, and 792 hidden units resulting in 10,067, 1 14 total parameters. We plot the effect on the average mixing time encountered during learning in Figure 8. We use this aggregate metric in this case only to simplify our analysis given that we are also varying the number of parameters, but it is definitely preferable to plot mixing times as a function
23

---

Context Length of
(a) Context Length of 2
4009
MLP
LSIM
2500-
3400
Average Mixing Time
Average Miting Tine
2000
2009
1500
1000
MILP
-0.2
-0.1
60
0.2 Reward Raie
0.1
0.3
0.4
0.5
0.2
0.1
0.0
0.1
0.2
0.3
0.4
0.5
Reward Rate
(c) Context Length 4
Context Length of
Tabular
3000
3000
LSTM
2500
2500
Avenage Mixing Time
Average Mixing Time
2000
2000
1300
1.000
1000
1000
Tabulae MLP
500
Transfrrer
-0.1
0.0
0.1
0.2
0.3
0.4
0.5
-0.1
0.1
0.2
0.5
044
0.5
Reward
Reward
(e) Context Length 7
(l) Context Length of
3500
Tabulax
Tobuler
3000
1000
LSTM
LSTM
Average Mixing Time
2500
Transformer
Average Mixing Time
2500
Trunsforer
2000-
2000
1.00
1000
0-
-.1
0.0
0.2
0.3
0.5
0.0
0.2
0.3
05
Reward
Figure 7: Policy Architecture vs. Encountered Mixing Times. We plot the average mixing time for each choice of policy architecture between tabular, MLP, LSTM, and Transformer models with averages binned by the reward rate of the policy.
24

---

2500
Average Mixing Time
k=2
2000
k-10
1500
1000
500
10"
10
10
Number of Parameters
Figure 8: Transformers: Average Encountered Mixing Time vs. Number of Parameters. We plot the average encountered mixing time during learning for Transformer models trained at number of different model sizes regulated by the number of hidden units. Averages are taken over 100 random seeds and we also provide 95% confidence intervals.
of the reward rate as we have primarily done throughout this paper. Because mixing times based the e-return mixing time depend strongly on the reward rate, aggregate metries like this largely wash out the effect of other factors. In light of this, the nearly 2x increase in the average encountered mixing time over the course of the entire learning period between & = 2 and & = 10 for Transformers with over 10M parameters is quite staggering and showcases the significant impact of the analysis in Theorem 2 even in this simple domain.
To test the hypothesis of the model size increase being correlated with paying more attention to the full context, we reran our experiments for 136 hidden units and 448 hidden units (10x more parameters) at k 10 while keeping track of the attention maps computed in the decoder with respect to the interaction history at past time-steps. During the mixing time evaluation phase, we averaged the attention maps computed at each step and then averaged over evaluations weighted by the number steps between evaluations. The average attention weights for the 136 hidden unit model were: / 9 0.048, f. - 8 : 0.052, L 7 0.061, /. 6 0.078, 0.212. In contrast, the average
Transformer Attention Map Analysis:
t t- 5 0.099, - 4: 0.123.,..... 123, f - 3 0.148, t - 2: 0.179, and f attention weights for the 448 hidden unit model were: - 0.089, - 8 0.080, t 7 0.078, t t-6 6 0.086, t - 5 0.102, f - 4 0.119, attention weight goes up for all time-steps in the interaction history window greater than 4 steps in the past. Indeed, the entropy associated with the average distribution grows over 4% from 2.99 bits for the 136 hidden unit model to 3.12 bits for the 448 hidden unit model.
0.133, t- 2: 0.149, and - 1: 0.163. The average
B.3 Offline Learning with Decision Transformers (Section 4)
Documentation for the SimpleCrrosinng99N// environment that we used for our experiments can be found at https::/minigrr/ farama org/ environments/ weaiitt///iiiiiiiii CrossingEnv/, which is publicly available following the Apache 2.0 license.
Environment Details:
We model our code and architecture following the code provided for the Atari experiments of the original Decision Transformers paper [19], which was released under an MIT license. Our architecture and optimization follows all of the details from their paper with the only exception being the use of a convolutional architecture that his been found more suited to the Minigrid domain where observations layers with filter sizes of 16, 32, and 64 respectively. 2 x 2 max pooling layers and ReLU activations follow each convolutional layer. The result a 64 dimensional embedding (as opposed to 128 for Atari) due to the smaller input size. Following Chen et al. [19], we used a batch size of 128, 6 layers, attention heads, epochs of training, GeLU nonlinearities within the Transformer, 512 20 warm-up tokens, 2 500000 * & final tokens, and dropout rate of 0.1. We also used the same Adam optimization with learning rate of 0.001, 3 = (0.9, 0.95), a gradnorm clip range of 1.0, and 0. weight decay. It is important to note that transformation is applied to the rewards before is sent to
Decision Transformer Details:
3 x dimensional [30; 31]. There are three convolutional
25

---

the Transformer model referred to as the return to go for policy or trial i namely Ii.t := Σ firr which is sent along with the current time step t instead. The advantage of this representation the ability to prompt the model for high returns to elicit high performing behaviors. We can consider this as special case of the paradigm which models take as input some transformation of the interaction history up to bounded context length.
Our random behavior policies take in the current observation input and process it with the same convolutional encoder as the Decision Transformer model. This representation is simply processed by linear layer, mapping then selects the argmax action. The maximum length of any episode is 325, the minimum length is 14, and the average length is 324.4 (as most random policy episodes do not arrive at the goal state).
Random Behavior Policies:
a value for each action. The policy
Our learning behavior policies are implemented with the REINFORCE algorithm [32]. It is quite common in the offline RL literature to use data drawn from replay buffer during learning. However, the correspondence of these to actual policies used can get complicated with an off-policy algorithms such as DQN or Rainbow. Mcanwhile, REINFORCE only performs updates after each episode and can ensure that the actions in each episode are drawn entirely from on-policy sampling. We opted for this approach due to its simplicity and to control for other potential complicating factors in the learning process such as policies that change or randomly explore in the middle of an episode. The learning rate of REINFORCE is set to 0.1 with SGD optimization, the entropy regularization coefficient is set to 0.001, and a discount factor is set to y = 0.99. The maximum length of any episode is 325, the minimum length is 14, and the average length is 240.5. We observe that by 10,000 episodes of training the policy is consistently solving the task, but we consider just the first 1,000 episodes to make the comparison more precise with the random behavior policy setting and to increase the variety in the behavior distribution (by putting some emphasis on early learning too).
Learning Behavior Policies:
We aim at a similar evaluation of the mixing time as the previous section. However, we also take advantage of the episodic structure of this problem promote computational efficiency. The policy is rolled out for 1000 episodes to compute the estimated reward rate per step and reward rate per episode. Then we consider 1000 random orderings of these episodes keeping track of both the final step not within e 0.01 precision of the estimated reward rate per step and the final episode not within e 0.01 precision of the estimated reward rate per episode. We report the average mixing time over these 1000 random orderings to again highlight mixing times experienced in practice rather than worst case mixing times.
Mixing Time Calculation:
Evaluation of Decision Transformer Models: In Figure 9 we take a deeper look at the general- ization ability of models trained with decision transformers as function of the context length k. We now elaborate a bit on our discourse from the main text. One may wonder if it is necessary for Decision Transformers to achieve near optimal training accuracy in order to achieve good downstream performance. To answer this question, we consider a Decision Transformer that we trained on random behavior policies (seed 0) that achieved optimal training accuracy at k = 25. This particular setting is interesting because this is the smallest ic found to be sufficient and because the deterministic hehavior policies allow us to identify the best possible training accuracy given the data (100%). We then evaluate this model across a variety of return to go prompts stepping by 0.01 from 0 to 1.0 and report the average performance of these models across 1, 1,000 episodes. In Figure 9 we plot the reward rate of these models ranked according to performance (along the axis) as well as each behavior policy used to generate the data evaluated the same way. We plot our results in Figure which validate our hypothesis about the importance of fitting the training data. Decision Transformer with k = 25 can achieve performance as good as any behavior policy used to generate the data. Meanwhile, models with smaller k fail to achieve the same performance and models with larger k seem to overfit dataset of this size.
C Why Average Reward RL?
Issues with Discounting: As typically implemented. the popular discounted reward setting of RL does not correspond to the maximization of any objective function over set of policies [39] and the policy gradient is not the gradient of any function [40]. Moreover, these fundamental issues do not resolve as the discount factor approaches [39] and discounting does not influence the ordering of policies, suggesting it likely has no role to play in the definition of the control problem [21]. In
26

---

0.20
Behavior Policies
k (59%)
0.15
(82%)
Reward Rate
k 10 (82%)
(1 (100%)
0.10
50 (100%)
100 (100%)
0.05
0.00
0.0
0.2 Policy Rank Percentile
0.4
0.6
0.8
1.0
Figure 9: Decision Transformers Context Length vs. Reward During Evaluation. We highlight learning results for Decision Transformers in the crossing environment with data generated from random behavior policies with a context length & of 1. This figure details the evaluation performance when prompting each policy with 100 return to go values stepping from 0 1.0. This is also compared to the actual distribution of performance for the behavior policies that generated the data. Each line color represents different context length & for the Decision Transformer model with the training accuracy rounded to the nearest percent included in parenthesis.
contrast, the average reward per step objective, is well-suited for analysis of partially observable [1], larger scale [14], and continual RL problems [21; 27; 41]. We thus adopt this formulation in our paper enable more rigorous analysis of agents.
Evaluation vs. Optimization: is important to note that we only use the average reward perspective evaluate agent's 's behavior and place restriction on discounting being used within the agent's learning process as it is in our experiments. So, to argue that the discounted setting would be more relevant for our goals, is essentially
propose that what practitioners actually care about is the discounted return rather than the undiscounted return or average reward of policies. As highlighted by Schwartz [42] it is quite rare for actual evaluation results reported in papers consider discount factors. Meanwhile, optimizing for the average reward actually also inherently optimizes for the undiscounted return [42]. That said, even if we did really care about the discounted return, this is corner case subsumed within the formulation of Theorem 2. This is because the discount factor or time-step number can be considered to be potentially unobserved element of the state space that is used to modulate the function that produces immediate rewards. As such, our choice of the average reward per step formulation does not at all limit the generality of our conclusions Theorem In fact, it helps ensure that the result is as general as possible.
Bojun [43] proves that every finite horizon task has a unique steady- state distribution under any policy. This means that Assumption 1 and Corollary are guaranteed for the case of episodic tasks, which is subsumed within our framework. For episodic tasks, the mixing time is related to the number of episodes a policy must be rolled out for before we get reliable measure of performance. Indeed, we consider episodic tasks within our experiments using the Minigrid Crossing environment in Section 4.
Infinite vs. Finite Horizons:
D New Architectures and Algorithms
Theorem 2 shows that what leads to potentially high mixing times is when models leverage monolithic representation that is highly sensitive to large part of the interaction history at all times. This is particularly descriptive of how vanilla transformers work, but there multiple already existing research directions that seem well suited to scaling to high context lengths while providing less history sensitivity at each step, which we will now outline.
Hierarchical RL: In hierarchical RL frameworks such as options [44] when applied to neural networks [45] with potentially many levels of abstraction [46] and/or complex weight sharing
27

---

[47; 48; 49; 50], it should be possible in domains with temporal coherence to approximate a policy with longer context length by multiple sub-policies with a smaller context length. This is particularly relevant for approaches to option learning that learn an independent lower dimensional representation space for each option as in [31], or associate options with subsets of the state space [51; 52; 53; 54].
Hybrid Transformer Working Memory Architectures: Recent work has aimed to improve the effective context length of transformers by augmenting them with some kind of bounded working memory component [55; 56; 57]. While these papers are typically solely motivated by computational efficiency, they effectively serve the role of extending Transformers longer contexts while making them less sensitive to experiences in this context that are not reflected in the memory. This effec- tively brings Transformers closer to some of the incremental design patterns of RNÑs, which our experiments indicate experience lower mixing times. Approaches that have less sensitivity to the input space are also motivated by the problem of attention dilution [58; 59; 60; 61], which typically addressed with retrieval augmented approaches that only focus on subset of the interaction history time [62; 63; 64] or generated small description [65].
Tracking Policies: If we aim to learn a non-stationary tracking policy solution concept as argued by [66] we could potentially represent a stationary policy over longer context length with non- stationary set of small context length policies. In effect, this becomes similar to what is achieved by the hierarchical RL policies described above. One additional subtlety to highlight in this case that this solution is also limited by the mixing properties of the tracking policy parameters, so it would also be necessary to tune this approach move through parameter space as fast as possible.
The Role of Replay:
While not related to sensitivity of the policy to its input, it is worth noting that replay based learning algorithms have been theoretically shown to lead to better sample efficiency of learning with respect to the mixing time [67; 16]. It will be an interesting direction for future work to explore the extent which these learning benefits can be combined with the benefits of particular choice of policy class aimed at lowering mixing times with smaller context length. Moreover, considering the importance of our analysis continual learning settings, it will be interesting to see if similar mixing time benefits could be established for approaches that combine replay with meta-learning [68]. or approximate replay through the use of strong generative models [69].
28

---

NeurIPS Paper Checklist
1. Claims
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Question:
Answer: [Yes]
Justification: At the end of the introduction we concretely list three primary contributions of our paper, with each being the topic of its own section (Sections 2, 3, and respectively).
The only other content in the paper is the concluding section, which discusses promising opportunities for future work build off our ideas.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results be expected to generalize other settings.
It is fine include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: We comprehensively discuss the theoretical limitations of the assumptions we make in this paper in Appendix A. We also discuss limitations in the scope of our main theorem in the paragraph titled "When it really matters" at the end of Section 3.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
29

---

Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: We provide discussion of all assumptions in the main text, which we recap for clarity the beginning of Appendix A along with a discussion of the generality of each assumption. Then detailed proofs are provided in Appendix A with each theoretical statement being devoted own subsection building off these assumptions.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justificatiin: To augment the presentation in the main text, all remaining details needed to reproduce our experiments are provided in Appendix
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often good way to accomplish this, but reproducibility instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, other means that appropriate to the research performed.
also be provided via detailed
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
30

---

(d)
We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: We have provided the code needed to reproduce our experiments and a com/
README file in the repository with instructions for running them at https://githhe
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips) ca/
selliic/gudess//e Soaidees/CooeSummiseeee heibbmmmssooooiccrr for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to reproduce the results. See the NeurIPS code and data submission guidelines (https: //nips.co wwtpppccccccccccccccccccccccccct guides/ /CodeSubmii ssionPolicy) for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: To augment the presentation in the main text, all remaining details needed to understand and even reproduce our experiments are provided in Appendix
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented the core the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: Every figure showing results throughout our paper provides enror bar shading so that readers can assess statistical significance. Only Figure does not include error bars and that is because that plot is taking deeper look single interesting random seed.
31

---

Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report I-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information the puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: We provide this information in the "compute infrastructure" paragraph the beginning of Appendix B.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (c.g., preliminary or failed experiments that didn't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https http://eeeup/peeeeee
Answer: [Yes]
Justification:
The authors have read the code of ethics and made every effort comply.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative socictal impacts of the work performed?
32

---

Answer: [NA]
Justification: Our paper contributes fundamental research on existing models, so the only impact of our work promoting better understanding of the limitations and evaluation of these existing models. We are not enabling any new technologies that do not already exist that could potentially have societal impact.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, is not needed to point out that generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: This paper poses such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (c.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer:
[Yes]
Justification: We discuss these details in the "software libraries" paragraph at the start of Appendix B.
Guidelines:
33

---

The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, if possible, include URL.
The name the license (e.g.. CC-BY 4.0) should be included for each asset. For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: This paper does not release new assets.
Guidelines:
The answer NA means that the paper does release assets.
Researchers should communicate the details the dataset/code/model as part their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
This paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: This paper does not involve crowdsourcing research with human subjects.
34

---

Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with
human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
35


Reviews:
[{"review_id": "UttSH3TCDG", "reviewer": "Reviewer_FWjm", "summary": "I am not a theorist and so did not feel able to provide a detailed review of the theoretical parts of this work. \n\nThe authors discuss the trade-off between context length and mixing time of MDPs, demonstrating that a longer context length leads to longer mixing times and hence greater difficulty in evaluating a given policy. This is an intuitive observation -- if my next action conditions on an observation many timesteps ago, I must wait at least that many timesteps before such a dependency can even be shown, let alone reliably evaluated. However, the authors introduce novel theoretical analysis formally demonstrating this connection. They then perform empirical experiments providing evidence validating their theoretical findings in simple settings and early evidence that transformers have a higher mixing time than other architectures.", "questions": "See weaknesses", "limitations": "The authors discuss the limitations of their theoretical analysis in the paper. It would be nice to flesh out this discussion with more examples of concrete environments or settings where their theoretical analysis is relevant or not.", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* The paper is very well written. Despite my lack of theoretical RL knowledge, I was able to follow the arguments of the paper well.\n* The implications of their argument that there is a link between context length and mixing time are potentially interesting. \n* The evaluation of their theoretical findings is convincing, with Figure 3 nicely demonstrating their central result.", "weaknesses": "* It is not clear to me that the larger mixing times are a significant problem in practice. The authors provide a rough intuition of when high mixing times would be relevant, but this example is not grounded in a real-world benchmark or evaluation. This makes it impossible to assess the practical impact of the authors' theoretical observations for policy evaluation, especially given the toy nature of the environments they evaluate on. \n* I'm not sure i understand the point being made by Figure 9. The authors show different values of context length, $k$, and show that for $k=25$ good performance can be attained, but for lower or higher values of $k$ performance degrades. Is this just a straightforward example of over/underfitting? I'm not sure what this has to do with mixing times? Clearly for most problems $k$ will have an optimal value much higher than the $k$ value of the ensemble of poicies used to generate the data, but I'm not sure I understand what this demonstrates about evaluation.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We wanted to begin by thanking you for your thoughtful review of our paper. We really appreciate your praise of the implications of our theoretical findings, the evaluation of these findings, and the writing quality of our paper. We also found your intuitive explanation of our findings to be quite insightful. We really like the way you put it and will be sure to reference this kind of explanation in the final draft. \n\n**Confusion Regarding Figure 9:** Thank you for pointing out to us how confusing this was in our submission. What wasn’t clearly articulated in our description from lines 316 to 317 was that this experiment was really motivated as a response to a particular potential criticism of our analysis in section 4. One could potentially argue that while a larger context length is needed to fit the set of behavior policies than each had individually, this is simply “overfitting” to the training data and not useful in achieving downstream performance. Our intention in Figure 9 is to provide a concrete counter example to refute this argument. What we believe must have made this even more confusing is that our main point in Figure 9 could only be understood by reading the in-text descriptions of the training accuracy at each context length value. In the final draft, we will update the label to include the training accuracy corresponding to each value of $k$. Figure 9 then chiefly shows that the models with the best generalization have 100% training accuracy. You are also right that another trend highlighted in Figure 9 is that $k$ values that are unnecessarily large seem more prone to overfitting, which isn’t very surprising. We believe these points are interesting to highlight, but also acknowledge that they are only indirectly connected to mixing times through the findings of our paper and simply serve as supporting evidence for the need for growing the context length when building foundation models. \n\n**Practical Impact:** Thank you for mentioning your confusion about the relevance of high mixing times and our analysis in practice. We believe the settings of greatest relevance to our work are those related to continual or multi-task environments where agents are evaluated as generalists over a number of skills rather than just solving a single narrow task. As such, we believe that focus on the difficulties presented by high mixing times is timely in the age of foundation models. As mentioned at the end of section 3.1, our analysis will not have relevance in problems where there are few state variables that each impact every observation. However, composite tasks that test a number of sub-skills naturally tend to have many total state variables with relatively few impacting each observation. So, for example, simple Atari domains such as Pong and Breakout will not suffer from high mixing times, but i.e. continual learning over multiple Atari games including Pong and Breakout will as a result of the sparsity of causal impact of variables across games. Another example would be open world environments that feature a number of natural sub-tasks and sub-regions such as RPGs like the Pokémon series of games or the Legend of Zelda games. Moreover, broadly speaking AI assistant tasks that include providing help on a number of topics rather than just one should also suffer from issues with high mixing times. We will provide more detail along these lines in the final draft. In the submitted draft, we wanted to prioritize focusing on our novel results, but it probably would help readers to contextualize our analysis more deeply with recent prior work on this topic such as [1]. In this paper it was shown why continual learning tasks with multiple subtasks have high mixing times and in-depth experiments on mixing times were conducted for the Atari games. It was found that the common setting in the literature results in practical $\\epsilon$-return mixing times on the order of 10 billion time steps even when only considering 6 Atari games. [1] also argues theoretically how high mixing times are a primary driver of poor performance and training instability on continual learning tasks. In light of these prior results in the literature and the complexity of analyzing large scale and long horizon games, we opted for smaller scale experiments in our paper to allow for greater control and more interpretable insights. \n\n[1] Riemer, M., Raparthy, S.C., Cases, I., Subbaraj, G., Puelma Touzel, M. and Rish, I. “Continual learning in environments with polynomial mixing times.” Advances in Neural Information Processing Systems, 2022.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "SXuOpAx2Ym", "reviewer": "Reviewer_bcCn", "summary": "Authors analyzed the relationship between mixing time (policy evaluation time) with increase in context length for non-markovian, POMDPs. Certainly, increasing context, increases mixing time.", "questions": "See the above", "limitations": "See the above", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 1, "strengths": "Authors propose a tighter upper-bound for mixing times of multi-dimensional mdp with longer contexts", "weaknesses": "Reviewer is not clear why it is of any significance, this seems like a trivial observation that may not even need a theoretical justification. \nReviewer is not sure whether it adds any new knowledge to the field.\n\nAs the context length is increased, we increase dimensionality of the problem, certainly increases the cardinality, so the reviewer doesn't find this result of any significance, unless reviewer is in-correct with his understanding of the draft.\nI don't intend to dis-respect the amount of work put in by the authors, however, I merely wonder how this direct observation needs a (tighter) upper bound, ideally in the current scenario, lower bound is of more importance. Experimental evaluation also doesn't offer any key insights, since context lengths considered are quite small (up to 100)\n\nI have looked at main draft thoroughly and did not look at the appendix.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for taking the time to review our paper. We are sorry to see that you are so negative about the value of our contributions to the community. We believe that there are some key points of confusion that may have led you to this perspective that we will attempt to address below. We will also reference other reviews when helpful to contrast with other perspectives on our paper. \n\n**Our contribution:** After reading through your summary of our contributions, we would like to clarify a potential confusion about the novelty of our results. Our result in Theorem 1 does indeed provide a tighter bound on a prior result, but this is only an intermediate result of our paper necessary for proving Theorem 2, which is our main contribution. Theorem 2 is entirely novel. We are not aware of any paper from the prior literature that has even mentioned the fact that the context length of the interaction history sent to a policy would have an impact on its mixing time. When the reviewer writes “I merely wonder how this direct observation needs a (tighter) upper bound,” the answer is that the tighter bound of Theorem 1 is essential for being able to prove the novel result of Theorem 2. \n\n**Trivial observation:** This comment seems at the heart of your criticisms of our paper. As such, we will attempt to break down your argument piece by piece to provide you with clarification on why this may be much harder to show than you realize. The first aspect of your argument “As the context length is increased, we increase dimensionality of the problem” requires some immediate clarification. Often in RL when we talk about “the dimensionality of the problem,” we are talking about the size of the environment. However, increasing the context length sent to a policy has no direct impact on the environment. Indeed, prior work at this conference [1] has established that as environments are scaled up, their mixing time also increases. We must reiterate though that the environment is not scaled up when the context length is increased, so this result is not at all related to our theoretical contribution. The reviewer also mentions that increasing the context length “increases the cardinality, so the reviewer doesn't find this result of any significance.” The cardinality that is being increased here is the set size of the possible interaction histories $h^{(k)}_t$. But we do not see how the reviewer jumps from the fact that this set size is increasing to the fact that the mixing time is going up. Notice that this $\\epsilon$-mixing time is defined in terms of states and not observations, so it is not obvious that an input that does not explicitly include the state being sent to a policy will impact the $\\epsilon$-mixing time. We know that the reviewer mentioned not looking through the appendix, but the proofs of Theorem 1 and 2 are not direct results of cardinality arguments, rather they have to do with the underlying causal structure of the Markov chain induced by the policy in the environment.  If the reviewer has a better idea of how to show this result in a trivial way that goes beyond a vague intuition, we would kindly request that they spell it out in more detail. As suggested by reviewer hXkG, we will provide proof sketches to aid comprehension in the main text of the final draft.\n\n**Bound Tightness:** As other reviewers also suggested, in the final draft we will add an analysis of bound tightness to section 2.3 for the problems highlighted in Figure 1. In every case $\\beta=0.1$. In the irrelevant variables example, $g=1$ for policies that only condition on $x$, and $g=2$ for policies that condition on both $x$ and $y$. As a result, the bound should be 10 times tighter for the set of policies that only condition on $x$ than the set of policies that condition on both $x$ and $y$. In practice the maximum $\\epsilon$-return mixing time is 7.6 times smaller. In the independent subtasks example, $g=2$ for policies that always condition on $x$, $y$, and $z$. Meanwhile, $g=1$ for policies that focus on the relevant variable $x$ or $y$ depending on the value of $z$. So the bound should once again be 10 times tighter, but the maximum $\\epsilon$-return mixing time is 3.6 times smaller in this case. One reason that we believe these bounds are not as tight as possible for this experiment has to do with using the $\\epsilon$-return mixing time that is not sensitive to state changes that don’t have an impact on the reward. In these examples, we adopted a simple sparse reward where many states get a reward of 0. For the final draft we will also include a variant of the experiment where the reward is set to be different at each state such that the Markov chain over states must fully mix for the $\\epsilon$-return mixing time to also mix. We expect that this setting will verify these bounds to be even significantly tighter than what we have measured here in practice.\n\n**Adding Knowledge to the Field:** We should note that reviewer FWjm also felt that the result in Theorem 2 was intuitive, but thought it was still a contribution to provide the formal analysis and that it was interesting to link context lengths with mixing times. Reviewer LbpG also felt that the contribution was significant. Likewise, reviewer hXkG felt that the contribution had a high degree of significance and that the paper’s biggest strength was its potential influence on subsequent work because its insights are timely. Moreover, reviewer aEyh praised our new angle at an important and interesting problem, remarking that the “community will benefit from this work being available to them.” \n\n**Insights from Evaluation:** Again we should note that reviewer FWjm called our evaluation convincing and reviewer LbpG felt the experiments were comprehensive.\n\n[1] Riemer, M., Raparthy, S.C., Cases, I., Subbaraj, G., Puelma Touzel, M. and Rish, I. “Continual learning in environments with polynomial mixing times.” Advances in Neural Information Processing Systems, 2022.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "FECW9fiINo", "reviewer": "Reviewer_LbpG", "summary": "In the quest for applying RL to more realistic tasks in the long term, the authors interrogate the relationship between context length and mixing times. They provide an extensive theoretical analysis leading to a novel finding regarding the trade-off between learning with large context and slower evaluation. The authors also perform empirical evaluation on a toy-example as well as Minigrid environments to provide evidence of the proposed theories.", "questions": "I am interested in hearing the author's thoughts on what architectural changes could be introduced to tackle the context length and mixing times trade-off", "limitations": "The work discusses the limitations of existing approaches. While the work is essential, it could only be possible with significant compute resources.", "rating": 7, "confidence": 2, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "The work is novel, providing both a theoretical and empirical analysis. The toy example provides evidence for high mixing times and is followed by analysis in larger scale environments: the experimentation is comprehensive. This is a significant contribution and the paper is well-written.", "weaknesses": "While the theoretical contribution is the main contribution, the main paper should ideally include more empirical evidence which has been moved to the supplementary material. It would be ideal if this was extensively explored in a larger journal paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We wanted to begin by thanking you for the kind words in your review regarding a number of aspects of our paper including the novelty, comprehensive empirical analysis, and writing quality. We really appreciate this validation of our work. We also wanted to make sure to engage with all of your concerns and suggestions. \n\n**Experiments in Appendix:** We definitely agree that space limits us in being able to provide our experiments in the main text, but we do feel that it should be possible to provide a concise overview of the key findings in the space limitations of a standard conference paper. This would allow us to reach the biggest audience possible (given that a much smaller community would be interested to read through an extended version). We were wondering if there were any experiments in particular that you felt would make a significant addition to the main text. If so, we would be happy to see if the paper could be restructured to accommodate for this.   \n\n**Possible Architectural Changes:** We really appreciate this question as it cuts to the chase of the core motivations of our work. Most work on RL that even acknowledges the challenges of high mixing times does so with a defeatist mentality, assuming that problems with high mixing times are unavoidably harder and that there is basically nothing that could be done about it. If our work makes any contribution to the community, we hope it will be to highlight that this isn’t actually true and that the policy class we choose to optimize over itself can have a big impact on mixing properties. We hope that readers of our paper come away from it asking exactly the question that you did.  What our paper shows in Theorem 2 is that what leads to potentially high mixing times is when our model leverages a monolithic representation that is highly sensitive to a large part of the interaction history at all times. This is particularly descriptive of how generic transformers work, but there are multiple already existing research directions that seem well suited to scaling to high context lengths while providing less history sensitivity at each step: \n\n- **Hierarchical RL:** In hierarchical RL frameworks such as options [1] it should be possible in domains with temporal coherence to approximate a policy with a longer context length by multiple sub-policies with a smaller context length. This is particularly relevant for approaches to option learning that learn an independent lower dimensional representation space for each option as in [2]. \n\n- **Hybrid Transformer Working Memory Architectures:** Recent work has aimed to improve the effective context length of transformers by augmenting them with some kind of bounded working memory component [3, 4, 5]. While these papers are typically solely motivated by computational efficiency, they effectively serve the role of extending Transformers to longer contexts while making them less sensitive to experiences in this context that are not reflected in the memory. This effectively bring Transformers closer to some of the incremental design patterns of RNNs, which our experiments indicate experience lower mixing times. This makes sense because they must commit to a single representation of the history at each step and do not have the capability to constantly reassess their history representation dynamically as the context changes the way Transformers do. As such, even state-space models [6, 7, 8] or hybrid alternatives could be attractive in achieving the performance of Transformers while potentially alleviating mixing time concerns during learning.  \n\n - **Tracking Policies:** If we aim to learn a non-stationary tracking policy solution concept as argued by [9] we could potentially represent a stationary policy over a longer context length with a non-stationary set of small context length policies. In effect, this becomes similar to what is achieved by the hierarchical RL policies described above. One additional subtlety to highlight in this case is that this solution is also limited by the mixing properties of the tracking policy parameters, so it would also be necessary to tune this approach to move through parameter space as fast as possible.  \n\nWe will make sure to add this additional detail into the discussion section of the final draft where possible.   \n\n[1] Sutton, Richard S., Doina Precup, and Satinder Singh. \"Between MDPs and semi-MDPs: A framework for temporal abstraction in reinforcement learning.\" Artificial intelligence 112.1-2 (1999): 181-211.\n\n[2] Abdulhai, Marwa, et al. \"Context-specific representation abstraction for deep option learning.\" Proceedings of the AAAI Conference on Artificial Intelligence. Vol. 36. No. 6. 2022.\n\n[3] Pham, Kha, et al. \"Generative pseudo-inverse memory.\" International Conference on Learning Representations. 2022.\n\n[4] Das, Payel, et al. \"Larimar: Large Language Models with Episodic Memory Control.\" Forty-first International Conference on Machine Learning. 2024.\n\n[5] Munkhdalai, Tsendsuren, Manaal Faruqui, and Siddharth Gopal. \"Leave no context behind: Efficient infinite context transformers with infini-attention.\" arXiv preprint arXiv:2404.07143 (2024).\n\n[6] Gu, Albert, and Tri Dao. \"Mamba: Linear-time sequence modeling with selective state spaces.\" arXiv preprint arXiv:2312.00752 (2023).\n\n[7] Dao, Tri, and Albert Gu. \"Transformers are SSMs: Generalized models and efficient algorithms through structured state space duality.\" arXiv preprint arXiv:2405.21060 (2024).\n\n[8] Samsami, Mohammad Reza, et al. \"Mastering Memory Tasks with World Models.\" The Twelfth International Conference on Learning Representations. 2024. \n\n[9] Sutton, Richard S., Anna Koop, and David Silver. \"On the role of tracking in stationary environments.\" Proceedings of the 24th international conference on Machine learning. 2007.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "OJg5QKQIE2", "reviewer": "Reviewer_kzFL", "summary": "This work analyzes the mixing time of a policy in average-reward reinforcement learning problems. It shows that the mixing time is related to the structure of the underlying dynamic Bayesian network (DBN), specifically how the state variables are grouped into strongly connected components by the policy. Additionally, the context (trajectory history) that is fed to a policy has an impact on the mixing time: the longer the context length, the harder to mix. The mixing time is empirically verified in several toy problems. The paper also demonstrates that the Transformer architecture is more prone to using longer context and thus harder to mix.", "questions": "Please refer to the weakness above.", "limitations": "The authors adequately addressed the limitations", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- Theoretically analyze the mixing time for the average-reward setting and its relation to the input context length\n- Clear writing and presentation in most places\n- Empirical verification using modern model architectures", "weaknesses": "Multiple places require further clarification.\n\n1. It is challenging to connect Fig.1 to the main point of L167-174 (as well as Thm.1). For example, it is unclear what $\\beta,g$ are in this example and how the bound in Thm.1 is related to the estimated mixing time. This remains unclear even after reading the details in the appendix.\n\n2. L243: It is unclear what’s the difference between $\\Theta$ and $\\Theta’$.\n\n3. To better interpret Thm.2, it would be helpful to point out how $\\beta_k$ and $g_k$ evolves as $k$ increases in the main text. This is discussed in the appendix (L663) but it is important and it would be better to see it in the main text instead.\n\n4. L286: \"It appears that attention mechanisms make it easier to focus on the full context rather than i.e. only the recent parts and that this capability is predictably enhanced when the model capacity is increased.\" It will be helpful to visualize this, e.g., seeing how dense the attention map is w.r.t. different reward rates.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thoughtful review of our paper. We appreciate that you highlighted our theoretical contribution and want to thank you for your praise regarding the writing quality and empirical verification of our theoretical results. Below we have tried to provide clarity regarding each of your questions in the weaknesses section. We hope that you will consider raising your score if we have adequately addressed your concerns. \n\n**Question 1:** Thank you for pointing out your confusion regarding $\\beta$ and $g$ in Figure 1. In the lines between 180 and 213 we went over the main facts needed to derive these terms in combination with the figure caption, but you are definitely right that it is much better to state them clearly as readers will not be as familiar with this kind of analysis and may find it difficult to follow along. In every case $\\beta=0.1$. In the irrelevant variables example, $g=1$ for policies that only condition on $x$, and $g=2$ for policies that condition on both $x$ and $y$. In the independent subtasks example, $g=2$ for policies that always condition on $x$, $y$, and $z$. Meanwhile, $g=1$ for policies that focus on the relevant variable $x$ or $y$ depending on the value of $z$. We will be sure to substantially update this section in the final draft to eliminate any potential confusion. \n\n**Question 2:** Thank you for mentioning that line 243 was confusing. Our motivation for introducing the prime notation was simply to draw a distinction between the parametric class of policies over the state space $\\theta’ \\in \\Theta’$ and policies over interaction history windows $\\theta \\in \\Theta$. The point being that the set of possible policies are not the same or directly comparable due to the different functional form resulting from different input spaces. We will add a footnote to the final draft to clarify this point.  \n\n**Question 3:** Thank you for this insightful comment highlighting this potential area of confusion. We will definitely include this context about how $\\beta_k$ and $g_k$ evolve with $k$ from line 663 of the appendix in the main text to better explain Theorem 2. We should not have assumed readers will figure this out themselves as this provides unneeded cognitive load when interpreting Theorem 2 and we want the results to be easily accessible to as wide as audience as possible. \n\n**Question 4:** Thank you also for you comment about visualizing attention maps to validate our conclusion on line 286. This is a really fabulous suggestion in order to validate whether this intuition we presented actually is consistent with the results. Unfortunately, we did not log data related to attention when we conducted these experiments and doing so will require us to rerun the experiments highlighted in Figure 8 from scratch. This corresponds to 1,000 individual experiments with specific seeds, requiring significant computational resources and time. For the final draft we will rerun these experiments while keeping track of the average attention attributed to each element of the interaction history window across the 8 heads in each layer. It will be interesting to see if the spread of attention is indeed more even on average for Transformer models that experience higher average mixing times. Theorem 2 shows theoretically that what leads to potentially high mixing times is when our model is highly sensitive to a large part of the interaction history at all times. Following from this insight, it seems natural that we may see higher mixing times for Transformers than LSTMs and for bigger Transformers than smaller Transformers. However, we definitely agree that this conclusion will be much stronger if each element of this hypothesis is empirically validated. We really appreciate this suggestion and are excited about being able to provide more clarity about the mechanisms at play in Figure 8 for the final draft. \n\nPlease let us know if we can provide any further clarity regarding any of these points during the discussion period.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "GZ2W4KFJWS", "reviewer": "Reviewer_hXkG", "summary": "This work presents theoretical and empirical evidence that increasing the context length informing a policy also increases the amount of time required to evaluate the capability of the policy accurately (in other words for the distribution over states independent of initial state to stabilize). Theorem bounds the mixing time with an inverse relationship between the minimum probability of reaching the same state configuration from two different states, and the number of strongly connected states: demonstrating that with more states with overlapping paths from the policy, the longer the mixing time. Theorem 2 then extends this by incorporating context length, where increasing context length increases the number of strongly connected components. Finally, some experiments are presented which show a general trend in agreement with the theory.", "questions": "I have asked some questions in the context of the review above. I have no further questions for the moment, but would appreciate if those questions could be answered.\n\nThe main points I need clarified are those raise under the Quality heading in Weaknesses. If these are addressed then I would increase my score to advocate acceptance.", "limitations": "I would highlight the need for proof sketches again here, as well as empirical evaluations of the bounds. Both of which are necessary to establish the limitations and intuition of the theoretical results while reading.", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "# Originality\nWhile this work is based on Kearns and Koller, the authors justify their additions and extend the results far enough to be considered original. In addition the connection to the context length is new and original. The evaluation of foundation models also seems new.\n\n# Quality\nThe work provides clear theorem statements and in one case a formal statement of the assumptions. Experiments seem valid and do depict a discrepancy in mixing time, as predicted by the theory. Overall the aim of the work is clear and the structure of the work makes sense for evaluating or achieving this aim.\n\n# Clarity\nThe work is well written and structured which aids reading. Notation is introduced and defined where necessary\n\n# Significance\nThis work offers a high degree of significance as it demonstrates an important limitation in the evaluation of policies. I believe the potential significance of this work and the subsequent work it may influence is it's biggest strength.  I also agree with the authors that some insight into the reliability of our assessment of foundations models would be useful and timely.", "weaknesses": "# Clarity\nMy biggest concerns of this work are on clarity. It defense of the work - a lot of ground is covered which makes it difficult to structure perfectly. However, in cases throught out terminology is introduced without definition (such as \"aperiodic\" and \"unichain\" in Assumption 1 - while I appreciate these are not new terms and I understand what is being said, it still hinders clarity), some notation and concepts are introduced without a clear purpose (for example why bring up coupling time in the middle of DBNs when it doesn't even appear in Theorem 1? If it is a necessary concept for the theorem then this needs to be clear in the theorem. The link between the causal parent state variables and Theorem 2 is similarly unclear and hinders clarity). Finally, the figures, their captions and corresponding examples are not as helpful as they should be. The examples using Figure 6 in Appendix A make no reference to strongly connected components or any of the more technical details in the Theorems. The example which corresponds to Figure 1 is helpful but the figure itself and the corresponding caption lack detail or any clear direction. This could easily be improved as the figure appears to be making the point that the more coordination between state variables the greater the number of strongly connected components. This was a missed opportunity to clarify one of the more opaque concepts in the mathematics. An example where the paper does provide clear explanations comes on lines 263 to 268. \n\n# Quality\nI have two concerns on quality\nFirstly, the mechanism through which the insight of this work are obtained is not presented. This is due to the absence of proof sketches or any following explanation after the Theorems. A particularly clear example of this is in Theorem 1 where the variable $l$ is not in the bound at all, yes immediately after the statement is made based on the number of strongly connected components (what $l$ quantifies). Another example, is where Lemma 1 is referenced in the Appendix with no explanation or statement of it, when discussing variable subsets. A final weaker example is how $k$ does not appear in Theorem 2 in its own right. While, in this case, it is discussed how $k$ affects $\\beta$ and $g$ it would still be useful to get a concise discussion of the exact mechanism of how $k$ actually leads to a larger mixing time. An example where the paper achieves such a type of statement is again on lines 263 to 268, and specifically \"When observations are only caused\nby a subset of the state variables at each step, there is more potential to break the problem up into independent subtasks...\".\n\nSecondly, there is not a very clear connection between the theory and experiments. For example, Figure 3 has reward rate on the x axis and this seems to have the biggest effect on mixing time (relative to $k$ at least), yet reward rate hardly factors into the theory and is not present in the bounds. The reward rate also seems to affect the results and conclusions as there is not a consistent trend in Figure 3 or 4 across al reward rates. Figure 5 is a highlight of the experiments for me as the distinction between random and learned policies has a clearer connection to strongly connected components and the types of variables present in the theory. So once again the paper does have moments where things are clear and the connections more evident. There just appears to be a need for more consistency. Even on the foundation model results - as necessary as they are - it appears in this case the main point is to just say that foundation models need long enough contexts to disambiguate the various policies they were trained on and that this means they will have a long mixing time. If this is indeed the point, then I think this needs to be clearer as the main point and connection to the theory. Finally, no empirical evidence for the tightness of the bounds is given and typically this is useful and necessary as a source of intuition on how closely the theory fits with what can be expected in practice.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your detailed review of our paper. We very much appreciate your kind words regarding the originality, quality, clarity, and significance of our work. We have attempted to address your key concerns and all points of confusion mentioned in your review below. \n\n**Definitions of aperiodic and unichain:** Thank you for pointing out that this can be a potential source of confusion for readers. We actually did have a relevant part of the appendix where we go through the assumptions and implications of these assumptions for the theoretical findings of our paper at the beginning of Appendix A. For example, we did discuss the implications of the unichain assumption in this section (line 478). However, we can also definitely take this opportunity to provide detailed definitions for all terminology used that we did not have space for in the main text. We will also make sure to direct readers to this section. \n\n**Proof Sketches:** Thank you for your feedback on this point. We didn’t include proof sketches in the submitted draft due to space constraints, but we agree that this can greatly improve comprehension of our results for readers without expecting them to go through our detailed proofs in the appendix. You are right that the reason we mentioned the coupling time was because of its importance in the proof of Theorem 1, but we did not make this link clear enough in the main text. Likewise, the causal parent state variables play a key role in the proof of Theorem 2 and this link was not clear enough in the main text. We believe that providing these sketches is a great idea and will let us better explain to readers why we introduce some important notation, terminology, and assumptions. \n\n**Connections to $\\beta$ and $g$ in Figures 1 and 6:** We really like the idea to provide explicit information on $\\beta$ and $g$ in the description of these figures. See below when we discuss bound tightness an example of these values for Figure 1. In section B.1 of the appendix of the submission, we did make an initial attempt to explain state variable structure of these domains, but it would definitely also be help to take it a step further to make the connection with our theoretical findings more concrete. \n\n**$\\ell$ in Theorem 1:** Thank you for bringing this up. We introduced $\\ell$ chiefly to contrast with the prior result from Kearns and Kholler (line 153). It would indeed be better to introduce $\\ell$ at this point after the theorem statement. \n\n**$k$ in Theorem 2:** Thank you also from bringing this up. On a similar point reviewer kzFL also suggested that it could be useful to provide line 663 from the appendix to make the connection with $k$ clearer in the main text. We will make this clear in the statement of Theorem 2 and subsequent proof sketch in the final draft. \n\n**Reward rate hardly factors into the theory:** In the theory we mostly do not mention reward rates in order to keep our results applicable to general $\\epsilon$-mixing times. However, our experiments are highly dependent on the reward rate because we use the $\\epsilon$-return mixing time in order to not overly focus on spurious state features. $\\epsilon$-return mixing time itself is fundamentally grounded in the reward rate (line 108). \n\n**Consistency of trends in Figures 3 and 4:** There is a difference between the input required to express a policy and the input available to a policy. For example, let’s consider the irrelevant variable example from Figure 1 where it is only necessary for the policy to leverage $x$ and not $y$. Because the optimal policy only depends on $x$, this policy is included within the set of policies that only depend on $x$. However, the optimal policy is also included within the set of policies that take both $x$ and $y$ as input. Even though the functional form of this policy depends on both $x$ and $y$, the actions of this particular policy are only conditionally dependent on $x$. So regardless of the functional form used to search for the optimal policy, the optimal policy is the same in both cases. It receives the same reward rate and encounters the same mixing time. As a result, the optimal reward rate in Figures 3 and 4 is associated with the same optimal policy regardless of the context length and model class. What is more interesting is the mixing times associated with policies experienced along the way for each model class. \n\n**Evidence of bound tightness:** Thank you for this great suggestion. We will add an analysis of bound tightness to section 2.3 for the problems highlighted in Figure 1. In every case $\\beta=0.1$. In the irrelevant variables example, $g=1$ for policies that only condition on $x$, and $g=2$ for policies that condition on both $x$ and $y$. As a result, the bound should be 10 times tighter for the set of policies that only condition on $x$ than the set of policies that condition on both $x$ and $y$. In practice the maximum $\\epsilon$-return mixing time is 7.6 times smaller. In the independent subtasks example, $g=2$ for policies that always condition on $x$, $y$, and $z$. Meanwhile, $g=1$ for policies that focus on the relevant variable $x$ or $y$ depending on the value of $z$. So the bound should once again be 10 times tighter, but the maximum $\\epsilon$-return mixing time is 3.6 times smaller in this case. One reason that we believe these bounds are not as tight as possible for this experiment has to do with using the $\\epsilon$-return mixing time that is not sensitive to state changes that don’t have an impact on the reward. In these examples, we adopted a simple sparse reward where many states get a reward of 0. For the final draft we will also include a variant of the experiment where the reward is set to be different at each state such that the Markov chain over states must fully mix for the $\\epsilon$-return mixing time to also mix. We expect that this setting will verify these bounds to be even significantly tighter than what we have measured here in practice.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "4IVSTfhThd", "reviewer": "Reviewer_aEyh", "summary": "This paper studies the interaction between policies that operate given some context and the mixing time of the policies. It provides a novel connection between the context length and the mixing time, indicating that increased context length might lead to increased mixing times and thus, need for longer, more expensive evaluation. The paper provides a theoretical connection in form of two theorem statements, provides toy examples and ultimately demonstrates some of the findings in an experimental section.", "questions": "Q1: In Figure 3, why does k=10 end at reward rate 0? Does it never achieve higher reward? If so, can you explain why?\n\nQ2: Do you have the fitting curves for the transformer models? Did they train to convergence? How did you tune the hyperparameters? Transformers can be a little tricky to get right and I’m curious what efforts were made to ensure proper function fitting.\n\nQ3: In Figure 4, do you have any results on the state-occupancy distributions of the policies at different reward rates? Or the similarity of these policies? It seems that the random policies which mix well and achieve 0 reward have similar mixing times and the optimal policy which possibly always executes the same short trajectory has similar mixing times. I am currently not 100% convinced that this is a function of context length.\n\nQ4: In Figure 4, do you have an idea why the transformer has much larger mixing times than the MLP? I’m assuming both see the same context length. \n\nQ5: I might be misunderstanding Figure 5, but why would we expect (no matter the context length) a model to be able to distinguish between 1000 different random policies unless it can hold the length of the trajectory in memory?\n\nQ6: It seems to be that there is a difference between the context required to express a policy and the context available to the model. How can this be modeled within the framework? The optimal policy of the MDP described in 6a clearly does not require a very long context. Why would the model with large context not simply learn to ignore any token farther away than like 3 steps (given enough data)?", "limitations": "The limitations of the theoretical assumptions have been discussed. However, I believe a more critical treatment of the experimental results would benefit the community to discern where we can be certain about the correctness of the correlation.", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 4, "contribution": 3, "strengths": "Problem Statement\n* Understanding the complexities of context-based reinforcement learning is an important and interesting problem\n\nClarity\n* The paper was a pleasure to read to due its high quality of writing\n* The paper is very well structured and the story thread is clear, the reader is walked through various sections with simple examples that help foster understanding\n\nRelated Work\n* A good amount of related work is spread throughout the paper and given the high quality of writing I do not believe a separate section for related work is needed.\n\nTheoretical results\n* While the theoretical results are not ground breaking in their complexity, they offer a nice novel spin on some known results and provide a new angle of looking at context lengths which I think is great.\n\nExperiments\n* Simple toy examples verify some of the findings proposed by the theory.", "weaknesses": "Clarity\n* Some of the Figures in the experimental section took me a bit to understand, specifically how the reward rates were correlated with the mixing time in Figures 3 and 4.\nTheory\n* The provided bounds are upper bounds and just because they are slightly tighter for lower context lengths does not mean that any algorithm will follow the trends of the these bounds. It could very well be that any context length complexity is below the complexity of the smallest bound. However, I think this is a good first stab at establishing a theoretical trend that might carry into practice.\n\nExperiments\n* This is I think the weakest part of the submission as I do not find all the experiments conclusive. I think there might be other confounding factors at work that have not been determined yet. Some of my concerns might be alleviated when my open questions. \n  * In both Figure 3 and 4, it seems that policies with high reward rates converge to the same mixing time independent of the context. This questions whether the claim of correlation between the two is correct. \n  * In Figure 4, there seems to be a clear difference between different models but it is unclear to me whether this is a separate issue or whether it can actually be attributed to context length since the MLP seems to have significantly lower mixing time than the Transformer but they use the same context lengths. \n  * In Figure 9, the increased context length might not only raise the mixing time but also introduces a vast amount of additional parameters that need to be fit. This makes it unclear if the mixing time or other potential causes lead to the observed behavior.\n\nOverall, I think this paper provides a novel idea and some new insights into analysing context lengths and even though not all the experiments are conclusive, I am leaning towards acceptance since I believe that the community will benefit from this work being available to them. I believe this is an interesting direction that should be further explored. I am arguing for a weak acceptance for now but am willing to raise my score if my open questions can be addressed.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your comprehensive review of our paper. We really appreciate your kind words about the importance of our problem statement, the clarity of our writing, the quality of our theoretical results, and the value of our work to the research community. We have tried to address each point of confusion that you mentioned below: \n\n**Q6:** While you asked this question last, we wanted to address this one at the onset as it is related to a few other points of confusion that you mentioned. We really like the way that you put “there is a difference between the context required to express a policy and the context available to the model.” We will update Sections 2 and 3 to emphasize this point. For example, let’s consider the irrelevant variable example from Figure 1 where it is only necessary for the policy to leverage $x$ and not $y$. Because the optimal policy only depends on $x$, this policy is included within the set of policies that only depend on $x$. However, the optimal policy is also included within the set of policies that take both $x$ and $y$ as input. Even though the functional form of this policy depends on both $x$ and $y$, the actions of this particular policy are only conditionally dependent on $x$. So regardless of the functional form used to search for the optimal policy, the optimal policy is the same in both cases. It receives the same reward rate and encounters the same mixing time. This is why we focused on $\\epsilon$-optimal policies that depend on both $x$ and $y$, making the point that there are very good policies with higher mixing times than any policy that only depends on $x$. So yes, any policy that receives input that is a sufficient statistic of the environment state and trains for long enough with a principled learning algorithm will arrive at the same optimal policy. The differences is not at this theoretical asymptotic limit, but rather at the points encountered along the way. In real-world problems of sufficient complexity, arriving at the truly optimal policy is generally considered infeasible.  \n\n**Convergence in Figure 3 and 4:** As described above, the optimal reward rate is associated with the same optimal policy regardless of the context length and model class. As noted in the caption of Figure 3, policies with context length $k=1$ do not have sufficient context to achieve the optimal policy. However, every policy class that takes in a greater context length, does eventually achieve the same optimal policy by the end of training. What is more interesting is the mixing times associated with policies experienced along the way for each model class. \n\n**Transformers vs. MLP (Q4):** To clarify all models in Figure 4 use the same context length and Figure 7 of the appendix includes a chart for each of the other context lengths. Keep in mind that we were careful to control for the number of parameters so that the architecture is the main source of variation. The comparison with MLPs is not straightforward and we should note that they do have the highest mixing times relative to Transformers when looked at across different context lengths, experiencing similar mixing times to Transformers i.e. for $k=10$. We hypothesize that these results indicate that MLPs are on average a bit less sensitive to the full context than Transformers when considered over the full course of learning.   \n\n**Number of Parameters:** We are not sure that we were fully able to follow this comment. Figure 9 is from experiments using the Decision Transformer architecture, which does not introduce more parameters when dealing with longer contexts due to parameter sharing within the sequence model. \n\n**Q1:** Thank you for pointing this out. This was a simple artifact of our plotting library which filtered out the points at reward rates of 0.4 and 0.5 for $k=10$ just because there was a bit less data (because of slower learning). When we take out this filter, we predictably see that the mixing time is above the other context lengths at a reward rate of 0.4 and, as expected, converges to the same optimal policy mixing time as the other context lengths. We are really sorry about this mistake and will update the final draft to address this omission.\n\n**Q2:** Yes. Our models train to convergence in all experiments. For our simple RGB world experiments, we did a search over Adam learning rates with gradient clipping until finding one that consistently converged at every context length. For our Decision Transformers experiments, we followed the code from the original paper, which included gradient clipping, Adam optimization, and a learning rate schedule that consists of linear warmup following by cosine learning rate decay. These details were copied from the original paper and we did a grid search over the initial learning rate. We will add these details in Appendix B.\n\n**Q3:** As explained above, the optimal policy is the same regardless of the function class, so convergence of all models around a reward rate of 0.5 is expected. We hypothesize that when multiple model classes have the same reward rate and same mixing time, they are likely learning the same simple policy i.e. always go right regardless of the state. In general, all policies that do not depend on the state of the environment at all will receive a reward of 0 for this problem and would also have the minimal mixing time bound following Theorem 2. We also really like your idea to look at the state occupancy directly and will do a deeper investigation into what is learned by each model at a reward rate of 0 in the appendix of the final draft.  \n\n**Q5:** There is indeed no guarantee that it will always be possible to achieve a training accuracy of 100 percent with limited context. This is actually exactly the point we are trying to make in this figure. Our argument is that models like decision transformers must leverage context lengths far greater than the context lengths used for any behavior policy i.e. $k=1$ in this example.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We wanted to begin by thanking all six reviewers assigned to our paper. We were very happy to see so many positive comments about our work. We were also grateful to see so many constructive comments that will help make our work even better in the final draft. We have provided a detailed rebuttal for each reviewer in order to address all questions and concerns mentioned in their reviews. We hope that if there are still any concerns that reviewers feel were not sufficiently addressed, we will have an opportunity to follow up on these points during the discussion period. 

A benefit of receiving such a large and diverse set of reviews is that we also received multiple follow up questions unique to each reviewer. These questions can generally be addressed with relatively modest adjustments to the presentation of the paper from the submitted draft. However, there were also a few topics that multiple reviewers asked about:

- **Dependence on reward rate in Figures 3 and 4:** In the responses to each reviewer, we tried to provide some context to help them understand why it is expected that the $\epsilon$-return mixing time will depend on the reward rate and why it is expected for mixing times to converge across different models at particular values of the reward rate. We believe that some small adjustments to the presentation of the theoretical results in Sections 2 and 3 will be sufficient to make these points clearer for future readers. 

- **Evidence of bound tightness:** By providing additional background required for analyzing bound tightness for the settings in Figure 1 of Section 2.3, we can provide direct validation of the degree of tightness. Additional experiments would be needed to show if the bound could be empirically verified to be even tighter, but these experiments would be simple variants of the experiments already provided in the submitted draft. As such, we believe it will be straightforward to provide increased clarity for readers on this topic. 

- **Connecting the domains considered to the concepts of $\beta$ and $g$:** The information needed to derive this information was already provided for Figure 1 in Section 2 and we will simply make the direct connection clearer in revisions. We will also make edits to Figure 6 and Section B.1 of the appendix so that this connection is clear in our function approximation experiments as well. We had already explained how these environments could be described in terms of state variables, so providing this analysis is a simple extension of our discussion of these environments in the submitted draft.

- **Clarifying the dependence on $k$ in Theorem 2:** As pointed out by reviewer kzFL, this detail is already included in the appendix and it will be simple to provide improved clarity for readers by presenting this alongside Theorem 2 in the main text.

Meta-review: The paper studies the trade-off between increasing context length in reinforcement learning (RL) models and the resulting increase in mixing time, which affects how quickly the model can be evaluated and trained. The paper presents theoretical insights and empirical evidence showing that although larger context lengths can improve policy performance, they also make reliable evaluation more challenging, especially in large-scale and complex environments.

This paper will provide refreshing insights to some and technical confirmation of a hunch to others. Overall, the reviewers appreciated both the theoretical and the empirical results and their significance. There were concerns whether high mixing times are a significant issue in real-world applications. I find the argument authors gave convincing: “We believe the settings of greatest relevance to our work are those related to continual or multi-task environments where agents are evaluated as generalists over a number of skills rather than just solving a single narrow task. As such, we believe that focus on the difficulties presented by high mixing times is timely in the age of foundation models.” I suggest the authors incorporate these arguments into their work. I also suggest the authors to incorporate the feedback from the reviewers such as adding a proof sketch and visualizing attention maps of transformer.
Decision: Accept (poster)