Balancing Context Length and Mixing Times for Reinforcement Learning at Scale

Authors: matthew  riemer
khimya  khetarpal
mila  janarthanan
sarath  chandar
école  polytechnique de montréal
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 05:50:53

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.

Main paper: 
 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 in comparison to what is achievable with an infinite context length [1] . Meanwhile, there are two obvious downsides to increasing the context length an 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 is 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 a trend towards training on internet-scale datasets [11; 12; 13] that far exceed the quantity of any single human'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 it 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 of TD algorithms also depends on the mixing time. 3Addressing 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 1 we 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 architecture can impact the mixing time and is complementary to recent insights about how changes to 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 a tradeoff to consider between increasing the context length and increasing the mixing time. When we increase the context length sent to an 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 an agent'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:1. We present theoretical analysis with supporting toy examples shedding light on how the input sent to a policy can influence its mixing time in Section 2.2. In Section 3 we formally establish the connection between the context length considered by a 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 a 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 as a result require more interaction for reliable evaluation.To reproduce our experiments see https://github.com/mattriemer/ContextLengthMixing.Question: Do the main claims made in abstract and introduction accurately reflect the paper's contributions and scope? 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 4 respectively).The only other content in the paper is the concluding section, which discusses promising opportunities for future work to build off our ideas. Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 
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.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 an agent's behavior and place no restriction on using discounting within the learning process (as we do in our experiments). 4 We consider a finite, discrete-time, infinite horizon Markov Decision Process (MDP) [20; 21] : M = S, A, T, R , where S is the set of states, A is the set of actions, R : S × A → [0, R max ] is the reward function, and T : S × S × A → [0, 1] is the environment transition probability function. At each time step, an agent perceives a state s ∈ S and takes an action a ∈ A drawn from a policy π : S × A → [0, 1]. The agent then receives a scalar reward drawn from the function R(s, a) and with probability T (s |s, 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 1 All stationary policies are aperiodic and unichain, resulting in a Markov chain with a single recurrent class that is recurrent in the Markov chain of every policy. 5Any RL problem may be modified such that Assumption 1 holds by adding an arbitrarily small positive constant to all transition probabilities in T (s |s, a) [23] . An important corollary is that the steady-state distribution µ π induced by the policy π is independent of the initial state:Corollary 1 All policies π induce a unique steady-state distribution µ π (s) = lim t→∞ P π (s t = s|s 0 ) that is independent of the initial state such that s∈S µ π (s) a∈A π(a|s)T (s |s, a) = µ π (s ) ∀s ∈ S. Corollary 1 implies that long-term rewards and thus the average reward per step objective ρ(π) can be defined independently of the current state [21] :ρ(π) := lim m→∞ 1 m m t=1 E π R(s t , a t ) = lim t→∞ E π R(s t , a t ) = s∈S µ π (s) a∈A π(a|s)R(s, a) .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 |s) = a∈A π(a|s)T (s |s, 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 t π mix ( ) as the -mixing time of the chain induced by π: where d TV is the total variation distance between the two distributions. The conventional mixing time is defined as t π mix ≡ t π mix (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 a mismatch with respect to the reward rate. The -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 s 0 as ρ(π, s 0 , m) := E π [ 1 m m t=1 r t |s 0 ], then we can define the -return mixing time as:t π ret ( ) := min m |ρ(π, s 0 , m ) -ρ(π)| ≤ , ∀s 0 ∈ S and ∀m ≥ 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 a policy at a given threshold of precision .Multidimensional State Space: We will also assume that problems have a state spaces comprised of multiple dimensions. Formally, each state is separated into n ≥ 1 variables i.e., s = [s 1 , ..., s n ] with s i ∈ S i for all variables i ∈ {1, ..., n} such that S ⊆ S 1 × ... × S n . 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, a) = Π i∈{1,..,n} T (s i |s, a). In Appendix A.1 we discuss why this setting is strictly more general than the typical discrete MDP setting and does not limit the scope of our results.Dynamic Bayesian Network (DBN) Structure: A Bayesian Network (BN) is a probabilistic graphical model that represents a set of variables and their conditional dependencies via a directed acyclic graph (DAG) and a dynamic Bayesian network (DBN) is a 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 s 1 , ..., s n that is also influenced by an action variable a dictated by the policy [15] . The dependency graph D π for the DBN is a directed cyclic graph whose nodes are s 1 , ..., s n and where there is then a directed path from s i to s j in D π iff s i t influences s j t for some t > t under policy π. While we can ignore actions as 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 D π should influence itself and thus have a self-loop. If there is a direct path from s i to s j and from s j to s i , the variables i and j are said to be in the same strongly connected component. In contrast, if there is either a directed path from s i to s j or s j to s i , but not both, the variables i 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 Γ π 1 , ..., Γ π 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 Γ π j to Γ π i . We can then define g π = max i |Γ π i | as the maximum number of variables in any strongly connected component of D π and g = max π g π as the maximum possible within a given policy class or parameterization π. In Appendix A.1 we provide examples to help readers better understand how this works and demonstrate that this formulation does not limit the scope of our results.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 τ 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 , let m be the number of steps such that for any two starting states in S if we can say that P (τ > m) ≤ then the Markov chain is -mixed at time m. For this analysis we also must introduce the parameter β i,π , which defines the minimum amount of common probability mass between any two state configurations for any state variable i so that β = min π min i β i,π ∈ [0, 1] further corresponds to the minimum over all state variables i and possible parameterizations of the policy π:β i,π = min s1,s2∈S s i ∈S i min T π (s i |s 1 ), T π (s i |s 2 ) .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 1 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 π (s |s) induced by policy π has maximal strongly connected components in D π with a maximum size of g state variables and a minimum of β common probability mass between any two state configurations for any state variable, the mixing times t π ret ( ) and t π mix ( ) can be upper bounded:t π ret ( ) ∈ t π mix ( ) ∈ O 1 β g log(1/ ) .Proof Sketch: Due to the sorting of the strongly connected components, our analysis is based on coupling each of the Γ π i 's in succession. Because it is possible that multiple Γ π i 's couple at the same step, every step where the Markov chain does not fully couple must be a step where some Γ π i does not couple. Our proof proceeds in the following high-level steps: 1) The probability of Γ π i coupling at a given step once Γ π 1 , ..., Γ π i-1 have all already coupled is ≥ β g . 2) Thus the probability of Γ π i not coupling at a step when Γ π 1 , ..., Γ π i-1 have all already coupled is ≤ (1 -β g ). 3) So the joint probability of Γ π i not coupling for m i ≥ 0 steps when Γ π 1 , ..., Γ π i-1 have all already coupled is ≤ (1 -β g ) mi . 4) If τ > m then i=1 m i = m and the joint probability that m-steps have been spent not coupling in some Γ π i has a probability bound independent of the particular allocation of m into individual m i . Thus we can conclude that P (τ > m) ≤ (1 -β g ) m . 5) Leveraging the identity that 1 -x ≤ e -x for x ≥ 0, we find that P (τ > m) ≤ (e -β g ) m . 6) The Markov chain is -mixed if P (τ > m) ≤ , so it must be -mixed if (e -β g ) m ≤ , which implies that m ≥ 1 β g log(1/ ). 7) Finally, we note the relationship between t π ret ( ) and t π mix ( ) following Lemma 1 of [24] . Bound Tightness: Theorem 1 is a tighter version of Theorem 5.4 presented by Kearns and Koller [15] by a factor of 8 from their result of O 8 β g log(1/ ) . 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 a 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 it 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 at 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 β > 0, in Corollary 2 of Appendix A, we demonstrate the extension of the definition for β to the probability mass in common over c > 1 steps, which must be > 0 for some value of c 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 β and g, which is important when we see significant variation in β i 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 π(•|s), which takes full states s = s 1 , ..., s n ∈ S as input can be associated with β n and g n . Meanwhile, a policy in an alternative class that only considers some potentially time dependent arbitrary subset of n (t) ≤ n state variables as input at any moment in time t such that s 1 , ..., s n (t) ⊆ s 1 , ..., s n can be associated with β n and g n . Lemma 1 in Appendix A then shows that the mixing time bound gets strictly tighter because 1/βg n n ≤ 1/β gn n .We illustrate this phenomenon more concretely through two toy MDP examples detailed in Figure 1 .We will consider the mixing times for the two examples highlighted in Figure 1 over the space of deterministic policies. As in Riemer et al. [14] , we relax the definition of the -return mixing time using an average over start states (rather than a 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 x and y that are both influenced by the agent's actions. As a result, x and y 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 y variable in its decision about actions. If the policy is only conditioned on the variable x, x influences y through its actions, but y no longer has an influence on x. 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 β held constant. Based on Theorem 1, 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 -return mixing time at a reward rate precision of = 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 x has a mixing time above 44.9. Meanwhile, there is a 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 a policy conditioned on both variables with a mixing time as high as 340.3.To test if 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 = x0, but edited to -0.25 when x = x1 and -1.5 when x = x2 rather than 0. This setting was chosen to have no 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 The above figure details the three relevant variables x, y, and z that we will consider for our toy examples. Note the action a0 is interpreted as a1 or vice versa with a 10% failure probability, which is the same as the rate at which z switches regardless of the agent's actions. Irrelevant Variables Example: In this case, we consider only the variables x and y (z is not needed in this example) where the reward is +1 if x = x0 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 x evolves (according to the diagram) with y remaining constant when z = z0 and the variables y evolves with x remaining constant when z = z1. The reward is +1 if z = z0 and x = x0 or if z = z1 and y = y0, or 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 a large part of the change in the state. So, mixing is still substantially faster with respect to the reward (t π ret ( )) than the state (t π mix ( )). MDPs with Independent Subtasks (Figure 1 ): We can now consider an example with three variables x, y, and z where the agent's actions influence both x and y. A policy conditioned on all three variables has a strongly connected component of size one including z (because z is not influenced by the agents actions or x or y) and a strongly connected component of size two including x and y such that g = 2. On the other hand, if the policy were able to condition only on x when z = z0 and only condition on y when z = z1, then the actions that still influence both x 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 β held constant. One for each variable x, 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 on only the relevant variables for the specific subtask even despite the small scale of this toy problem. As for the previous example, we compute the -return mixing time at a reward rate precision of = 0.05, reflecting a minimum of 10% relative estimation error. The optimal policy has a mixing time of 16.0 and only depends on the variable x when z = z0 and variable y when z = z1. In fact, no policy structured like this has a mixing time above 267.1. Meanwhile, there is a policy dependent on all variables that is within the reward rate precision threshold of the optimal policy with a 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.Each state is separated into n variables i.e., s = [s 1 , ..., s n ] with s i ∈ S i for all variables i ∈ {1, ..., n} such that S ⊆ S 1 × ... × S n and T (s |s, a) = Π i∈{1,..,n} T (s i |s, a).• Needed for Theorem 1, Corollary 2, Lemma 1, and Theorem 2.• This assumption is necessary to make it meaningful to talk about the Markov chain induced by a policy as a DBN over state variables and the strongly connected component structure of such a DBN. of size k is also caused by a subset of the state variables such that Par(h (k) ) ⊆ {1, .., n}.i not coupling at a step when Γ π 1 , ..., Γ π i-1 have all already coupled is ≤ (1 -β g ).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.1. 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 no 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 a 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 a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a 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 of 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, a 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 important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
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 a 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.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, T, O . This adds to the definition of an MDP M an observation space O and an observation function O(o|s) that maps states s ∈ S to observations o ∈ O. 6 The state s ∈ S of such a POMDP ultimately serves as a theoretical quantity that is never actually observed from the agent's perspective and the observations o ∈ O that it 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 environment can be called its interaction history h t := {o 0 , a 0 , r 0 , o 1 , a 1 , r 1 , ..., o t }. 7 We can also consider the case where the agent only maintains a finite window of history with size k i.e., h (k) t := {o t-k+1 , a t-k+1 , r t-k+1 , x t-k+2 , a t-k+2 , r t-k+2 , ..., o t } where our notation is chosen so that a memoryless policy that only processes the current observation corresponds to k = 1. We will henceforth call k the context length of the interaction history. A given POMDP P is non-Markovian to order k P which implies that for every combination of state s, action a, and history windowh (k) , T (•|s, a) = T (•|h (k) , a) and R(•|s, a) = R(•|h (k) , a) ∀k ≥ k P .Optimality in POMDPs: An advantage of using a 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 t → ∞, π θ (•|h t ) processes an infinite length sequence h t . In practice, we must limit the context length sent to the policy. The best policy π θ (•|h (k) t ) over a finite context length has the same optimal reward rate as the best policy π θ (•|s t ) over true state inputs when k 8 However, this implies the context length must scale with k P , which could become quite large in highly non-Markovian real-world environments where achieving optimality is difficult.≥ k P such that lim t→∞ max θ∈Θ E π θ (•|h (k) t ) [r t ] = lim t→∞ max θ ∈Θ E π θ (•|st) [r t ].The Effect of Context Length: As described in the previous section, we are interested in multidimensional state spaces consisting of n variables and in particular problems where n is large. A given observation o generated from O(o|s) may be then caused by only a subset of the state variables Par(o) ⊆ {1, .., n} where P ar(o) denotes the causal parent state variables of o. As such, a finite history window h (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, a finite context length is reflective of a maximum n k (t) sized subset of the state variables at a time t where n ≥ n k (t) ≥ n k (t) if context length k ≥ k for all t. This is because P ar(h (k) ) ⊆ P ar(h (k ) ) ⊆ {1, .., n}. Theorem 2 then follows from the combined results of Lemma 1 in Appendix A and Theorem 1.Theorem 2 (Limiting Mixing Times with Context Length): If the Markov chain induced by a policy conditioned on a finite interaction history window π(•|h (k) ) with a context length of k has k maximal strongly connected components in D π with a maximum size of g k variables and a minimum of β k common probability mass between any two state configurations, the mixing times t π ret ( ) and t π mix ( ) can be bounded for any k ≥ k:t π ret ( ) ∈ t π mix ( ) ∈ O 1 β g k k log(1/ ) ∈ O 1 β g k k log(1/ )where the dependence on the context length implies that 0≤ β k ≤ β k ≤ 1 and g k ≥ g k ≥ 1.Proof Sketch: Lemma 1 in the appendix 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 in which k ≥ k for all t to the results of Theorem 1 and Lemma 1: 1) We consider the causal parent state variables of each observation, action, and reward to conclude that P ar(h (k) ) ⊆ P ar(h (k ) ) ⊆ {1, .., n}, which implies that n ≥ n k (t) ≥ n k (t). 2) Through Lemma 1 we show that by rule of Cartesian products over subsets that 0 ≤ β k ≤ β k ≤ 1. 3) Through Lemma 1 we also demonstrate that g k ≥ g k ≥ 1 because causal connections in D π are only added and not removed when the context length is increased. 4) This then implies that 1/βg k k ≥ 1/β g k k ,which is sufficient to prove Theorem 2 using Theorem 1 because is independent of k.! ! " ! ! !"# " !"# … ! $ " $ ! $"# " $"# … Figure 2:POMDPs with Local Observation Structure. An example of an environment with a multidimensional state space where near term observations are only causally dependent on a subset of the variables. At time T >> t we see different dimensions of the state space influencing observations and thus these separate subsets do not need to be in the same strongly connected component.Formalizing the Tradeoff: Theorem 2 formally establishes a 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 k P is high. However, any increase in k leads to a looser mixing time bound.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 environments where k P → 1 or environments with no local structure where all state variables influence all observations. However, if changing the context length k does actually lead to a change in the strongly connected component structure, the impact on the mixing time will be exponential. In Figure 2 we present an illustration of the kinds of problems where the chosen value of k may make a large difference in the value of n k (t). 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 as in our example in the previous section. Context Length. We plot the average mixing time and 95% confidence intervals encountered during 1 billion steps of learning over 100 seeds at each context length k. We bin the average mixing time computation by the nearest 0.1 increment of the reward rate of the policy. k = 1 is not visible because the reward rate is always 0.While Theorem 2 draws a clear connection between context length and mixing times, there still remains a question about if this analysis is too conservative and it is unclear if these policies will actually be encountered during learning. To test this question empirically, we consider the simple RGB world environment in Figure 6a , which is modeled after the classic T-maze [28] environment. We conducted comprehensive experiments leveraging tabular Q-learning 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 important 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 a sufficient statistic of the environment state as input i.e. 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.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 1M 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 a reward rate that others never learn to achieve. In this case, there is not a clear basis for comparison. We also ran experiments evaluating the role of the number of parameters. The extra capacity did not have a statistically significant effect on the mixing times for MLP and LSTM models. Meanwhile, extra capacity resulted in higher encountered mixing times throughout learning and a larger effect when increasing k for Transformer models (see Figure 8 ). 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. In Appendix B we take a 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. [19] is a popular approach that achieved state of the art performance in offline RL by treating learning as a sequence modeling problem akin to language modeling. The objective is to imitate the behavior of all policies π i in the dataset using a window of their interaction history h (k) . See Appendix B for additional details.Fitting Training Data: Foundation models may struggle to model the training data when the context window length k used for training is less than the maximum context window of any policy k D := max i∈{1,...,N } k i . The reason is because our model must not only produce the behavior of each agent π i , but also, must store additional context to disambiguate π i , from all π j ∀j ∈ {1, ..., N } \ 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. 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 D 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 k D = 1. In Figure 5 we plot the training accuracy achieved by a Decision Transformer model as a function of its context length k. We report that the training accuracy is only optimal in general for k ≥ 50, which is significantly higher than k D = 1 with k = 25 reaching optimal performance for some random seeds. We also considered 1,000 episodes generated by different stochastic behavior policies of REINFORCE based learning agents [32] that learn following every episode. Figure 5 demonstrates that it 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 to 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 a particular random policy seed where k = 25 is able to achieve optimal training performance across a variety of return to go prompts stepping by 0.01 from 0 to 1.0 and report the average performance across 1,000 episodes. We plot our results in Figure 9 which validate 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. Moreover, increasing context length k, while important for optimization, raises the mixing time. The average mixing time with precision = 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 are more prone to overfitting than the minimal context length that achieves 100% training accuracy.• Needed for Theorem 2 only because it is the only result in this section concerned with policies over h (k) . • This assumption is not strictly necessary to state in the sense that it also covers the corner case where h (k) depends on all n 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).Multidimensional State Space:The assumption that there are n state variables s = [s 1 , ..., s n ] subsumes the corner case where n = 1 leaving only one variable. Moreover, the assumption that T (s |s, a) = Π i∈{1,..,n} T (s i |s, 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 a ∈ S a and s b ∈ S b do not involve independently of each other, a new variable can be made to represent the combination of these variables s c ∈ S c := S a × S b . 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 1 and Theorem 2. However, as explained in the paragraph on When it Really Matters in Section 3, when n = 1 the bound on the mixing time does not change with k 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 ∈ A is an acyclic directed graph with two layers. The first layer is 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 a given policy can be seen as DBN formed as a 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 Γ π 1 , ..., Γ π is also not restrictive because = 1 is a 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 = 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 in our experiments naturally possess the desired structure to make our bounds in Theorem 1 and Theorem 2 meaningful despite their simplicity. See Figure 6 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 a 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 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, a wall, an empty space or the location of the goal. Each variable can then only be influenced by the variables adjacent to it across a single time step. Thus 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 a limited radius of the agent's location. To reflect all state variables, it is indeed necessary that multiple observations are combined.Theorem 1 (Strongly Connected State Variables Bound): If the Markov chain T π (s |s) induced by policy π has maximal strongly connected components in D π with a maximum size of g state variables and a minimum of β common probability mass between any two state configurations for any state variable, the mixing times t π ret ( ) and t π mix ( ) can be upper bounded:t π ret ( ) ∈ t π mix ( ) ∈ O 1 β g log(1/ ) .Proof Sketch: Due to the sorting of the strongly connected components, our analysis is based on coupling each of the Γ π i 's 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 Γ π i does not couple. Our proof proceeds in the following high-level steps:1. The probability of Γ π i coupling at a given step once Γ π 1 , ..., Γ π i-1 have all already coupled is ≥ β g .i not coupling for m i ≥ 0 steps when Γ π 1 , ..., Γ π i-1 have all already coupled is ≤ (1 -β g ) mi .4. If τ > m then i=1 m i = m and the probability that m-steps have been spent not coupling in some Γ π i has a probability bound independent of the particular allocation of m into m i . Thus we can conclude that P (τ > m) ≤ (1 -β g ) m .5. Leveraging the identity that 1 -x ≤ e -x for x ≥ 0, we find that P (τ > m) ≤ (e -β g ) m .6. The Markov chain is -mixed if P (τ > m) ≤ , so it must be -mixed if (e -β g ) m ≤ , which implies that m ≥ 1 β g log(1/ ).7. Finally, we note that t π ret ( ) ∈ t π mix ( ) following Lemma 1 of Kearns and Singh [24] .Additional Details: Due to the sorting of the strongly connected components, our analysis will be based on stabilizing the Γ π i 's in succession as in Kearns and Koller [15] . For example, if we assume that Γ π 1 , ..., Γ π i-1 have all stabilized by time t, all the variables in Γ π i must then couple at the same time for i to stabilize. This event then happens at time t with probability ≥ β g . 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 ≥ (β g ) . It is not that a strongly connected component i can'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 i was the focus of the successive stabilization for m i ≥ 0 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 i=1 m i = m. This is because every step where a component i did stabilize, either the next component i + 1 did not stabilize at the same step, or we consider a series of components j > i where each stabilizes immediately i.e. m j = 0 until the next component j + 1 does not stabilize. Note that if only i -1 components have stabilized after m steps that implies that m j = 0 for all j > i, so this still describes the setting where not all components actually have the opportunity to stabilize despite m i being defined for all i ∈ {1, ..., }. The probability that the system has not stabilized by m steps P (τ > m) is then equal to the probability across all allocations of m i that each component i did not couple for m i ≥ 0 steps when all components j < i had already stabilized. The probability of i stabilizing at a given step if i -1 also stabilized at or before that step is ≥ β g , so the probability of not stabilizing at a given step is ≤ (1 -β g ), and the probability of not stabilizing for each of m i steps is ≤ (1 -β g ) mi . Then we can consider the joint probability of this across all components to upper bound P (τ > m) noting that the total probability of all possible allocations of m i must be ≤ 1:EQUATIONThe dependence between the strongly connected components in the joint probability is incorporated through interdependence between the m i 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=1 m i = m. In the final inequality of the above equation, we note that 1 -x ≤ e -x for x ≥ 0. The system is -mixed after m steps if P (τ > m) ≤ therefore by the transitive property it is then also -mixed if:EQUATIONWe then proceed by taking the logarithm of both sides of this equation:EQUATIONWe can then multiply each side by negative one, which also reverses the inequality:EQUATIONFinally, we divide both sides by β g and yield the result that it must be -mixed if:EQUATIONThis then implies that:EQUATIONThe final step of the proof is simply to note that the result from Lemma 1 of Kearns and Singh [24] that t π mix ( ) ∈ Ω(t π ret ( )) for any policy π.Corollary 2 (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 D π with a maximum size of g state variables and a minimum of β c common probability mass between any two state configurations for any state variable after c ≥ 1 steps, the mixing times t π ret ( ) and t π mix ( ) can be upper bounded:t π ret ( ) ∈ t π mix ( ) ∈ O c β g c log(1/ )A potential criticism of β from [15] is the potential for β = 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 1 step. For policies following Assumption 1, there must be a common recurrent set 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 β c 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 β c = β i,c ∈ [0, 1]:EQUATIONIt is easy to show this because the probability of a strongly connect component coupling at each step is the same as in the proof of Theorem 1, but we must simply account for the difference in definitions of m = mc and β c = β because the probabilities are a function of c steps now rather than 1 step before. So we can simply plug these substitutions into our equations from Theorem 1 to show that m c ≥ 1β g c log(1/). The result is then a simple multiplication of the bound from Theorem 1 by c given the new definition of β c .Lemma 1 (Conditioning on Subsets of State Variables): Consider the Markov chains induced by two policies drawn from different policy classes. The first policy is π(•|x), which takes n(t) states variables x = s 1 , ..., s n(t) as input at time t, has maximal strongly connected components with a maximum size of g state variables and a minimum of β common probability mass between any two state configurations for any state variable. The second policy π (•|x ), only takes some arbitrary subset of n (t) ≤ n(t) state variables as input at any moment in time t such that x = s 1 , ..., s n (t) ⊆ s 1 , ..., s n(t) has maximal strongly connected components with a maximum size of g state variables and a minimum of β common probability mass between any two state configurations for any state variable. The mixing times in the policy classes of π and π can then be upper bounded such that the bound on π is tighter:t π ret ( ) ∈ t π mix ( ) ∈ O 1 β g log(1/ ) ∈ O 1 β g log(1/ )The policy class over π removes edges in the DBN that were previously present in the class over π and cannot add any. If these edges were within a pre-existing 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 ≥ . If removing edges resulted in splitting the largest strongly connected component, we can say that g < g and g = g otherwise. This implies that g ≥ g .To understand the relation between β and β we must make our notation a bit more detailed than what was presented in the main text. β i can then be stated as follows while making the parameterization of π explicit:EQUATIONLikewise, β i can be defined in a similar fashion:EQUATIONwhere X (t) = S 1 × ... × S n(t) and X (t) = S 1 × ... × S n (t) such that X (t) ⊆ X (t) for all t while x j and x j are analogously obtained from the full state s j for j ∈ {1, 2}. The main difference between these formulations is then minimizing over X (t) × ∆(A) in the case of β i and X (t) × ∆(A) in the case of β i . By rule of Cartesian products over subsets, we then know that X (t) × ∆(A) ⊆ X (t) × ∆(A) and because the minimization over a subset must yield a result larger or the same size as a minimization of the full set, β i ≤ β i for all state variables i. This also implies that β ≤ β as it also must hold for the smallest value over state variables.Bringing it all together, 0 ≤ β ≤ β ≤ 1, and g ≥ g ≥ 1, so 1/β g ≤ 1/β g . So, when combined with the results from Theorem 1, this result has been proven.Theorem 2 (Limiting Mixing Times with Context Length): If the Markov chain induced by a policy conditioned on a finite interaction history window π(•|h (k) ) with a context length of k has k maximal strongly connected components in D π with a maximum size of g k variables and a minimum of β k 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 a 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 a grid search over the initial learning rate.Mixing Time Calculation: Following Riemer et al. [14] , we compute a relaxed version of thereturn mixing time that is averaged over start states to be more reflective of mixing times encountered rather than the worst case. This is 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 x and y and only 3 total states when ignoring y. As a result, there are 2 9 = 512 deterministic policies across both variables and 2 3 = 8 deterministic policies over the relevant variable. 0.474 is the highest reward rate estimate for any policy, so if the reward rate has not mixed to a precision of = 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 x, y, and z and only 6 total states when ignoring the irrelevant variable for the subtask. As a result, there are 2 18 = 262,144 deterministic policies across all variables and 2 6 = 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 a precision of = 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,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.Statistical Significance: The lightly shaded regions around each data point in Figures 4, 5 , and 7 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σ/ √ n where σ is the standard deviation and n is the number of data points at a particular binning of the reward rate.Experimental Protocol: 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 1 billion steps (corresponding to about 7 days for the highest context length k = 10). Meanwhile, our function approximation experiments were run for 6 hours corresponding to about 1 million steps for the Transformer models, 6 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 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] .Neural Network Architecture Details: 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 a linear layer deriving an output for each action's value. ReLU activations are used at each hidden layer. Our LSTM model is sent a sequence of 5 dimensional inputs for each step where 3 dimensions represent the observation, one dimension represents the action, and one dimension represents the reward. The current observation is then sent to the LSTM with an action of -1 and reward of 0 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.Mixing Time Calculation: 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 a 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 Times Across Context Lengths: Our results in the main text only plot the mixing time for the context length k = 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 k = 2), 404 hidden units resulting in 331,686 total parameters (at k = 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,154,252 total parameters (at k = 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 a context length of 1 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,054,722 total parameters, 454 hidden units resulting in 3,308,754 total parameters, and 792 hidden units resulting in 10,055,234 total parameters. We again 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. • Transformer Scaling. For the Transformer model we considered: 88 hidden units resulting in 127,338 total parameters, 136 hidden units resulting in 301,242 total parameters, 256 hidden units resulting in 1,058,562 total parameters, 448 hidden units resulting in 3,228,738 total parameters, and 792 hidden units resulting in 10,067,114 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 of the reward rate as we have primarily done throughout this paper. Because mixing times based on the -return mixing time depend strongly on the reward rate, aggregate metrics 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 k = 2 and k = 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.Transformer Attention Map Analysis: 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 of steps between evaluations. The average attention weights for the 136 hidden unit model were: t -9 : 0.048, t -8 : 0.052, t -7 : 0.061, t -6 : 0.078, t -5 : 0.099, t -4 : 0.123, t -3 : 0.148, t -2 : 0.179, and t -1 : 0.212. In contrast, the average attention weights for the 448 hidden unit model were: t -9 : 0.089, t -8 : 0.080, t -7 : 0.078, t -6 : 0.086, t -5 : 0.102, t -4 : 0.119, t -3 : 0.133, t -2 : 0.149, and t -1 : 0.163. The average 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.Environment Details: Documentation for the SimpleCrossingS9N1-v0 environment that we used for our experiments can be found at https://minigrid.farama.org/environments/minigrid/ CrossingEnv/, which is publicly available following the Apache 2.0 license.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 are 3×7×7 dimensional [30; 31] . There are three convolutional layers with filter sizes of 16, 32, and 64 respectively. 2 × 2 max pooling layers and ReLU activations follow each convolutional layer. The result is 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, 8 attention heads, 5 epochs of training, GeLU nonlinearities within the Transformer, 512 * 20 warm-up tokens, 2 * 500000 * k final tokens, and a dropout rate of 0.1. We also used the same Adam optimization with a learning rate of 0.001, β = (0.9, 0.95), a gradnorm clip range of 1.0, and 0.1 weight decay. It is important to note that a transformation is applied to the rewards before it is sent to the Transformer model referred to as the return to for policy or trial i namely g i,t := Ti τ =t r i,τ which is sent along with the current time step t instead. The advantage of this representation is the ability to prompt the model for high returns to elicit high performing behaviors. We can consider this as a special case of the paradigm in which models take as input some transformation of the interaction history up to a bounded context length.Random Behavior Policies: Our random behavior policies take in the current observation as input and process it with the same convolutional encoder as the Decision Transformer model. This representation is simply processed by a linear layer, mapping it to a value for each action. The policy 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).Learning Behavior Policies: 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 a 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. Meanwhile, 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 γ = 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).Mixing Time Calculation: 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 to 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 = 0.01 precision of the estimated reward rate per step and the final episode not within = 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.In Figure 9 we take a deeper look at the generalization ability of models trained with decision transformers as a 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 k found to be sufficient and because the deterministic behavior 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,000 episodes. In Figure 9 we plot the reward rate of these models ranked according to performance (along the x axis) as well as each behavior policy used to generate the data evaluated in the same way. We plot our results in Figure 9 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 on a dataset of this size.Issues with Discounting: As typically implemented, the popular discounted reward setting of RL does not correspond to the maximization of any objective function over a 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 1 [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 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 to enable more rigorous analysis of agents.Evaluation vs. Optimization: It is important to note that we only use the average reward perspective to evaluate an agent's behavior and place no 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 to 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 to 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 a corner case subsumed within the formulation of Theorem 2. This is because the discount factor or time-step number can be considered to be a 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 in Theorem 2. In fact, it helps ensure that the result is as general as possible.Infinite vs. Finite Horizons: Bojun [43] proves that every finite horizon task has a unique steadystate distribution under any policy. This means that Assumption 1 and Corollary 1 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 a reliable measure of performance. Indeed, we consider episodic tasks within our experiments using the Minigrid Crossing environment in Section 4.Theorem 2 shows that what leads to potentially high mixing times is when models leverage 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, 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 [47; 48; 49; 50] , it should be possible in domains 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 [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 to longer contexts while making them less sensitive to experiences in this context that are not reflected in the memory. This effectively brings Transformers closer to some of the incremental design patterns of RNNs, 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 is typically addressed with retrieval augmented approaches that only focus on a subset of the interaction history at a time [62; 63; 64] or a 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 a longer context length with a nonstationary 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.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 to which these learning benefits can be combined with the benefits of a particular choice of policy class aimed at lowering mixing times with a smaller context length. Moreover, considering the importance of our analysis to 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] .Question: For each theoretical result, does the paper provide the set of assumptions and a complete (and correct) proof?Answer: [Yes] Justification: We provide discussion of all assumptions in the main text, which we recap for clarity at 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 its own subsection building off these assumptions.• 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 crossreferenced. • All assumptions should be clearly stated or referenced in the statement of any theorems.• The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a 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. 
experimental result reproducibility: 
 Question: Does the paper fully disclose all the information needed to reproduce the main experimental 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)?Justification: To augment the presentation in the main text, all remaining details needed to reproduce our experiments are provided in Appendix B.• The answer NA means that the paper does not include experiments.• If the paper includes experiments, a 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 a 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 a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it 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 one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed. • While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. Guidelines:• The answer NA means that the paper does include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence 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).• It should be clear whether the error bar is the standard deviation or the standard error of the mean. • It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 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. 
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 defeatist 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 a 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 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 will not suffer from high mixing times, but i.e. continual learning over multiple Atari games will [14] as a result of the sparsity of causal impact of variables across games. 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.Evaluation of Foundation Models: 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 is very important for researchers to consider as these models are increasingly being deployed in the real-world.Appendix Overview: In this Appendix we provide important details that there was not enough space to include in the main text. First in 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 C 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 1 and Theorem 2 in the main text. We also highlight Corollary 2, which is a key extension of Theorem 1 that relaxes the definition of β to be over c-steps rather than a single step. Moreover, we highlight Lemma 1, which is a key intermediate result to bridge the gap between the results in Theorem 1 and Theorem 2. For each theoretical result, we first state the result and then provide a detailed proof. Our proofs build off the definitions and notation established in the main text.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 to a 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 our definitions of the mixing time, the average reward, and coupling time in 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). 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?Answer: [Yes] Justification: We provide this information in the "compute infrastructure" paragraph at 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 (e.g., preliminary or failed experiments that didn't make it into the paper). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines?Justification: The authors have read the code of ethics and made every effort to 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 circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [NA] Justification: Our paper contributes fundamental research on existing models, so only impact of our work is promoting a 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 not tied to particular applications, let alone deployments. However, if there is a 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, it is not needed to point out that a 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 as 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. • If 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). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)? Answer: [NA] Justification: This paper poses no such risks. Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a 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 to 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 a best faith effort. 12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected? Answer: [Yes] Justification: We discuss these details in the "software libraries" paragraph at the start of Appendix B. Guidelines:Guidelines:• The answer NA means that the paper does not involve 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 we 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 as the institution conducting the review.See Riemer et al.[14] for a discussion of the difficulties presented for learning when mixing times are high.See Appendix C for additional explanation about why this is the 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 a more detailed discussion of this assumption and its implications.Note that any non-stationary MDP can alternatively be viewed as a POMDP (see Proposition 2 of Khetarpal et al.[27]) and by the same logic any non-stationary POMDP can alternatively be viewed as a stationary POMDP, so there is no loss of generality in making the assumption that the POMDP is stationary.Many exclude reward from the history representation in POMDPs, but we include it to be as general as possible. That said, our theoretical results do not require this particular choice of the history representation.Note the distinction between Θ , the possible policies over st, and Θ, the possible policies over h (k) t . 
 
 Appendix:  


Decision: 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.

Reviews: review: paper_summary: paper_summary: I am not a theorist and so did not feel able to provide a detailed review of the theoretical parts of this work. 

The 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.
strength_weakness: strength_weakness: * 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. 
* 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.
questions: questions: See weaknesses
limitations: 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.
paper_summary: paper_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.
strength_weakness: strength_weakness: Reviewer is not clear why it is of any significance, this seems like a trivial observation that may not even need a theoretical justification. 
Reviewer is not sure whether it adds any new knowledge to the field.

As 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.
I 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)

