Original: You Can’t Count on Luck: Why Decision Transformers Fail in Stochastic Environments
Revised: You Can’t Count on Luck: Why Decision Transformers Fail in Stochastic Environments

Original: Abstract
Revised: Abstract

Original: Recently, methods such as Decision Transformer [1] that reduce reinforcementlearning to a prediction task and solve it via supervised learning (RvS) [2] havebecome popular due to their simplicity, robustness to hyperparameters, and strongoverall performance on offline RL tasks.
Revised: Recently, methods such as Decision Transformer [1] that reduce reinforcement learning to a prediction task and solve it via supervised learning (RvS) [2] have become popular due to their simplicity, robustness to hyperparameters, and strong overall performance on offline RL tasks.

Original: However, simply conditioning a probabilistic model on a desired return and taking the predicted action can fail dramaticallyin stochastic environments since trajectories that result in a return may have onlyachieved that return due to luck.
Revised: However, simply conditioning a probabilistic model on a desired return and taking the predicted action can fail dramatically in stochastic environments since trajectories that result in a return may have only achieved that return due to luck.

Original: In this work, we describe the limitations of RvSapproaches in stochastic environments and propose a solution.
Revised: In this work, we describe the limitations of RvS approaches in stochastic environments and propose a solution.

Original: Rather than simplyconditioning on returns, as is standard practice, our proposed method, ESPER,conditions onlearned average returns which are independent from environmentstochasticity.
Revised: Rather than simply conditioning on the return of a single trajectory as is standard practice, our proposed method, ESPER, learns to cluster trajectories and conditions on average cluster returns, which are independent from environment stochasticity.

Original: Doing so allows ESPER to achieve strong alignment between targetreturn and expected performance in real environments.
Revised: Doing so allows ESPER to achieve strong alignment between target return and expected performance in real environments.

Original: We demonstrate this inseveral challenging stochastic offline-RL tasks including the challenging puzzlegame 2048, and Connect Four playing against a stochastic opponent.
Revised: We demonstrate this in several challenging stochastic offline-RL tasks including the challenging puzzle game 2048, and Connect Four playing against a stochastic opponent.

Original: In all testeddomains, ESPER achieves significantly better alignment between the target returnand achieved return than simply conditioning on returns.
Revised: In all tested domains, ESPER achieves significantly better alignment between the target return and achieved return than simply conditioning on returns.

Original: ESPER also achieveshigher maximum performance than eventhe value-based baselines.
Revised: ESPER also achieves higher maximum performance than even value-based baselines.

Original:  Introduction
Revised:  Introduction

Original: Offline reinforcement learning (RL) [3, 1, 4, 5] is a promising approach to train agents withoutrequiring online experience in an environment, which is desirable when online experience is expensiveor when offline experience is abundant.
Revised: Offline reinforcement learning (RL) [3, 1, 4, 5] is a promising approach to train agents without requiring online experience in an environment, which is desirable when online experience is expensive or when offline experience is abundant.

Original: A recent trend in offline RL has been to use simple approachesthat do RL via Supervised Learning (RvS) (e g., [1, 2, 6–8]) rather than typical value-based approaches.
Revised: A recent trend in offline RL has been to use simple approaches that do RL via Supervised Learning (RvS) (e g., [1, 2, 6–8]) rather than typical value-based approaches.

Original: RvS algorithms such as Decision Transformer [1] train a model to predict an action based on thecurrent state and an outcome such as a desired future return.
Revised: RvS algorithms such as Decision Transformer [1] train a model to predict an action based on the current state and an outcome such as a desired future return.

Original: These agents ask the question “if Iassume the desired outcome will happen, in my experience what action do I typically take next.”
Revised: These agents ask the question “if I assume the desired outcome will happen, in my experience what action do I typically take next.”

Original: Thesemethods are popular due to their simplicity, strong performance on offline benchmark tasks [1, 2],and similarity to large generative models (e.g., [9–12]) that continue to show stronger performanceon more tasks when training larger models on more data.
Revised: These methods are popular due to their simplicity, strong performance on offline benchmark tasks [1, 2], and similarity to large generative models (e.g., [9–12]) that continue to show stronger performance on more tasks when training larger models on more data.

Original: However, as we will show, methods that condition on outcomes such as return can make incorrectdecisions in stochastic environments regardless of scale or the amount of data they are trainedon.
Revised: However, as we will show, methods that condition on outcomes such as return can make incorrect decisions in stochastic environments regardless of scale or the amount of data they are trained on.

Original: This is because implicitly these methods assume that actions that end up achieving a particulargoal are optimal for achieving that goal.
Revised: This is because implicitly these methods assume that actions that end up achieving a particular goal are optimal for achieving that goal.

Original: This assumption is not true in stochastic environments,where it is possible that the actions taken in the trajectory were actually sub-optimal and that theoutcome was only achieved due to lucky environment transitions.
Revised: This assumption is not true in stochastic environments, where it is possible that the actions taken in the trajectory were actually sub-optimal and that the s). outcome was only achieved due to lucky environment transitions.

Original: For example, consider the gamblingenvironment in Figure 1.
Revised: For example, consider the gambling environment in Figure 1.

Original: Though there may be many episodes in which an agent gets a positive 
Revised: Though there may be many episodes in which an agent gets a positive return from gambling ( a 0 or a 1 ), gambling is sub-optimal since it results in a negative return in expectation while a 2 always results in a positive return. Since RvS takes all of these trajectories as expert examples of how to achieve the goal, RvS will act sub-optimally.

Original: [Figure] This limitation of RvS approaches in stochastic environments is well-hidden by the majority ofbenchmark tasks for offline RL (e.g., [13, 14]), which tend to be deterministic or near-deterministic.
Revised: This limitation of RvS approaches in stochastic environments is well-hidden by the majority of benchmark tasks for offline RL (e.g., [13, 14]), which tend to be deterministic or near-deterministic.

Original: Locomotion tasks in MuJoCo [15] and Atari games in the Arcade Learning Environment [16] are twosuch examples.
Revised: Locomotion tasks in MuJoCo [15] and Atari games in the Arcade Learning Environment [16] are two such examples.

Original: While deterministic tasks can be solved by replaying promising action sequences,stochastic tasks are significantly less trivial to solve and often more realistic, requiring reactivepolicies to be learned [17].
Revised: While deterministic tasks can be solved by replaying promising action sequences, stochastic tasks are significantly less trivial to solve and often more realistic, requiring reactive policies to be learned [17].

Original: Many real-world tasks are stochastic, either inherently or due to partialobservability, such as having a conversation, driving a car, or navigating an unknown environment.
Revised: Many real-world tasks are stochastic, either inherently or due to partial observability, such as having a conversation, driving a car, or navigating an unknown environment.

Original: Intheir current form, approaches such as Decision Transformer [1] are likely to behave in unexpectedways in such scenarios.
Revised: In their current form, approaches such as Decision Transformer [1] are likely to behave in unexpected ways in such scenarios.

Original: One way to view why RvS doesn’t work in the gambling environment is that when conditioning ontrajectories that achieve a positive reward, the model doesn’t get to see any of the trajectories wherethe same sequence of actions leads to a negative reward.
Revised: One way to view why RvS doesn’t work in the gambling environment is that when conditioning on trajectories that achieve a positive reward, the model doesn’t get to see any of the trajectories where the same sequence of actions leads to a negative reward.

Original: Due to these unrealistic dynamics, there isno policy that would generate this set of trajectories in the real environment, so it doesn’t make senseto treat them as expert trajectories.
Revised: Due to these unrealistic dynamics, there is no policy that would generate this set of trajectories in the real environment, so it doesn’t make sense to treat them as expert trajectories.

Original: Our insight is that there are certain functions of the trajectoryother than return that, when conditioned on, will better preserve the dynamics of the environment.
Revised: Our insight is that there are certain functions of the trajectory other than return that, when conditioned on, will better preserve the dynamics of the environment.

Original: Our approach, called ESPER, realizes this by conditioning on outcomes that are fully determinedby the actions of the agent and independent of the uncontrollable stochasticity of the environment .
Revised: Our approach, called ESPER, realizes this by conditioning on outcomes that are fully determined by the actions of the agent and independent of the uncontrollable stochasticity of the environment .

Original: While trajectory return is not such an outcome, we show that the expected return of behavior shownin a trajectory is, and how to learn such a value.
Revised: While trajectory return is not such an outcome, we show that the expected return of behavior shown in a trajectory is, and how to learn such a value.

Original: The contributions of this work are as follows:
Revised: The contributions of this work are as follows:

Original: • We show that RvS-style agents can only reliably achieve outcomes that are independent from environment stochasticity and only depend on information that is under the agent’s control.
Revised: • We show that RvS-style agents can only reliably achieve outcomes that are independent from environment stochasticity and only depend on information that is under the agent’s control.

Original: • We propose a method to learn environment-stochasticity-independent representations of trajectories using adversarial learning , on top of which we label each trajectory with theaverage return for trajectories with this representation.
Revised: • We propose a method to learn environment-stochasticity-independent representations of trajectories using adversarial learning , on top of which we label each trajectory with the average return for trajectories with this representation.