I have looked at main draft thoroughly and did not look at the appendix.
questions: questions: See the above
limitations: limitations: See the above
paper_summary: paper_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.
strength_weakness: strength_weakness: 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.
questions: 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: limitations: The work discusses the limitations of existing approaches. While the work is essential, it could only be possible with significant compute resources.
paper_summary: paper_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.
strength_weakness: strength_weakness: Multiple places require further clarification.

1. 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.

2. L243: It is unclear what’s the difference between $\Theta$ and $\Theta’$.

3. 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.

4. 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.
questions: questions: Please refer to the weakness above.
limitations: limitations: The authors adequately addressed the limitations
paper_summary: paper_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.
strength_weakness: strength_weakness: # Clarity
My 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. 

# Quality
I have two concerns on quality
Firstly, 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
by a subset of the state variables at each step, there is more potential to break the problem up into independent subtasks...".

Secondly, 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.
questions: 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.

The 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: 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.
paper_summary: paper_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.
strength_weakness: strength_weakness: Clarity
* 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.
Theory
* 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.

Experiments
* 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. 
  * 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. 
  * 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. 
  * 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.

Overall, 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.
questions: 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?

Q2: 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.

Q3: 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.

Q4: 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. 

Q5: 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?

Q6: 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: 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.

Comments: title: Re: Response to Rebuttal by Reviewer hXkG (Part 2)
Re: Response to Rebuttal by Reviewer hXkG (Part 1)
Response
Re: A quick question for authors
Response to Rebuttal by Reviewer hXkG
Response noted
A quick question for authors
comment: **Consistency of trends in Figures 3 and 4:** Thank you for pointing out that this is still confusing. Our overall point is that at some reward rates (i.e. the optimal one) there is only one potential policy that achieves that reward rate that can be learned. When the policy of two different architectures or context lengths learn the exact same behavior, both their reward rates and mixing times should be expected to be the same. What is more interesting is analyzing the difference between the mixing times when the policies for two different architectures or context lengths arrive at a similar reward rate while still displaying different behavior (i.e. the policies learned are not exactly the same). We would argue that we don’t actually see instances of findings flipping in our experiments. Rather, we see two trend that may appear as inconsistency, but really are perfectly consistent with the theory of our paper: 