Original: • We introduce several stochastic offline RL benchmark tasks and show that while RvS agents consistently underperform on the conditioned returns, our approach (ESPER) achieves significantly stronger alignment between target return and actual expected 
Revised: • We introduce several stochastic offline RL benchmark tasks and show that while RvS agents consistently underperform on the conditioned returns, our approach (ESPER) achieves significantly stronger alignment between target return and actual expected performance.

Original: performance. ESPER gets state-of-the-art performance on these tasks, solving all tasks 
Revised: ESPER gets state-of-the-art performance on these tasks, solving all tasks with near-maximum performance and surpassing the performance of even strong value-based methods such as CQL [4].

Original:  Approach2.1 Problem Setup
Revised:  Approach 2.1 Problem Setup

Original: We model the environment as an MDP, defined as the tuple ( S, A, T, R, γ ) .
Revised: We model the environment as an MDP, defined as the tuple ( S, A, T, R, γ ) .

Original: S is a set of states; Ais a set of actions; the transition probabilities T : S ×
Revised: S is a set of states; A is a set of actions; the transition probabilities T : S ×

Original: A × S → [0 , 1] defines the probability of73 the environment transitioning from state s to s ′ given that the agent acts with action a ; the rewardfunction R : S ×
Revised: A × S → [0 , 1] defines the probability of the environment transitioning from state s to s ′ given that the agent acts with action a ; the reward function R : S ×

Original: A → R maps a state-action transition to a real number; and 0 ≤ γ ≤ 1 is the75 discount factor, which controls how much an agent should prefer rewards sooner rather than later.
Revised: A → R maps a state-action transition to a real number; and 0 ≤ γ ≤ 1 is the discount factor, which controls how much an agent should prefer rewards sooner rather than later.

Original: The performance a policy, defined by π ( a | s ) , is typically measured using the cumulative discountedreturn (cid:80) t γ t r t that the policy achieves in an episode.
Revised: The performance a policy, defined by π ( a | s ) , is typically measured using the cumulative discounted return (cid:80) t γ t r t that the policy achieves in an episode.

Original: Central to our work is algorithms that do reinforcement learning via supervised learning (RvS) suchas Decision Transformer [1].
Revised: Central to our work is algorithms that do reinforcement learning via supervised learning (RvS) such as Decision Transformer [1].

Original: These approaches train a model using supervised learning on a datasetof trajectories to predict p D ( a | s, R ) - the probability of next action conditioned on the current stateand on the agent getting a cumulative discounted return R = (cid:80) t γ t r t .
Revised: These approaches train a model using supervised learning on a dataset of trajectories to predict p D ( a | s, R ) - the probability of next action conditioned on the current state and on the agent getting a cumulative discounted return R = (cid:80) t γ t r t .

Original: At evaluation time, the modelis conditioned on a desired target return and the agent takes the actions predicted by the model.
Revised: At evaluation time, the model is conditioned on a desired target return and the agent takes the actions predicted by the model.

Original: To understand the problem with RvS in stochastic environments, we build off of a general frameworkwhere a goal is presented to a policy, the policy acts in the environment, and the agent is scoredon how well the goal was achieved as described in Furuta et al.
Revised: To understand the problem with RvS in stochastic environments, we build off of a general framework where a goal is presented to a policy, the policy acts in the environment, and the agent is scored [Figure] gamble on how well the goal was achieved as described in Furuta et al.

Original: Mathematically, the agent’sobjective can be expressed as minimizing some distance metric D ( · , · ) between a target goal z andthe information statistics I ( τ ) of the trajectory generated by a conditional policy π ( a | s, z ) :
Revised: Mathematically, the agent’s objective can be expressed as minimizing some distance metric D ( · , · ) between a target goal z and the information statistics I ( τ ) of the trajectory generated by a conditional policy π ( a | s, z ) :

Original: [Equation]
Revised: [Equation]