1.	At some reward rates multiple context lengths or architectures learn the exact same policy and estimated mixing times are the same or their differences are not statistically significant. 
2.	Sometimes a context length or architecture will learn to achieve a value of the reward rate that other policies and architectures never learn to achieve. In this case there is not a clear basis for comparison at these values. 

Please also let us know whether this explanation helps clarify your confusion on the trends in Figures 3 and 4. While at some points in the figures multiple architectures or context lengths arrive at the same solution or are not directly comparable, when they are comparable and different, we see consistent orderings between the different policy classes considered. 

**Additional Experiments on Bound Tightness:** Because of your interest in this experiment, we thought you may like to hear about our intermediate progress. We have tried a setting of the irrelevant variables example where reward is still $1.0$ when $x = x0$, but edited to $-0.25$ when $x = x1$ and $-1.5$ when $x = x2$ rather than $0$. This setting was chosen to have no 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 when only focusing on the relevant variable, rather 7.6 times smaller with the previous sparse rewards. Note that our reward still has no dependence on the irrelevant variable $y$ and is thus invariant to a large part of the change in the state space. We will also try changing the reward based on $y$ in a way that does not impact preferences over policies in order to close the gap further between $t_{ret}^\pi(\epsilon)$ and the mixing time over the full state space $t^\pi_{mix}(\epsilon)$. We believe this is likely to push our results even closer to the theoretical bound of a 10 times ratio for this problem. We definitely appreciate why you are more interested in the independent subtasks example, but we were unfortunately not able to get new results in time for the close of the discussion phase. As discussed in Appendix B.1, the search over the policy space is very expensive for that problem, so it was not feasible to complete in time. We certainly understand if you feel that these results are too incomplete to impact your review, but still wanted to speak to our progress just to give you a better overall picture of where things stand.
Thank you for providing such a thorough response to our rebuttal. We really appreciate your willingness to work with us constructively to improve our paper during the discussion period. As you suggested, we have provided first draft proof sketches for Theorems 1 and 2 using the notation we have established in the main text. 

**Proof Sketch for Theorem 1:** Due to the sorting of the $\ell$ strongly connected components, our analysis is based on coupling each of the $\Gamma^\pi_i$'s in succession. Because it is possible that multiple $\Gamma^\pi_{i}$’s couple at the same step, every step where the Markov chain does not fully couple must be a step where some $\Gamma^\pi_{i}$ does not couple. Our proof proceeds in the following high-level steps:

1.	The probability of $\Gamma^\pi_{i}$ coupling at a given step once $\Gamma^\pi_1, ..., \Gamma^\pi_{i-1}$ have all already coupled is $\geq \beta^g$.

2. Thus the probability of $\Gamma^\pi_{i}$ not coupling at a step when $\Gamma^\pi_1, ..., \Gamma^\pi_{i-1}$ have all already coupled is $\leq (1 - \beta^g)$. 

3.	So the joint probability of $\Gamma^\pi_{i}$ not coupling for $m_i \geq 0$ steps when $\Gamma^\pi_1, ..., \Gamma^\pi_{i-1}$ have all already coupled is $\leq (1-\beta^g)^{m_i}$. 

4.	If $\tau > m$, then $\sum_{i=1}^\ell m_i = m$ and the joint probability that $m$-steps have been spent not coupling in some $\Gamma^\pi_{i}$ has a probability bound independent of the particular allocation of $m$ into individual $m_i$. Thus we can conclude that $P(\tau > m) \leq (1-\beta^g)^m$. 