Original: where p ( z ) is the goal distribution, p π z ( τ ) is the trajectory distribution ( s 1 , a 1 , s 2 , a 2 , . [Equation]
Revised: where p ( z ) is the goal distribution, p π z ( τ ) is the trajectory distribution ( s 1 , a 1 , s 2 , a 2 , . . . )

Original: obtainedby running rollouts in an MDP following policy π ( a | s, z ) .
Revised: obtained by running rollouts in an MDP following policy π ( a | s, z ) .

Original: In this section, we refer to the goalpresented to an agent as z and the function that calculates which goal was achieved in a particulartrajectory τ as its statistics I ( τ ) .
Revised: In this section, we refer to the goal presented to an agent as z and the function that calculates which goal was achieved in a particular trajectory τ as its statistics I ( τ ) .

Original: As discussed in Furuta et al.
Revised: As discussed in Furuta et al.

Original: [18], this framework includes a numberof popular algorithms, including RvS algorithms that condition on target returns (e.g., [1, 2, 8, 6]) orgoals (e.g., [7, 19]).
Revised: [18], this framework includes a number of popular algorithms, including RvS algorithms that condition on target returns (e.g., [1, 2, 8, 6]) or goals (e.g., [7, 19]).

Original: For example, in Decision Transformer [1] cumulative discounted reward is usedas the trajectory statistics I ( τ ) = (cid:80) t γ t r t and a reasonable choice for D ( · , · ) could be the squarederror between the target and achieved return.
Revised: For example, in Decision Transformer [1] cumulative discounted reward is used as the trajectory statistics I ( τ ) = (cid:80) t γ t r t and a reasonable choice for D ( · , · ) could be the squared error between the target and achieved return.

Original: Furuta et al.
Revised: Furuta et al.

Original: [18] justify using the supervised learning approach (RvS) to minimize Equation 2.1by claiming that trajectories with statistics I ( τ ) act as optimal examples of how to act to achieve
Revised: [18] justify using the supervised learning approach (RvS) to minimize Equation 2. by claiming that trajectories with statistics I ( τ ) act as optimal examples of how to act to achieve

Original: D ( I ( τ ) , z ) = 0 for an agent whose goal is to achieve z
Revised: D ( I ( τ ) , z ) = 0 for an agent whose goal is to achieve z

Original: = I ( τ ) .
Revised: = I ( τ ) .

Original: However, in stochastic environments,the actions taken in these trajectories may not be optimal since a trajectory may achieve a goal dueto uncontrollable randomness in the environment rather than from the agent’s own actions.
Revised: However, in stochastic environments, the actions taken in these trajectories may not be optimal since a trajectory may achieve a goal due to uncontrollable randomness in the environment rather than from the agent’s own actions.

Original: Infact, as discussed in Paster et al.
Revised: In fact, as discussed in Paster et al.

Original: [19], not only can the actions taken in these “optimal trajectories”actually be sub-optimal, but they can be arbitrarily bad depending on the policy used to collect thedata, the environment dynamics, and the choice of distance metric.
Revised: [19], not only can the actions taken in these “optimal trajectories” actually be sub-optimal, but they can be arbitrarily bad depending on the policy used to collect the data, the environment dynamics, and the choice of distance metric.

Original: In this section, we propose to sidestep this issue by limiting the types of statistics I ( τ ) we areinterested in to those that are not affected by uncontrollable randomness in the environment.
Revised: In this section, we propose to sidestep this issue by limiting the types of statistics I ( τ ) we are interested in to those that are not affected by uncontrollable randomness in the environment.

Original: We showfirst that RvS policies trained to reach such goals, under the assumption of infinite data and modelcapacity, learn optimal policies for achieving such goals and that goals that are independent fromenvironment stochasticity are the only goals that RvS can reliably achieve.
Revised: We show first that RvS policies trained to reach such goals, under the assumption of infinite data and model capacity, learn optimal policies for achieving such goals and that goals that are independent from environment stochasticity are the only goals that RvS can reliably achieve.

Original: Importantly, we arguethat limiting the form of trajectory statistics I ( τ ) in this way is not prohibitive and that desirablequantities can be transformed to fit this property in an intuitive way (e.g., return turns into expectedreturn when optimized to be independent of environment stochasticity).
Revised: Importantly, we argue that limiting the form of trajectory statistics I ( τ ) in this way is not prohibitive and that desirable quantities can be transformed to fit this property in an intuitive way (e.g., return turns into expected return when optimized to be independent of environment stochasticity).

Original: Finally, we propose anapproach for automatically finding such trajectory statistics using adversarial learning.
Revised: Finally, we propose an approach for automatically finding such trajectory statistics using adversarial learning.

Original: 2.2 Stochasticity Independent Representations
Revised: 2.2 Stochasticity Independent Representations

Original: Definition 2.1 (Consistently Achievable) .
Revised: Definition 2.1 (Consistently Achievable) .

Original: A goal z is consistently achievable from state s [-
 underpolicy-] π ( a | s, z ) if E τ ∼ p πz ( τ | s 0 )
Revised: A goal z is consistently achievable from state s [-
 underpolicy-]0 under policy π ( a | s, z ) if E τ ∼ p πz ( τ | s 0 )

Original: [ D ( I ( τ ) , z )] = 0 .
Revised: [ D ( I ( τ ) , z )] = 0 .

Original: A natural question is under which circumstances will the RvS approach actually minimize Equation 2.1 to zero?

Original: Clearly the approach works empirically on deterministic environments [1, 2] sincesimply replaying an action sequence that achieved a goal once will achieve it again.

Original: While not

Original: Adversarial Clustering
Revised: Adversarial Clustering

Original: Phase
Revised: Phase

Original: Estimate Average Returns
Revised: Estimate Average Returns

Original: Phase
Revised: Phase

Original: Train RvS Agent

Original: [Figure]
Revised: [Figure]

Revised: c d e

Revised: A natural question is under which circumstances will the RvS approach actually minimize Equation 2.1 to zero?

Revised: Clearly the approach works empirically on deterministic environments [1, 2] since simply replaying an action sequence that achieved a goal once will achieve it again.

Original: a b c d e all statistics of a trajectory will lead to consistently achievable goals under policies trained with
Revised: While not all statistics of a trajectory will lead to consistently achievable goals under policies trained with

Original: RvS, the supervised learning approach will minimize Equation 2.1 if I ( τ ) is independent fromenvironment stochasticity .
Revised: RvS, the supervised learning approach will minimize Equation 2.1 if I ( τ ) is independent from environment stochasticity .

Original: Theorem 2.1.
Revised: Theorem 2.1.

Original: Let π D be a data collecting policy used to gather data used to train an RvS policy, and assume this policy is trained such that π ( a | s, z )
Revised: Let π D be a data collecting policy used to gather data used to train an RvS policy, and assume this policy is trained such that π ( a | s, z )

Original: = p π D ( a | s, I ( τ ) = z ) .
Revised: = p π D ( a | s, I ( τ ) = z ) .

Original: Then, for anygoal z such that p π D ( I ( τ ) =
Revised: Then, for any goal z such that p π D ( I ( τ ) =

Original: z | s 0 ) > 0 , goal z is consistently achievable iff p π D
Revised: z | s 0 ) > 0 , goal z is consistently achievable iff p π D

Original: ( s t +1 | s t , a t ) =
Revised: ( s t +1 | s t , a t ) =

Original: p π D ( s t +1 | s t , a t , I ( τ )) .
Revised: p π D ( s t +1 | s t , a t , I ( τ )) .

Original: Proof.
Revised: Proof.

Original: [Sketch] The forward direction of the theorem is proven by rewriting E τ ∼ p πz ( τ | s 0 )
Revised: [Sketch] The forward direction of the theorem is proven by rewriting E τ ∼ p πz ( τ | s 0 )

Original: [ D ( I ( τ ) , z )]to be taken over p π z ( τ | s 0 , I ( τ )
Revised: [ D ( I ( τ ) , z )] to be taken over p π z ( τ | s 0 , I ( τ )

Original: =
Revised: =

Original: z ) using the independence assumption combined with the form ofthe policy.
Revised: z ) using the independence assumption combined with the form of the policy.

Original: For the other direction and the full proof, refer to Appendix A.1.
Revised: For the other direction and the full proof, refer to Appendix A.1.

Original: Remarks.
Revised: Remarks.

Original: The above theorem has two major implications.
Revised: The above theorem has two major implications.

Original: First is that in stochastic environments,many of the most common statistics such as final states or cumulative discounted rewards are oftencorrelated with environment stochasticity and therefore are unlikely to work well with the behavioralcloning approach.
Revised: First is that in stochastic environments, many of the most common statistics such as final states or cumulative discounted rewards are often correlated with environment stochasticity and therefore are unlikely to work well with the behavioral cloning approach.

Original: For example, in the gambling environment discussed in Figure 1, the achievedreward is correlated with the uncontrollable randomness present when gambling, and therefore whenfiltering an offline dataset for trajectories that achieve a high reward, the resulting trajectories willhave unrealistic environment dynamics because we ignore the unsuccessful attempts (i.e., the agentalways wins the money when gambling) and an agent behavioral cloned on these trajectories will notactually end up achieving a high reward consistently.
Revised: For example, in the gambling environment discussed in Figure 1, the achieved reward is correlated with the uncontrollable randomness present when gambling, and therefore when filtering an offline dataset for trajectories that achieve a high reward, the resulting trajectories will have unrealistic environment dynamics because we ignore the unsuccessful attempts (i.e., the agent always wins the money when gambling) and an agent behavioral cloned on these trajectories will not actually end up achieving a high reward consistently.

Original: Theorem 2.1 also gives a solution to ensure asymptotically that RvS policies will consistently achievegoals: use trajectory statistics that are independent from the stochasticity in environment transitions.
Revised: Theorem 2.1 also gives a solution to ensure asymptotically that RvS policies will consistently achieve goals: use trajectory statistics that are independent from the stochasticity in environment transitions.

Original: This requires the statistic function I ( τ ) to return the same value for different trajectories where theonly difference is in the environment transitions, making I ( τ ) essentially act as a cluster assignmentwhere each trajectory within a cluster has the same value of I ( τ ) and has environment transitionssampled from the true distribution p ( s t +1 | s t , a t ) .
Revised: This requires the statistic function I ( τ ) to return the same value for different trajectories where the only difference is in the environment transitions, making I ( τ ) essentially act as a cluster assignment where each trajectory within a cluster has the same value of I ( τ ) and has environment transitions sampled from the true distribution p ( s t +1 | s t , a t ) .

Original: These clusters can be thought of as datasets eachgenerated by distinct policies in the same environment.
Revised: These clusters can be thought of as datasets each generated by distinct policies in the same environment.

Original: For example, in the gambling environmenteach cluster could include trajectories that all took the same action.
Revised: For example, in the gambling environment each cluster could include trajectories that all took the same action.

Original: Ultimately, we want to use a target average return to control an RvS agent.
Revised: Ultimately, we want to use a target average return to control an RvS agent.

Original: Luckily, given any I ( τ )that satisfies Theorem 2.1, we can transform it by using the average returns of trajectories withineach cluster as the statistic instead.
Revised: Luckily, given any I ( τ ) that satisfies Theorem 2.1, we can transform it by using the average returns of trajectories within each cluster as the statistic instead.

Original: In the next section, we describe a practical algorithm for learningan I ( τ ) that satisfies Theorem 2.1, finding the average return of trajectories with the same value of
Revised: In the next section, we describe a practical algorithm for learning an I ( τ ) that satisfies Theorem 2.1, finding the average return of trajectories with the same value of

Original: I ( τ ) (i.e. in the same “cluster”), and training an RvS policy conditioned on these estimated averagereturns.
Revised: I ( τ ) (i.e. in the same “cluster”), and training an RvS policy conditioned on these estimated average returns.

Original: 2.3 Learning Stochasticity-Independent Representations
Revised: 2.3 Learning Stochasticity-Independent Representations

Original: In this section we propose a method for learning such trajectory representations, which we call
Revised: In this section we propose a method for learning such trajectory representations, which we call

Original: ESPER, short for e nvironment- s tochasticity-inde pe ndent r epresentations.
Revised: ESPER, short for e nvironment- s tochasticity-inde pe ndent r epresentations.

Original: Rather than use trajectoryreturn as the statistic I ( τ ) , which is often dependent on environment stochasticity and therefore may result in unexpected behavior, ESPER uses a neural network to learn statistics that are independentfrom environment stochasticity.
Revised: Rather than use trajectory return as the statistic I ( τ ) , which is often dependent on environment stochasticity and therefore may result in unexpected behavior, ESPER uses a neural network to learn statistics that are independent from environment stochasticity.

Original: It does so in three phases: first, it uses an auto-encoder framework tolearn a discrete representation I ( τ ) acting as a cluster assignment for each trajectory along with anadversarial loss that changes the representation in order to hurt the predictions of a dynamics model;second, a model learns to predict average trajectory return from the learned representation; and third,an RvS agent is trained to predict actions taken when conditioned on a state and estimated averagereturn.
Revised: It does so in three phases: first, it uses an auto-encoder framework to learn a discrete representation I ( τ ) acting as a cluster assignment for each trajectory along with an adversarial loss that changes the representation in order to hurt the predictions of a dynamics model; second, a model learns to predict average trajectory return from the learned representation; and third, an RvS agent is trained to predict actions taken when conditioned on a state and estimated average return.

Original: For some trajectory τ = ( s 1 , a 1 , r 1 , s 2 , a 2 , r 2 , . . . )
Revised: For some trajectory τ = ( s 1 , a 1 , r 1 , s 2 , a 2 , r 2 , . . . )

Original: , ESPER trains the following parameterized modelsin addition to a vanilla RvS policy π ξ :
Revised: , ESPER trains the following parameterized models in addition to a vanilla RvS policy π ξ :

Original: Clustering model:
Revised: Clustering model:

Original: [Equation]
Revised: [Equation]

Original: Action predictor:
Revised: Action predictor:

Original: [Equation]
Revised: [Equation]

Original: Return predictor:
Revised: Return predictor:

Original: [Equation]
Revised: [Equation]

Original: Transition predictor:
Revised: Transition predictor:

Original: [Equation]
Revised: [Equation]

Original: Adversarial clustering.
Revised: Adversarial clustering.

Original: Here, we will focus on learning representation in the form of clusteringassignments.
Revised: Here, we will focus on learning representation in the form of clustering assignments.

Original: The goal is to learn a clustering assignment per trajectory
Revised: The goal is to learn a clustering assignment per trajectory

Original: I ( τ ) to minimize thediscrepancy between the true environment transition p ( s t +1 | s t , a t ) and the estimated transitionp ϕ ( s t +1 | s t , a t , I ( τ )) from the replay buffer when conditioned on I ( τ ) .
Revised: I ( τ ) to minimize the discrepancy between the true environment transition p ( s t +1 | s t , a t ) and the estimated transition p ϕ ( s t +1 | s t , a t , I ( τ )) from the replay buffer when conditioned on I ( τ ) .

Original: Such a representation willsatisfy Theorem 2.1.
Revised: Such a representation will satisfy Theorem 2.1.

Original: In the first phase of training, discrete cluster assignments produced by theclustering model are trained by taking alternating gradient steps with the following two losses toensure that information present in I ( τ ) does not improve the performance of the dynamics model:[Equation] policy reconstruction
Revised: In the first phase of training, discrete cluster assignments produced by the clustering model are trained by taking alternating gradient steps with the following two losses to ensure that information present in I ( τ ) does not improve the performance of the dynamics model:

Original: [Equation]
Revised: [Equation]

Original: adversarial

Original: [Equation] dynamics prediction [Equation] In this step, the dynamics prediction loss is predicting the next state given the current state, action,and cluster assignment I ( τ ) while the adversarial loss tries to change I ( τ ) in order to hurt thisprediction.
Revised: In this step, the dynamics prediction loss is predicting the next state given the current state, action, and cluster assignment I ( τ ) while the adversarial loss tries to change I ( τ ) in order to hurt this prediction.

Original: Since a constant representation minimizes this loss, we use a policy reconstruction loss toencourage I ( τ ) to contain information about the policy used to generate the trajectory and encouragethe formation of more than one cluster.
Revised: Since a constant representation minimizes this loss, we use a policy reconstruction loss to encourage I ( τ ) to contain information about the policy used to generate the trajectory and encourage the formation of more than one cluster.

Original: β act and β adv are hyperparameters to balance the strength ofthis policy reconstruction loss with the adversarial loss.
Revised: β act and β adv are hyperparameters to balance the strength of this policy reconstruction loss with the adversarial loss.

Original: Estimate cluster average returns.
Revised: Estimate cluster average returns.

Original: After clustering, we learn to predict discounted returns R =180 (cid:80) t γ t r t for each cluster.
Revised: After clustering, we learn to predict discounted returns R = (cid:80) t γ t r t for each cluster.

Original: The return predictor parameterized by ψ is trained using the following loss:
Revised: The return predictor parameterized by ψ is trained using the following loss:

Original: [Equation]
Revised: [Equation]

Original: The above losses are described such that the model produces a single estimated return per trajectory.
Revised: The above losses are described such that the model produces a single estimated return per trajectory.

Original: However, in practice it can be extended to produce a value for each time-step in a trajectory by usingsuffixes of trajectories τ i,t = ( a t , s t , r t , . .
Revised: However, in practice it can be extended to produce a value for each time-step in a trajectory by using suffixes of trajectories τ i,t = ( a t , s t , r t , . .

Original: . , a T , s T , r T ) to generate return predictions ˆ R t for eachtime-step.
Revised: . , a T , s T , r T ) to generate return predictions ˆ R t for each time-step.

Original: Training policy on predicted returns.
Revised: Training policy on predicted returns.

Original: After learning expected future returns for each step in atrajectory, we use the dataset of ( s t , a t , ˆ R t ) triples to learn a policy π ξ ( a t
Revised: After learning expected future returns for each step in a trajectory, we use the dataset of ( s t , a t , ˆ R t ) triples to learn a policy π ξ ( a t

Original: | s t , ˆ R t ) parameterized byξ by predicting the action from the current state and estimated expected return by minimizing thefollowing loss just as in prior RvS works [2, 1]:
Revised: | s t , ˆ R t ) parameterized by ξ by predicting the action from the current state and estimated expected return by minimizing the following loss just as in prior RvS works [2, 1]:

Original: [Equation]
Revised: [Equation]

Original: 2.4 Implementation

Original: The clustering model is implemented using an LSTM [20], using truncated backpropagation throughtime for training on long trajectories (up to 1000 time-steps in our experiments).

Original: The action predictor,return predictor, transition predictor, and policy are all implemented using MLPs, though our method

Original: Task
Revised: Task

Original: Return-Conditioned RvS
Revised: Return-Conditioned RvS

Original: CQL
Revised: CQL

Original: ESPER (Ours)
Revised: ESPER (Ours)

Original: Gambling -0.05 (0.27) 1.0 (0.0) 1.0 (0.0)
Revised: Gambling -0.05 (0.27) 1.0 (0.0) 1.0 (0.0)

Original: Connect Four 0.24 (0.15) 0.61 (0.05) 0.99 (0.01) 0.56 (0.03) 0.7 (0.09) 0.81 (0.05)  could easily be extended to use a transformer policy as in Decision Transformer [1] (see section 3.2).
Revised: Connect Four 0.24 (0.15) 0.61 (0.05) 0.99 (0.01) 0.56 (0.03) 0.7 (0.09) 0.81 (0.05) 2.4 Implementation The clustering model is implemented using an LSTM [20], using truncated backpropagation through time for training on long trajectories (up to 1000 time-steps in our experiments). The action predictor, return predictor, transition predictor, and policy are all implemented using MLPs, though our method could easily be extended to use a transformer policy as in Decision Transformer [1] (see section 3.2).

Original: Our clustering outputs a discrete value sampled from a categorical distribution and we use gumbelsoftmax [21, 22] in order to backpropagate gradients through it.
Revised: Our clustering outputs a discrete value sampled from a categorical distribution and we use gumbelsoftmax [21, 22] in order to backpropagate gradients through it.

Original: We use normal distributions with unitvariances for the transition predictor.
Revised: We use normal distributions with unit variances for the transition predictor.

Original: Code for our implementation will be available upon publication.
Revised: Code for our implementation will be available upon publication at https://sites.google.com/view/esper-paper.

Original: See appendix A.6 for pseudocode for the adversarial clustering step of ESPER.
Revised: See appendix A.6 for pseudocode for the adversarial clustering step of ESPER.

Original: See appendix A.4 formore information about the implementation, including specific hyperparameters for our environmentsas well as best practices.
Revised: See appendix A.4 for more information about the implementation, including specific hyperparameters for our environments as well as best practices.

Original:  Experiments
Revised:  Experiments

Original: We designed our experiments to answer the following questions: (i) How well does return-conditioned
Revised: We designed our experiments to answer the following questions: (i) How well does return-conditioned

Original: RvS work on stochastic problems where returns are correlated with environment stochasticity?
Revised: RvS work on stochastic problems where returns are correlated with environment stochasticity?

Original: (ii)
Revised: (ii)

Original: Does ESPER improve performance, both in terms of maximum achievable performance as well asin correlation between the target and achieved returns?
Revised: Does ESPER improve performance, both in terms of maximum achievable performance as well as in correlation between the target and achieved returns?

Original: (iii) Does the degree to which the learnedrepresentations in ESPER are independent from environment stochasticity measurably affect theperformance of the learned agent?
Revised: (iii) Does the degree to which the learned representations in ESPER are independent from environment stochasticity measurably affect the performance of the learned agent?

Original: 3.1 Stochastic Benchmark Tasks
Revised: 3.1 Stochastic Benchmark Tasks

Original: Prior offline RL methods including Decision Transformer have been tested primarily on deterministicor near-deterministic environments such as continuous control tasks from the D4RL dataset [13] and
Revised: Prior offline RL methods including Decision Transformer have been tested primarily on deterministic or near-deterministic environments such as continuous control tasks from the D4RL dataset [13] and

Original: Atari games [16].
Revised: Atari games [16].

Original: To get a better idea of how these methods and our approach work on realistic,stochastic environments, we created three new benchmark tasks.
Revised: To get a better idea of how these methods and our approach work on realistic, stochastic environments, we created three new benchmark tasks.

Original: Since we are evaluating in an offlinesetting, each task consists of a stochastic environment as well as one or several data collection policiesin order to gather the offline dataset.
Revised: Since we are evaluating in an offline setting, each task consists of a stochastic environment as well as one or several data collection policies in order to gather the offline dataset.

Original: Gambling.
Revised: Gambling.

Original: An illustrative gambling environment described in Figure 1 with only 3 actions.
Revised: An illustrative gambling environment described in Figure 1 with only 3 actions.

Original: Theoffline dataset for this task is 100k steps collected with a random policy.
Revised: The offline dataset for this task is 100k steps collected with a random policy.

Original: Connect Four.
Revised: Connect Four.

Original: A game of Connect Four, a tile-based game where players compete to get four-in-arow, against an opponent that does not place a tile in the rightmost column with a probability of 20% .
Revised: A game of Connect Four, a tile-based game where players compete to get four-in-arow, against an opponent that does not place a tile in the rightmost column with a probability of 20% .

Original: The agent gets 1 reward for winning, 0 for a draw, and − 1 for losing.
Revised: The agent gets 1 reward for winning, 0 for a draw, and − 1 for losing.

Original: The offline dataset for this taskis 1M steps collected using a mixture of an ϵ -greedy policy and a policy that always places tiles inthe rightmost column.2048.
Revised: The offline dataset for this task is 1M steps collected using a mixture of an ϵ -greedy policy and a policy that always places tiles in the rightmost column.

Original: A simplified version of 2048 [23], a sliding puzzle game where alike numbers can be combinedto create tiles of larger value, where the agent gets a reward of 1 by creating a 128 tile and gets noreward otherwise.
Revised: A simplified version of 2048 [23], a sliding puzzle game where alike numbers can be combined to create tiles of larger value, where the agent gets a reward of 1 by creating a 128 tile and gets no reward otherwise.

Original: The offline dataset for this task is 5M steps collected using a mixture of a randomagent and an expert policy trained using PPO [24].
Revised: The offline dataset for this task is 5M steps collected using a mixture of a random agent and an expert policy trained using PPO [24].

Original: Detailed information about each environment can be found in Appendix A.2.
Revised: Detailed information about each environment can be found in Appendix A.2.

Original: 3.2 Baselines

Original: We compare our method primarily against RvS-R, the simple return-conditioned model describedin Emmons et al.

Original: [2] which trains an MLP to predict the next action given the current state and thetarget return.

Original: Our method uses the same model and training code, except rather than condition on

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: actions from environment stochasticity and is able achieve performance close to the target return.
Revised: actions from environment stochasticity and is able achieve performance close to the target return.

Original: Interestingly, the distribution of expected returns learned by ESPER match the performance-level of the data-collection policies.
Revised: Interestingly, the distribution of expected returns learned by ESPER match the performance-level of the data-collection policies.

Original:  target returns we condition on expected returns learned using ESPER.
Revised: We compare our method primarily against RvS-R, the simple return-conditioned model described in Emmons et al. [2] which trains an MLP to predict the next action given the current state and the target return. Our method uses the same model and training code, except rather than condition on target returns we condition on expected returns learned using ESPER.

Original: We opt to not use a transformerpolicy [1] since our tasks are all MDPs and prior work has found that MLPs can achieve competitiveor better performance on Markov tasks while using significantly less compute [2].
Revised: We opt to not use a transformer policy [1] since our tasks are all MDPs and prior work has found that MLPs can achieve competitive or better performance on Markov tasks while using significantly less compute [2].

Original: We also compareagainst a strong value-based approach for offline-RL called Conservative Q-Learning (CQL) [4],which has been found to have strong performance on both discrete and continuous offline RL tasks.
Revised: We also compare against a strong value-based approach for offline-RL called Conservative Q-Learning (CQL) [4], which has been found to have strong performance on both discrete and continuous offline RL tasks.

Original: 3.3 Performance in Stochastic Environments
Revised: 3.3 Performance in Stochastic Environments

Original: When comparing the performance of our approach, ESPER, with return-conditioned RvS, we foundthat ESPER consistently achieves stronger alignment between target return and average performance while also achieving a higher maximum level of performance when tuning the target return.
Revised: When comparing the performance of our approach, ESPER, with return-conditioned RvS, we found that ESPER consistently achieves stronger alignment between target return and average performance while also achieving a higher maximum level of performance when tuning the target return.

Original: As shown in Figure 5a, the return-conditioned baseline cannot achieve the maximum performanceof 1 while ESPER learns behaviors corresponding to average performances ranging from − 5 to 1 .
Revised: As shown in Figure 5a, the return-conditioned baseline cannot achieve the maximum performance of 1 while ESPER learns behaviors corresponding to average performances ranging from − 5 to 1 .

Original: In Connect Four (Figure 4), despite seeing many examples of winning trajectories, an RvS agentcannot achieve more than 0 .
Revised: In Connect Four (Figure 4), despite seeing many examples of winning trajectories, an RvS agent [Figure] r [Figure]

Revised: a

Revised: [Figure]

Original:  2 average return (corresponding to a win-rate of 60% ) while ESPER canboth win and lose with perfect accuracy depending on the target performance.
Revised: ESPER (Ours) [Figure] cannot achieve more than 0 . 2 average return (corresponding to a win-rate of 60% ) while ESPER can both win and lose with perfect accuracy depending on the target performance.

Original: In 2048 (Figure 5b),
Revised: In 2048 (Figure 5b),

Original: RvS cannot disentangle trajectories where the reward is high simply from luck from trajectories thattook good actions and can only win the game about 60% of the time.
Revised: RvS cannot disentangle trajectories where the reward is high simply from luck from trajectories that took good actions and can only win the game about 60% of the time.

Original: ESPER learns many modes ofbehavior, and can be controlled to win anywhere from 30% to 80% of the time.[Figure] [Figure] [Figure]
Revised: ESPER learns many modes of behavior, and can be controlled to win anywhere from 30% to 80% of the time.

Original: [Figure]

Original: [Figure]

Original: P e

Original: r

Original: f

Original: o r m a n ce ESPER (Ours) [Figure] Additionally, Figure 6 shows that this lack of performance for the return-conditioned agent isnot due to a lack of data .
Revised: Additionally, Figure 6 shows that this lack of performance for the return-conditioned agent is not due to a lack of data .

Original: In the figure, going from 5% data to 100% (5 million frames) does notaffect the performance of RvS while the performance of ESPER gets closer to matching the targetreturn with more data.
Revised: In the figure, going from 5% data to 100% (5 million frames) does not affect the performance of RvS while the performance of ESPER gets closer to matching the target return with more data.

Original: Finally, we compared the maximum performance of RvS and ESPER with thestate-of-the-art value-based offline RL algorithm, Conservative Q-Learning (CQL) [4].
Revised: Finally, we compared the maximum performance of RvS and ESPER with the state-of-the-art value-based offline RL algorithm, Conservative Q-Learning (CQL) [4].

Original: As shown in
Revised: As shown in

Original: Figure 3, ESPER and CQL achieve perfect performance on the gambling task while RvS cannot getpositive reward.
Revised: Figure 3, ESPER and CQL achieve perfect performance on the gambling task while RvS cannot get positive reward.

Original: Remarkably, on Connect 4 and 2048, while CQL can get respectable performance(better than the average return of the offline dataset), ESPER achieves state-of-the-art performance .
Revised: Remarkably, on Connect 4 and 2048, while CQL can get respectable performance (better than the average return of the offline dataset), ESPER achieves state-of-the-art performance .

Original: 3.4 Learned Representations and Behaviors
Revised: 3.4 Learned Representations and Behaviors

Original: Rather than simply condition on returns, ESPER conditions on learned expected return values.
Revised: Rather than simply condition on returns, ESPER conditions on learned expected return values.

Original: Inaddition to plotting the performance against the target return, we also use a histogram to show thedistribution of values (returns or learned expected returns) on which the agents are trained.
Revised: In addition to plotting the performance against the target return, we also use a histogram to show the distribution of values (returns or learned expected returns) on which the agents are trained.

Original: As shownin the histograms (Figure 4, Figure 5a, Figure 5b), ESPER often learns values for many more modesof behaviors corresponding to many different expected returns.
Revised: As shown in the histograms (Figure 4, Figure 5a, Figure 5b), ESPER often learns values for many more modes of behaviors corresponding to many different expected returns.

Original: Unlike RvS, when an agent is trainedon a particular expected return value, the actual average performance of the agent is often close tothis value.
Revised: Unlike RvS, when an agent is trained on a particular expected return value, the actual average performance of the agent is often close to this value.

Original: This is useful for offline settings where tuning the target return in an online environmentis not feasible, since with ESPER one can have confidence in which levels of performancethe model will be able to achieve .
Revised: This is useful for offline settings where tuning the target return in an online environment is not feasible, since with ESPER one can have confidence in which levels of performance the model will be able to achieve .

Original: In contrast, with return-conditioning, the performance of thereturn-conditioned policy may not correlate with the returns on which it was trained.
Revised: In contrast, with return-conditioning, the performance of the return-conditioned policy may not correlate with the returns on which it was trained.

Original: Finally, we empirically measured the relationship between the independence of learned trajectorystatistics and the performance of the agent.
Revised: Finally, we empirically measured the relationship between the independence of learned trajectory statistics and the performance of the agent.

Original: As shown in appendix A.3, we trained several agentswith different random seeds and different hyperparameters on the gambling task and measuredthe dynamics loss (Equation 2.3) and performance of each agent.
Revised: As shown in appendix A.3, we trained several agents with different random seeds and different hyperparameters on the gambling task and measured the dynamics loss (Equation 2.3) and performance of each agent.

Original: Indeed, agents with trajectorystatistics that enable the dynamics model to “cheat” and get a low loss performed worse thanthose with representations independent from environment stochasticity , empirically confirming
Revised: Indeed, agents with trajectory statistics that enable the dynamics model to “cheat” and get a low loss performed worse than those with representations independent from environment stochasticity , empirically confirming

Original: Theorem 2.1.
Revised: Theorem 2.1.

Original:  Related Work
Revised:  Related Work

Original: Offline Reinforcement Learning.
Revised: Offline Reinforcement Learning.

Original: Offline reinforcement learning is a framework where loggeddata is used to learn behaviors [3].
Revised: Offline reinforcement learning is a framework where logged data is used to learn behaviors [3].

Original: In contrast to most online RL methods, offline RL avoids theneed for often-expensive online data collection by using fixed datasets of trajectories generated byvarious policies and stitching together observed behaviors in order to optimize a reward function.
Revised: In contrast to most online RL methods, offline RL avoids the need for often-expensive online data collection by using fixed datasets of trajectories generated by various policies and stitching together observed behaviors in order to optimize a reward function.

Original: The main challenge in the offline setting is that offline RL agents cannot collect more data to reducetheir uncertainty and therefore offline RL approaches often have some type of value pessimism orpolicy constraints to keep the learned policy within the support of the offline dataset.
Revised: The main challenge in the offline setting is that offline RL agents cannot collect more data to reduce their uncertainty and therefore offline RL approaches often have some type of value pessimism or policy constraints to keep the learned policy within the support of the offline dataset.

Original: Conservative
Revised: Conservative

Original: Q-Learning [4] accomplishes this by learning a conservative lower-bound on the value of the currentpolicy and achieve state-of-the-art performance on most offline-RL benchmarks.
Revised: Q-Learning [4] accomplishes this by learning a conservative lower-bound on the value of the current policy and achieve state-of-the-art performance on most offline-RL benchmarks.

Original: Other methods relyon ensembles (e.g., [5]) or policy regularization (e.g., [25–27]).
Revised: Other methods rely on ensembles (e.g., [5]) or policy regularization (e.g., [25–27]).

Original: Several standardized benchmarks have emerged for testing offline RL agents, including D4RL [13] and RL Unplugged [14].
Revised: Several standardized benchmarks have emerged for testing offline RL agents, including D4RL [13] and RL Unplugged [14].

Original: In contrastto our work, prior methods for offline RL primarily evaluate on deterministic environments, such asrobotics and locomotion tasks as well as Atari tasks, which are near-deterministic [4].
Revised: In contrast to our work, prior methods for offline RL primarily evaluate on deterministic environments, such as robotics and locomotion tasks as well as Atari tasks, which are near-deterministic [4].

Original: Reinforcement Learning via Supervised Learning.
Revised: Reinforcement Learning via Supervised Learning.

Original: At the core of our work is the idea of reducingreinforcement learning to a prediction problem and solving it via supervised learning.
Revised: At the core of our work is the idea of reducing reinforcement learning to a prediction problem and solving it via supervised learning.

Original: This ideawas first proposed by Schmidhuber [6] and Kumar et al.
Revised: This idea was first proposed by Schmidhuber [6] and Kumar et al.

Original: [8], who proposed to learn behaviors bypredicting actions conditioned on a desired outcome such as return.
Revised: [8], who proposed to learn behaviors by predicting actions conditioned on a desired outcome such as return.

Original: Ghosh et al.
Revised: Ghosh et al.

Original: [7] proposed to usethe same approach to solve goal-conditioned tasks and Paster et al.
Revised: [7] proposed to use the same approach to solve goal-conditioned tasks and Paster et al.

Original: [19] proposed to predict multi-stepaction sequences, showing a connection between sequence modeling and optimal open-loop planning.
Revised: [19] proposed to predict multi-step action sequences, showing a connection between sequence modeling and optimal open-loop planning.

Original: Decision Transformer [1] proposed to use a transformer [28] to condition on trajectory history whenpredicting actions and tested the approach on offline RL rather than online, achieving results competitive with value-based offline RL approaches.
Revised: Decision Transformer [1] proposed to use a transformer [28] to condition on trajectory history when predicting actions and tested the approach on offline RL rather than online, achieving results competitive with value-based offline RL approaches.

Original: Several followup extensions to Decision Transformerwere explored, including using non-return goals [18] and using pretrained transformers [29].
Revised: Several followup extensions to Decision Transformer were explored, including using non-return goals [18] and using pretrained transformers [29].

Original: In arecent work, Emmons et al.
Revised: In a recent work, Emmons et al.

Original: [2] coined the term RvS (reinforcement learning via supervised learning)to describe such approaches.
Revised: [2] coined the term RvS (reinforcement learning via supervised learning) to describe such approaches.

Original: Notably found that a transformer is not necessary to perform well onmost tasks, showing that a two-layer MLP with sufficient regularization can actually outperformtransformers at a fraction of the computational cost.
Revised: Notably found that a transformer is not necessary to perform well on most tasks, showing that a two-layer MLP with sufficient regularization can actually outperform transformers at a fraction of the computational cost.

Original: No prior RvS approach has been thoroughlyevaluated in stochastic environments.
Revised: No prior RvS approach has been thoroughly evaluated in stochastic environments.

Original: Paster et al.
Revised: Paster et al.

Original: [19] describes a counter-example where an RvSagent in Ghosh et al.
Revised: [19] describes a counter-example where an RvS agent in Ghosh et al.

Original: [7] doesn’t converge in a stochastic environment, and proposes a solution.
Revised: [7] doesn’t converge in a stochastic environment, and proposes a solution.

Original: However, the approach can only be used to plan action sequences rather than reactive plans andtherefore won’t achieve maximal performance in many stochastic settings.
Revised: However, the approach can only be used to plan action sequences rather than reactive plans and therefore won’t achieve maximal performance in many stochastic settings.

Original: 4.1 Adversarial Learning
Revised: 4.1 Adversarial Learning

Original: ESPER uses an adversarial loss to ensure that trajectory statistics are independent from environment[Figure] stochasticity.
Revised: ESPER uses an adversarial loss to ensure that trajectory statistics are independent from environment stochasticity.

Original: Adversarial losses are also widely used in fairness (e.g., [30–32]), where it is used toprevent models from using sensitive attributes, and in generative modeling such as in GANs [33].
Revised: Adversarial losses are also widely used in fairness (e.g., [30–32]), where it is used to prevent models from using sensitive attributes, and in generative modeling such as in GANs [33].

Original:  Conclusion
Revised:  Conclusion

Original: As models are scaled up in pursuit of better performance [10], it is important to acknowledge thelimitations that cannot be fixed by more scale and data alone.
Revised: As models are scaled up in pursuit of better performance [10], it is important to acknowledge the limitations that cannot be fixed by more scale and data alone.

Original: Our work points out an issue withthe increasingly popular Decision Transformer [1] and other RL via supervised learning algorithmswhen making decisions in stochastic environments.
Revised: Our work points out an issue with the increasingly popular Decision Transformer [1] and other RL via supervised learning algorithms when making decisions in stochastic environments.

Original: We show why such errors occur and show thatthese approaches work if and only if the goals that they are trained to achieve are independent fromenvironment stochasticity.
Revised: We show why such errors occur and show that these approaches work if and only if the goals that they are trained to achieve are independent from environment stochasticity.

Original: We give a practical algorithm that transforms problematic goals such astrajectory returns into expected returns that satisfy this condition.
Revised: We give a practical algorithm that transforms problematic goals such as trajectory returns into expected returns that satisfy this condition.

Original: Finally, we validate our approachby testing it on several new stochastic benchmark tasks, showing empirically that our approach vastlyoutperforms return-conditioned models in terms of alignment between the target and average achievedreturn as well as maximum performance.
Revised: Finally, we validate our approach by testing it on several new stochastic benchmark tasks, showing empirically that our approach vastly outperforms return-conditioned models in terms of alignment between the target and average achieved return as well as maximum performance.

Original: We hope that our work gives insight into how predictivemodels can be used to make optimal decisions, even in stochastic environments, hopefully paving theway for the creation of more general and useful agents.
Revised: We hope that our work gives insight into how predictive models can be used to make optimal decisions, even in stochastic environments, hopefully paving the way for the creation of more general and useful agents.

Original: Limitations.
Revised: Limitations.

Original: There are several limitations and opportunities for future work.
Revised: There are several limitations and opportunities for future work.

Original: First, while our theorygives necessary and sufficient condition for the types of trajectory statistics that are appropriate to usewith RvS agents, there can be several possible functions that satisfy this condition.
Revised: First, while our theory gives necessary and sufficient condition for the types of trajectory statistics that are appropriate to use with RvS agents, there can be several possible functions that satisfy this condition.

Original: Choosing whichfunction to use, perhaps by exploiting more information about the task, is a promising direction forfuture research.
Revised: Choosing which function to use, perhaps by exploiting more information about the task, is a promising direction for future research.

Original: While we empirically validate our approach with return goals, we did not explore theuse other types of goals where conditioning the model on the mean may be less appropriate, such asvisual goals.
Revised: While we empirically validate our approach with return goals, we did not explore the use other types of goals where conditioning the model on the mean may be less appropriate, such as visual goals.

Original: Societal Impact.
Revised: Societal Impact.

Original: We believe that this work will result in a positive societal impact, since ourapproach can help to avoid unexpected behavior when controlling an agent trained using supervisedlearning.
Revised: We believe that this work will result in a positive societal impact, since our approach can help to avoid unexpected behavior when controlling an agent trained using supervised learning.

Original: However, we acknowledge that powerful automated decision making algorithms such as
Revised: However, we acknowledge that powerful automated decision making algorithms such as

Original: ESPER have the potential to make harmful decisions when trained with under-specified objectives oron biased data.Checklist
Revised: ESPER have the potential to make harmful decisions when trained with under-specified objectives or on biased data.

Original: The checklist follows the references.

Original: Please read the checklist guidelines carefully for information onhow to answer these questions.

Original: For each question, change the default [TODO] to [Yes] , [No] , or[N/A] .

Original: You are strongly encouraged to include a justification to your answer , either by referencingthe appropriate section of your paper or providing a brief inline description.

Original: For example:

Original: • Did you include the license to the code and datasets?

Original: [Yes] See Section ?? .•

Original: Did you include the license to the code and datasets?

Original: [No] The code and the data are

Original: proprietary.• Did you include the license to the code and datasets?

Original: [N/A]Please

Original: do not modify

Original: the questions and

Original: only use

Original: the

Original: provided

Original: macros for

Original: your answers. Note that the
Revised: Acknowledgements

Original: Checklist section does not count towards the page limit In your paper, please delete this instructions
Revised: The authors thank Qizhen Zhang, Forest Yang, Dami Choi, Joyce Yang, Harris Chan, Sheng Jia,

Original: block

Original: and only keep the Checklist section heading above

Original: along with the questions/answers below. For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’scontributions and scope?
Revised: Nikita Dhawan, Jonah Philion, Michael Zhang and Yongchao Zhou for insightful discussions and helpful feedback on drafts of the paper.

Original: [Yes](b) Did you describe the limitations of your work? [Yes] See section 5.(c) Did you discuss any potential negative societal impacts of your work? [Yes] Seesection 5.(d) Have you read the ethics review guidelines and ensured that your paper conforms tothem?
Revised: We gratefully acknowledge funding from the Natural

Original: [Yes]2.

Original: If you are including theoretical results...

Original: (a) Did you state the full set of assumptions of all theoretical results?

Original: [Yes]

Original: (b) Did you include complete proofs of all theoretical results?

Original: [Yes] See the appendix inthe supplemental materials. (a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material or as a URL)? [Yes] Limited code is
Revised: Sciences and Engineering Research Council of Canada (NSERC)

Original: available in the supplemental material.

Original: (b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes] See appendix.
Revised: [JB: 2020-06904], the Canada

Original: (c) Did you report error bars (e.g., with respect to the random seed after running exper-

Original: iments multiple times)?

Original: [Yes] Experiments are run on multiple seeds with standarddeviation error bars.

Original: (d) Did you include the total amount of compute and thetype of resources used (e.g., type [Table] of GPUs, internal cluster, or cloud provider)? [Yes] See appendix.
Revised: CIFAR AI Chairs Program, Microsoft Research, the Google Research Scholar Program and the

Original: If you are using existing assets (e

Original: g., code, data, models) or curating/releasing new assets...

Original: [Table]

Original: (a) If your work uses existing assets, did you cite the creators?

Original: [N/A]

Original: [Table] (b) Did you mention the license of the assets? [N/A] [Table]
Revised: Amazon Research Award.

Original: (c) Did you include any new assets either in the supplemental material or as a URL?

Original: [N/A]

Original: [Table]

Original: (d) Did you discuss whether and how consent was obtained from people whose data you’re

Original: using/curating?

Original: [N/A]

Original: (e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A] If you used crowdsourcing or conducted research with human subjects...
Revised: Resources used in preparing this research were provided, in part, by the

Original: (a) Did you include the full text of instructions given to participants and screenshots, if

Original: applicable? [N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
Revised: Province of Ontario, the Government of Canada through CIFAR and companies sponsoring the

Original: [N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]
Revised: Vector Institute for Artificial Intelligence ( www.vectorinstitute.ai/partners ).

Original: A AppendixA.1 Proof for Theorem 2.1 
Revised: A Appendix A.1 Proof for Theorem 2.

Original: Proof.
Revised: Proof.

Original: Let π ( a | s, z ) =
Revised: Let π ( a | s, z ) =

Original: p π D ( a | s, I ( τ ) = z ) .
Revised: p π D ( a | s, I ( τ ) = z ) .

Original: If I ( τ ) ⊥ s t +1 given s t , a t , then,E
Revised: If I ( τ ) ⊥ s t +1 given s t , a t , then,

Original: [Equation]
Revised: [Equation]

Original: The second equality follows from the independence assumption and the final equality holds becausethe expectation is conditioned on I ( τ ) =
Revised: The second equality follows from the independence assumption and the final equality holds because the expectation is conditioned on I ( τ ) =

Original: z .
Revised: z .

Original: To show the other direction, assume E τ ∼ p πz ( τ | s 0 )
Revised: To show the other direction, assume E τ ∼ p πz ( τ | s 0 )

Original: [ D ( I ( τ ) , z )] = 0 .
Revised: [ D ( I ( τ ) , z )] = 0 .

Original: Then, p π z ( τ | s 0 ) > 0 implies D ( I ( τ ) , z ) = 0 since distance D ( · , · ) is non-negative, which implies566 that when a trajectory has a non-zero probability, then I ( τ ) =
Revised: Then, p π z ( τ | s 0 ) > 0 implies D ( I ( τ ) , z ) = 0 since distance D ( · , · ) is non-negative, which implies that when a trajectory has a non-zero probability, then I ( τ ) =

Original: z .
Revised: z .

Original: Therefore, p π z ( τ | s 0 )
Revised: Therefore, p π z ( τ | s 0 )

Original: =
Revised: =

Original: p π z ( τ | s 0 , I ( τ ) =
Revised: p π z ( τ | s 0 , I ( τ ) =

Original: z ) , which finally implies that p π z ( s t +1 | s t , a t ) =568 p π z ( s t +1 | s t , a t , I ( τ ) =
Revised: z ) , which finally implies that p π z ( s t +1 | s t , a t ) = p π z ( s t +1 | s t , a t , I ( τ ) =

Original: z ) .
Revised: z ) .

Original: A.2 Benchmark Task DetailsA.2.1 An Illustrative Gambling Environment
Revised: A.2 Benchmark Task Details A.2.1 An Illustrative Gambling Environment

Original: To clearly illustrate the issue with prior approaches when conditioning on outcome variables that arecorrelated with environment stochasticity, we run our approach on a simple gambling environment.
Revised: To clearly illustrate the issue with prior approaches when conditioning on outcome variables that are correlated with environment stochasticity, we run our approach on a simple gambling environment.

Original: This environment, illustrated in Figure 1, has three actions: one which will result in the agent gainingone reward, and two gambling actions where the agent could receive either positive or negativereward.
Revised: This environment, illustrated in Figure 1, has three actions: one which will result in the agent gaining one reward, and two gambling actions where the agent could receive either positive or negative reward.

Original: A.2.2 Multi-Agent Game: Connect Four
Revised: A.2.2 Multi-Agent Game: Connect Four

Original: Connect Four is a popular two-player board game where players alternate in placing tiles in the hopeto be the first to get four in a row.
Revised: Connect Four is a popular two-player board game where players alternate in placing tiles in the hope to be the first to get four in a row.

Original: In this task, we consider a single agent version of Connect Fourwhere the opponent is fixed to be a stochastic agent.
Revised: In this task, we consider a single agent version of Connect Four where the opponent is fixed to be a stochastic agent.

Original: This is a realistic setting, since in the real world,the single greatest source of stochasticity will likely be other agents that the agent is interacting with.
Revised: This is a realistic setting, since in the real world, the single greatest source of stochasticity will likely be other agents that the agent is interacting with.

Original: An ideal agent will need to take this into account to make optimal decisions.
Revised: An ideal agent will need to take this into account to make optimal decisions.

Original: Since Connect Four can be optimally solved with search techniques, we set the opposing agent to beoptimal 1 , with a small chance that it won’t place a piece in the rightmost column when it is optimalto do so.
Revised: Since Connect Four can be optimally solved with search techniques, we set the opposing agent to be optimal 1 , with a small chance that it won’t place a piece in the rightmost column when it is optimal to do so.

Original: This creates an MDP with two ways to win: first, the agent can play optimally to the endof the game to guarantee a win (the first player always can win); second, to win quickly, the agentcan place four pieces on the rightmost column with the hope that the opponent will not block (whichhappens with a low probability).
Revised: This creates an MDP with two ways to win: first, the agent can play optimally to the end of the game to guarantee a win (the first player always can win); second, to win quickly, the agent can place four pieces on the rightmost column with the hope that the opponent will not block (which happens with a low probability).

Original: The offline dataset for this task is generated by using an epsilon-optimal agent with 50% probabilityand an exploiter agent that only places pieces on the right with 50% probability.
Revised: The offline dataset for this task is generated by using an epsilon-optimal agent with 50% probability and an exploiter agent that only places pieces on the right with 50% probability.

Original: [Figure]
Revised: [Figure]

Original: A.2.3 Stochastic Planning in
Revised: A.2.3 Stochastic Planning in

Original: 2048 [23] is a single player puzzle game where identical tiles are combined in order to build up tilesrepresenting different powers of two.
Revised: 2048 [23] is a single player puzzle game where identical tiles are combined in order to build up tiles representing different powers of two.

Original: With each move, a new tile randomly appears on the board, andthe game ends when no moves are available.
Revised: With each move, a new tile randomly appears on the board, and the game ends when no moves are available.

Original: A strong 2048 agent will consider the different possibleplaces new tiles will appear in order to maximize the potential for combining tiles in the future.
Revised: A strong 2048 agent will consider the different possible places new tiles will appear in order to maximize the potential for combining tiles in the future.

Original: Since the vanilla version of 2048 can require billions of steps to solve with reinforcement learning[34], we modified the game 2 by terminating the episode when a 128 tile is created.
Revised: Since the vanilla version of 2048 can require billions of steps to solve with reinforcement learning [34], we modified the game 2 by terminating the episode when a 128 tile is created.

Original: The agent getsone reward for successfully combining tiles in order to reach 128 and zero reward otherwise.
Revised: The agent gets one reward for successfully combining tiles in order to reach 128 and zero reward otherwise.

Original: The offline dataset for this task is generated using a mixture of trajectories from an agent trained with
Revised: The offline dataset for this task is generated using a mixture of trajectories from an agent trained with

Original: PPO [24] using the implementation in Stable Baselines 3 [35] and a random policy.
Revised: PPO [24] using the implementation in Stable Baselines 3 [35] and a random policy.

Original: A.3 Stochasticity-
Revised: A.3 Stochasticity-Independent Representations Perform Better

Original: Independent Representations Perform Better

Original: In order to verify that the degree to which the learned trajectory representations (cluster assignments)are independent of environment stochasticity affects the empirical performance of an agent, we setout to empirically measure this.
Revised: In order to verify that the degree to which the learned trajectory representations (cluster assignments) are independent of environment stochasticity affects the empirical performance of an agent, we set out to empirically measure this.

Original: We did so by training several agents in the gambling environmentwith different hyperparameters (specifically β adv ) and measured their performance.
Revised: We did so by training several agents in the gambling environment with different hyperparameters (specifically β adv ) and measured their performance.

Original: Results are shownin Figure 7.
Revised: Results are shown in Figure 7.

Original: Clustering assignments that enabled the dynamics model to predict the next state moreaccurately than should be possible lead to agents that performed worse.[Figure]
Revised: Clustering assignments that enabled the dynamics model to predict the next state more accurately than should be possible lead to agents that performed worse.

Original: A.4 Training Details 
Revised: A.4 Training Details A.4.1 Hyperparameters

Original: [Figure] A.4.1 Hyperparameters

Original: The most important hyperparameters to tune for our method are the tradeoff between reconstructingactions and removing dynamics information from the clusters, controlled by β act , and the numberof clusters, controlled by rep_size and rep_groups .
Revised: The most important hyperparameters to tune for our method are the tradeoff between reconstructing actions and removing dynamics information from the clusters, controlled by β act , and the number of clusters, controlled by rep_size and rep_groups .

Original: The trajectory representation (i.e. clusterassignment) is formed by sampling from rep_groups categorical variables of dimension rep_size/ rep_groups .
Revised: The trajectory representation (i.e. cluster assignment) is formed by sampling from rep_groups categorical variables of dimension rep_size / rep_groups .

Original: These values were tuned per environment using a simple grid search.
Revised: These values were tuned per environment using a simple grid search.

Original: Specifichyperparameter values for each environment can be found at Table 1.
Revised: Specific hyperparameter values for each environment can be found at Table 1.

Original: A.4.2 Computation
Revised: A.4.2 Computation

Original: Our experiments were run on T4 GPUs and running our algorithm took only around 1 hour per seed.
Revised: Our experiments were run on T4 GPUs and running our algorithm took only around 1 hour per seed.

Original: We used PyTorch [38] and experiments were tracked using Weights and Biases [39].
Revised: We used PyTorch [38] and experiments were tracked using Weights and Biases [39].

Original: Shared
Revised: Shared

Original: [Table] [Table]
Revised: [Table] [Table]

Original: [Table]
Revised: [Table]

Original: MLP
Revised: clustering_ [Table]

Original: [Table]
Revised: [Table]

Original: policy.hidden_layers clustering_model.
Revised: clustering_model.lstm

Original: [Table]
Revised: [Table]

Original: _layers clustering_model.lstm_
Revised: clustering_model

Original: [Table]
Revised: [Table]

Revised: lstm

Original: [Table]
Revised: [Table]

Original: clustering_model.lstm_layers action_predictor.hidden_layers return_predictor.hidden_layers transition_predictor.hidden_layers
Revised: _predictor.

Original: [Table]
Revised: [Table]

Original: activation_fn ReLU optimizer
Revised: return_predictor.

Original: [Table]
Revised: [Table]

Original: [Table]

Original: [Table]
Revised: [Table]

Original: Four
Revised: transition_predictor

Original: [Table] [Table]
Revised: [Table] [Table]

Original: For return-conditioned RvS [2], we used the same code as for ESPER, using returns rather thanthe learned estimated returns and otherwise using the same hyperparameters found in Table 1.
Revised: For return-conditioned RvS [2], we used the same code as for ESPER, using returns rather than the learned estimated returns and otherwise using the same hyperparameters found in Table 1.

Original: For
Revised: For

Original: Conservative Q-Learning (CQL) [4], we used the default implementation from d3rlpy, an offline deep
Revised: Conservative Q-Learning (CQL) [4], we used the default implementation from d3rlpy, an offline deep

Original: RL library [40].
Revised: RL library [40].

Original: A.6 ESPER Pseudocode
Revised: A.6 ESPER Pseudocode

Original: We provide pseudocode for ESPER in algorithm 1 and detailed pseudocode (roughly following thesyntax used in PyTorch [38]) for the clustering step of ESPER in algorithm2.
Revised: We provide pseudocode for ESPER in algorithm 1 and detailed pseudocode (roughly following the syntax used in PyTorch [38]) for the clustering step of ESPER in algorithm

Original: # s, a, seq_len: states, actions, sequence lengths (in case of early termination)
Revised: # s, a, seq_len: states, actions, sequence lengths (in case of early termination)

Original: # Models:
Revised: # Models:

Original: # encoder_mlp - MLP # temporal_encoder - LSTM # rep_mlp - MLP # act_mlp - MLP
Revised: # encoder_mlp - MLP # temporal_encoder - LSTM # rep_mlp - MLP # act_mlp - MLP

Original: # dynamics_mlp - MLP def ClusterAssignments(s, a): bsz, t =
Revised: # dynamics_mlp - MLP def ClusterAssignments(s, a): bsz, t =

Original: s.shape[:2] x = torch.cat((s, a), dim=-1).view(bsz, t, -1) x = torch.flip(x, dims=[1]) x = torch.encoder_mlp(x) x, hidden = temporal_encoder(x, init_hidden())
Revised: s.shape[:2] x = torch.cat((s, a), dim=-1).view(bsz, t, -1) x = torch.flip(x, dims=[1]) x = torch.encoder_mlp(x) x, hidden = temporal_encoder(x, init_hidden())

Original: x = torch.flip(x, dims=[1]) x = rep_mlp(x)
Revised: x = torch.flip(x, dims=[1]) x = rep_mlp(x)

Original: cluster_assignments = F.gumbel_softmax(x) return cluster_assignments def ActionPredictor(s, cluster_assignments): # cluster_assignments.shape =
Revised: cluster_assignments = F.gumbel_softmax(x) return cluster_assignments def ActionPredictor(s, cluster_assignments): # cluster_assignments.shape =

Original: [bsz, t, -1]
Revised: [bsz, t, -1]

Original: # For a timestep t, we sample a cluster assignment for a timestep # from 0, ..., t randomly past_assignments =
Revised: # For a timestep t, we sample a cluster assignment for a timestep # from 0, ..., t randomly past_assignments =

Original: sample_past_assignments(cluster_assignments) x = torch.cat((s, past_assignments), dim=-1) pred_next_action = act_mlp(x) return pred_next_action
Revised: sample_past_assignments(cluster_assignments) x = torch.cat((s, past_assignments), dim=-1) pred_next_action = act_mlp(x) return pred_next_action

Original: def TransitionPredictor(s, a, cluster_assignments): # cluster_assignments.shape =
Revised: def TransitionPredictor(s, a, cluster_assignments): # cluster_assignments.shape =

Original: [bsz, t, -1]
Revised: [bsz, t, -1]

Original: # For a timestep t, we sample a cluster assignment for a timestep # from 0, ..., t randomly past_assignments =
Revised: # For a timestep t, we sample a cluster assignment for a timestep # from 0, ..., t randomly past_assignments =

Original: sample_past_assignments(cluster_assignments) x = torch.cat((s, a, past_assignments), dim=-1) pred_next_s
Revised: sample_past_assignments(cluster_assignments) x = torch.cat((s, a, past_assignments), dim=-1) pred_next_s

Original: =
Revised: =

Original: dynamics_mlp(x) return pred_next_s
Revised: dynamics_mlp(x) return pred_next_s

Original: # training loop , a in dataloader:
Revised: # training loop for (s, a) in dataloader:

Revised: # get cluster assignments for the trajectories cluster_assignments =

Revised: ClusterAssignments(s, a) # predict actions based on

Original: [Table]
Revised: [Table]