5.	Leveraging the identity that $1-x \leq e^{-x}$ for $x \geq 0$, we find that $P(\tau > m) \leq (e^{-\beta^g})^m$. 

6.	The Markov chain is $\epsilon$-mixed if $P(\tau > m) \leq \epsilon$, so it must be $\epsilon$-mixed if $(e^{-\beta^g})^m \leq \epsilon$, which implies that it is  is $\epsilon$-mixed if $m \geq \frac{1}{\beta^{g}} log(1/\epsilon)$. 

7.	Finally, we note the relationship between $t_{ret}^\pi(\epsilon)$ and $t^\pi_{mix}(\epsilon)$ following Lemma 1 of [1].   

**Proof Sketch for Theorem 2:** Lemma 1 in the appendix considers the mixing time relationship of policy classes conditioned on subsets of the state variables that other policy classes are conditioned on. Our proof proceeds by applying our notation from Section 3 in which $k' \geq k$ for all $t$ to the results of Theorem 1 and Lemma 1:

1. We consider the causal parent state variables of each observation, action, and reward to conclude that 
$Par(h^{(k)}) \subseteq Par(h^{(k')})$, which implies that $n \geq n_{k'}(t) \geq n_k(t)$.

2.	Through Lemma 1 we show that by rule of Cartesian products over subsets $0 \leq \beta_{k'} \leq \beta_k \leq 1$. 

3.	Through Lemma 1 we also demonstrate that $g_{k'} \geq g_{k} \geq 1$ because causal connections in $\mathcal{D}^\pi$ are only added and not removed when the context length is increased.   

4.	This then implies that $1/\beta_{k'}^{g_{k'}} \geq 1/\beta_{k}^{g_{k}}$, which is sufficient to prove Theorem 2 using Theorem 1 because $\epsilon$ is independent of $k$. 

Definitely let us know whether these proof sketches helped provide clarity for you. We are happy to revise them based on your feedback as this is just an initial draft.

[1] Michael Kearns and Satinder Singh. Near-optimal reinforcement learning in polynomial time. *Machine learning,* 49(2):209–232, 2002.
Dear authors, 

thank you for responding to my questions so thoroughly. I also read the reviews by other reviewers. I believe that I have a better understanding of some of the points now. I do maintain that I think this paper introduces novel ideas. I also maintain that the relation between experiments and theory could be stronger as others have pointed out but I believe that this should at this point not be reason for rejection. I think this is a good starting point and I think we should encourage more work in this direction. Thus, I'm raising my score to 6.
Thank you for participating in the discussion period and for your constructive feedback that we believe has helped improve our paper. We were wondering if you could clarify further about what concerns regarding the presentation and experiments of our paper still remain after our response to your review. We just ask because many of your comments, especially regarding the presentation, seemed straightforward for us to address. We want to make sure that we understand how our paper can be improved to address your recent feedback.
I thank the authors for the clarifications. The paper can help reveal some potential issues theoretically, but the presentation and experiments can be better so I would like to maintain my score.
Thanks very much to the authors for their rebuttal and clarifying comments. I will maintain my score for now.
Thank you for your question. Intuitively, what we meant was $\rho(\pi, s_0, m) := E_\pi[\frac{1}{m} \sum_{t=1}^m r_t | s_0] $. Maybe it would have been clearer if we used the term reward rather than return?

The $\epsilon$-return mixing time was first proposed in [1] with the authors defining an analogous term to $\rho(\pi,s_0,m)$ in their different notation using definition 4 as well as equations 3 and 5 to accomplish this (see https://www.cis.upenn.edu/~mkearns/papers/KearnsSinghE3.pdf). Their definition is more formal than what we provided above. They begin by defining what they call a $T$ path in the MDP i.e. a particular $T$-step path within the MDP. $T$ here is analogous to what we call $m$ in our paper. Then they define the average undiscounted return per step along that $T$ path in equation 3. This allows them to define a term analogous to $\rho(\pi,s_0,m)$ in equation 5, which is the expected average undiscounted return per step where the expectation is taken over the likelihood that a particular $T$ path is executed by the policy $\pi$ from the starting state $s_0$. 

We appreciate you raising the potential confusion regarding $\rho(\pi, s_0, m)$ to our attention and will definitely include a definition in the final draft of our paper. We were hoping to ask your opinion on which of these two ways of writing it you feel would be most appropriate. The approach from [1] is definitely more formal and descriptive of how it is measured in practice. On the other hand, we also worry about bogging the readers down in notation like the notion of paths through the MDP, which could distract readers from our main contributions. 

[1] Michael Kearns and Satinder Singh. Near-optimal reinforcement learning in polynomial time. *Machine learning,* 49(2):209–232, 2002.
I thank the authors for their thorough response. My following thoughts are as follows:\
**Definitions of aperiodic and unichain**: Yes, explicitly pointing to the appendix here is necessary and will address my concern.\
**Proof Sketches**: Would the authors be able to provide these before the end of the discussion period? I appreciate that I have responded a couple of days into the discussion and so if these are still draft version I will be happy. It just seems like accepting this change without review is unwise.\
**Connections to $\beta$ and $g$ in Figures 1 and 6, $l$  in Theorem 1, $k$  in Theorem 2**: Noted with thanks.\
**Reward rate hardly factors into the theory**: I understand better now. Thank you. I do think this is a general weakness of the work as it stands as clearly in the experiments reward rate is having a material effect. Considering and experimental design or metric which does not have this issue or incorporating reward rate into the theory would help. However, I do not think this is grounds for rejection, subject to a clearer mention of the link in a revised draft. I would just put this under the "why not higher" part of my final score rather than a "reject vs accept" point.\
**Consistency of trends in Figures 3 and 4**: I appreciate the point the author is making. Correct me if I'm wrong then, but I think the point is that the nuisance input variable can change the mixing time but not optimal policy. If so, I understand this, but I'm still not clear how this addresses the fact that the reward rate can change the relative mixing time between context lengths (Fig 3) or architectures (Fig 4). For Fig 3 the switch between $k=4$ and $k=2$ is right at the end and within the variance bounds, so I could drop this point there. But I'm still unsure why reward rate can flip the findings.\
**Evidence of bound tightness**: This is an important addition. As can be see from those results, qualitatively the results hold but quantitatively there can be a big difference. The $3.6$ shrink rather than the expected $10$ is a bit of a concession. But I think the qualitative results are still meaningful and significant, so my main issue here is that it need to be added. I trust it will be in the subsequent draft. I notice that the authors did not use the figures page. If they are still able to and can provide results for the proposed experiment: "include a variant of the experiment where the reward is set to be different at each state" that would be helpful. However, without reviewing these results it is difficult to let it influence my assessment.

For the moment, I feel I am understanding the work better and will raise my confidence score. My rating will remain unchanged but if the authors provide the draft **proof sketches** and a bit more clarity on **Consistency of trends in Figures 3 and 4** I *will* increase my score to advocate acceptance.
Thank you for engaging with the review. I believe adding the discussion on architectural changes to the paper is essential. I think that space could be better used through subfigures so that Figure 5 and Figure 8 could be placed side by side, for example. I will stick to my original score.
Can you please define $\rho(\pi, s, m)$ mathematically in terms of rewards? Thank you.