You Can’t Count on Luck: Why Decision Transformers and RvS Fail in Stochastic Environments
Authors: Keiran Paster
Sheila A. McIlraith
Jimmy Ba
Venue: NeurIPS 2022 Accept (2022)
Keywords: representation learning
reinforcement learning
model-based reinforcement learning
decision transformer

Abstract: Recently, methods such as Decision Transformer that reduce reinforcement learning to a prediction task and solve it via supervised learning (RvS) have become popular due to their simplicity, robustness to hyperparameters, and strong overall performance on offline RL tasks. 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. In this work, we describe the limitations of RvS approaches in stochastic environments and propose a solution. Rather than simply conditioning on returns, as is standard practice, our proposed method, ESPER, conditions on learned average returns which are independent from environment stochasticity. Doing so allows ESPER to achieve strong alignment between target return and expected performance in real environments. We demonstrate this in several challenging stochastic offline-RL tasks including the challenging puzzle game 2048, and Connect Four playing against a stochastic opponent. In all tested domains, ESPER achieves significantly better alignment between the target return and achieved return than simply conditioning on returns. ESPER also achieves higher maximum performance than even the value-based baselines.


You Can't Count on Luck: Why Decision Transformers and RvS Fail in Stochastic Environments
Keiran Paster
Sheila A. Mcllraith & Jimmy Ba Department of Computer Science University of Toronto, Vector Institute {sheila, jba toroots edu toronto.....
Department of Computer Science University of Toronto, Vector Institute keirpecs. toronto. edu
Abstract
Recently, methods such Decision Transformer [I] that reduce reinforcement learning to a prediction task and solve it via supervised learning (RvS) [2] have become popular due to their simplicity, robustness hyperparameters,, and strong overall performance on offline RL tasks. However, simply conditioning a probabilis- tie model desired return and taking the predicted action fail dramatically in stochastic environments since trajectories that result in a return may have only achieved that return due to luck. In this work, describe the limitations of RvS approaches in stochastic environments and propose solution. Rather than simply conditioning on the return of 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. Doing so allows ESPER to achieve strong alignment between target return and expected performance in real environments. We demonstrate this in several challenging stochastic offline-RL RL tasks including the challenging puzzle game 2048, and Con- neet Four playing against stochastic opponent. In all tested domains, ESPER achieves significantly better alignment between the target return and achieved return than simply conditioning on returns. ESPER also achieves higher maximum performance than even value-based baselines.
Introduction
Offline reinforcement learning (RI.) [3, 1, 4, 5] is a promising approach to train agents without requiring online experience an environment, which is desirable when online experience is expensive or when offline experience is abundant. A recent trend in offline RL has been
simple approaches that do RL via Supervised Learning (RvS) (e.g.. [1, 2, 6-8]) rather than typical value-based approaches. RyS algorithms such Decision Transformer [1] train model to predict action based on the current state and an outcome such as a desired future return. These agents ask the question "if / assume the desired outcome will happen, in my experience what action do I typically take next." These methods are popular due to their simplicity. strong performance on offline benchmark tasks [1, 2]. and similarity large generative models (c.g., [9-12]) that continue to show stronger performance on more tasks when training larger models on more data.
However, we will show, methods that condition outcomes such as return can make incorrect decisions in stochastic environments regardless scale or the amount of data they are trained on. This is because implicitly these methods assume that actions that end up achieving particular goal are optimal for achieving that goal. This assumption not true in stochastic environments, where it is possible that the actions taken in the trajectory were actually sub-optimal and that the outcome was only achieved due to lucky environment transitions. For example, consider the gambling
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

environment in Figure 1. Though there may be many episodes in which an agent gets positive return from gambling (ag or a1), gambling is sub-optimal since results in a negative return in expectation while d2 always results in positive return. Since RvS takes all of these trajectories as expert examples of how to achieve the goal, RvS will act sub-optimally.
This limitation of RyS 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. Locomotion tasks in MuJoCo [15] and Atari games in the Arcade Learning Environment [16] are two such examples. 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]. 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, In their current form, approaches such as Decision Transformer [1] are likely to behave in unexpected ways in such scenarios.
One way to view why RyS 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 a negative reward. 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. 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. 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. While trajectory return not such an outcome, we show that the expected return of behavior shown in trajectory is, and how to learn such value. The contributions of this work are as follows:
We show that RyS-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.
We propose method to learn anirronmennt-tcchaaaiiittddddeee trajectories using adversarial learning, on top of which we label each trajectory with the average return for trajectories with this representation.
representations of
We introduce several stochastic offline RL benchmark tasks and show that while Rv$ agents consistently underperform on the conditioned returns, our approach (ESPER, short for hith--aacicccc---ieeeeeeeeee
representations) achieves significantly stronger alignment between target return and actual expected performance. 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].
2 Approach
2.1 Problem Setup
We model the environment as an MDP, defined as the tuple A,, R, 7). S is set of states; A is a set of actions; the transition probabilities T S x x S + [0, 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 x A - IR maps a state-action transition to a real number; and 0 < 7 < is the discount factor, which controls how much an agent should prefer rewards sooner rather than later. The performance a policy, defined by π (als). is typically measured using the cumulative discounted return Σ, 'rr that the policy achieves in an episode.
Central to our work is algorithms that do reinforcement learning via supervised learning (RvS) such as Decision Transformer [1]. These approaches train model using supervised learning on a dataset of trajectories to predict (a
(a)s. s. R) - the probability of next action conditioned on the current state and the agent getting cumulative discounted return R. Σ, 7'r,. At evaluation time, the model is conditioned on a desired target return and the agent takes the actions predicted by the model.
To understand the problem with RvS in stochastic environments, we build off of a general framework where a goal presented to a policy, the policy acts in the environment, and the agent is scored on how well the goal was achieved as described in Furuta al. [18]. Mathematically, the agent's

---

80
33%
33%
33%
Conditioned Return
E [r]
Actions
00
(1)
a22
-5
-5.0
{au}
r = -6
-2.5
{ai}
50%
50%
50%
50%
-0.17
{al, (2)
5
-5.0
{ao)
r
6
Left: A simple gambling environment with three actions where return conditiond algo- rithms such as Decision Transformer [1] will fail, even with infinite data. The optimal action, d2 will always grant the agent reward while gambling (ag and 01) give stochastic amount of reward. The numbers (33%) above each action represent the data collection policy. Right: The second column represents the performance of a policy behaviorally cloned from trajectories achieving the reward in the first column. For example, conditioned on the agent receiving r = 1, a third of trajectories take ai and two-thirds take 02. Averaging the returns -2.5 1/3 + 1 2/3 -0.17. No matter how much data the model is trained on, when conditioned on receiving a reward of 1, it will always gamble and take aj some of the time, rather than just taking ag, which guarantees the reward.
Figure 1:
objective can be expressed as minimizing some distance metric D(:-)) between a target goal z and the information statistics ((r) of the trajectory generated by conditional policy (a]8.
min
[D(](T),=)
(2.1)
where p(2) the goal distribution, Pr, (T) is the trajectory distribution ($1, a1, 82, ag..... obtained by rumning rollouts in an MDP following policy T(ass,)) 2). In this section, we refer to the goal presented to an agent as and the function that calculates which goal was achieved particular trajectory v as its statistics I(T). As discussed in Furuta et al. [18], this framework includes number of popular algorithms, including Rv$ algorithms that condition on target returns (e.g.. [1, 2, 8, 6)) or goals (e.g.. [7, 19]). For example, in Decision Transformer [I] cumulative discounted reward is used as the trajectory statistics I ((T) = E. y'r; and a reasonable choice for D(..)) could be the squared error between the target and achieved return.
Furuta et al. [18] justify using the supervised learning approach (RvS) to minimize Equation 2.1 by claiming that trajectories with statistics [(T) act as optimal examples of how to act to achieve D(I()). =) = 0 for an agent whose goal is to achieve z [(T). However, in stochastic environments, the actions taken in these trajectories may not be optimal since trajectory may achieve goal due to uncontrollable randomness in the environment rather than from the agent's own actions. In fact, as discussed in Paster et al. [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.
In this section, propose sidestep this issue by limiting the types of statistics I(r) interested in to those that are not affected by uncontrollable randomness in the environment. We show first that RvS policies trained 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. Importantly, we argue that limiting the form of trajectory statistics 7(r) in this way is not prohibitive and that desirable quantities can be transformed fit this property an intuitive way (e.g., return turns into expected return when optimized to be independent of environment stochasticity). Finally, we propose an approach for automatically finding such trajectory statistics using adversarial learning.
2.2 Stochasticity Independent Representations
Definition 2.1 (Consistently Achievable). A goal z is consistently achievable from state so under O.
policy x (als, 2) Eyrpt, (rleo) [P(())) =)] (1(T))
natural question is under which cireumstances will the RyS approach actually minimize Equa- tion 2.1 to zero? Clearly the approach works empirically deterministic environments [1, 2] since
3

---

Phase 1:
Phase 2:
Phose 3: Train RvS Agent
Adversarial Clustering
Estimate Average Returns
I(T)
R=5
D
assign Rogectory
chuster
A=
R=1
change oluster assignment to hurt prediction
predict dyramics conditione on cluster essignment
T(acs Rs 3)
1 [(7))
Figure 2: ESPER learns policy that conditions on a desired expected return. In phase I of the algorithm, a function is learned using adversarial learning that assigns each trajectory in the dataset to a cluster such that the cluster assignments do not contain information about the stochastic outcomes of the environment that can help a dynamics model cheat to predict next states. In phase 2, the average return in each cluster is calculated. In phase 3, an RvS agent is trained to predict the next action given the current state and the estimated average return.
simply replaying an action sequence that achieved goal once will achieve again. While not all statistics of trajectory will lead to consistently achievable goals under policies trained with RvS, the supervised learning approach will minimize Equation 2.1 if ((T is independent from environment stochasticity.
Theorem 2.1. Let TD be data collecting policy used to gather data used to train an RvS policy,
and assume this policy is trained such that π [ dg ......... St; 2) Pro (a, 80; ag....... ... St, [(r)
z). Then, for any goal 2 such that Pap (I( = z so) >> O, goal is consistently achievable iff Prp (5t 50: dg...., at-1) Pr (Se|S0, ........ 01-1, I (TT) 2).
Proof. [Sketch] The forward direction of the theorem is proven by rewriting Ey-p., (T(s6) [D((()), =)] to be taken over Pa, (r|s0. (T) = z) using the independence assumption combined with the form of the policy. For the other direction and the full proof, refer Appendix A.1.
Remarks. The above theorem has two major implications. 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. For example, the gambling environment discussed in Figure I, 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 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 high reward consistently.
Theorem 2.1 also gives a solution ensure asymptotically that RvS policies will consistently achieve goals: use trajectory statistics that are independent from the stochasticity in environment transitions. This requires the statistic function /(T) to return the same value for different trajectories where only difference is in the environment transitions, making 7(T) essentially act cluster assignment where each trajectory within cluster has the same value of (T) and has environment transitions sampled from the true distribution p(s++1|ss st, a.). These clusters can be thought of as datasets each generated by distinct policies the same environment. For example, in the gambling environment each cluster could include trajectories that all took the same action.
Ultimately, we want use target average return to control an RyS agent. Luckily, given any I(T) that satisfies Theorem 2. we can transform it by using the average returns of trajectories within each cluster as the statistic instead. In the next section, we describe a practical algorithm for learning an T(r) that satisfies Theorem 2.1, finding the average return of trajectories with the same value of T(r) (i.e. in the same "cluster"), and training an RvS policy conditioned on these estimated average returns.

---

2.3 Learning Stochasticiy-Indeeeneel Representations
In this section we propose method for learning such trajectory representations, which we call ESPER, short for return as the statistic (Tr), which is often dependent on environment stochasticity and therefore may result in unexpected behavior, ESPER uses neural network to learn statisties that are independent from environment stochasticity. does so three phases: first, uses an auto-encoder framework to learn discrete representation / (() acting as cluster assignment for each trajectory along with an adversarial loss that changes the representation in order to hurt the predictions a dynamics model; second, model learns 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.
representations, Rather than use trajectory
For some trajectory ($1, 01. r1, 82, 02 .2....). .), ESPER trains the following parameterized models in addition to vanilla RvS policy me:
Clustering model:
[(T) ~ (o(((T) 0 ~ po(arlse,
Action predictor:
Return predictor:
it = fx(((T))
Transition predictor:
St+1 ~ Po 150 ao...., St. ocs (T))
Adversarial clustering, Here, we will focus on learning representation in the form of clustering assignments. The goal is to learn clustering assignment per trajectory I(r) to minimize the discrepancy between the true environment transition p( [st, ax) and the estimated transition Po (8(++1// (8t; at, [(T)) from the replay buffer when conditioned on [(T). Such a representation will satisfy Theorem 2.1. 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 ((T) does not improve the performance of the dynamics model:
((0) B((r))
Bus log po (s, I(T)) + Paty log (o (sp ag. ................................ St, Ce. [(T))
policy reconstruction
adversarial loss
(2.2)
[(O) (RG))
log P& (8++1 1 |8: oo....,, St, at, ((T))
(2.3)
dynamics prediction
In this step, the dynamics prediction loss is predicting the next state given the current state, action, and cluster assignment [T) while the adversarial loss tries to change I(r) in order to hurt this prediction. Since constant representation minimizes this loss, we use policy reconstruction loss to encourage ((T) to contain information about the policy used to generate the trajectory and encourage the formation of more than one cluster. Aai and Buly are hyperparameters to balance the strength of this policy reconstruction loss with the adversarial loss.
Estimate cluster average returns. After clustering, we learn to predict discounted returns It = E, 2'r, for each cluster. The return predictor parameterized by ve is trained using the following loss:
(2.4) The above losses are described such that the model produces a single estimated return per trajectory. However, practice it can be extended produce a value for each time-step a trajectory by using suffixes of trajectories Ti, = (at, St, Tt....,, ., ar, ST. sT.rr) to generate return predictions RG for each time-step.
-( = ((())) ((
Training policy predicted returns. After learning expected future returns for each step in trajectory, we use the dataset of (st, a,, R₂) triples learn policy me (ar|s,, Rτ) parameterized by
Note that in practice, we condition the return predictor not only on the cluster assignment for the trajectory, but also on the first state and action. This improves performance while still satisfying Theorem 2.1.
5

---

Task
Return- Cetum-Conditioned RvS (DT)
CQL
ESPER (Ours)
Gambling
-0.02 (0.24)
1.0 (0.0)
1.0 (0.0)
Connect Four
0,8 (0.07)
0.61 (0.05)
0.99 (0.03)
2048
0.57 (0.05)
0.7 (0.09)
0.81 (0.06)
Figure 3: Comparing the maximum performance of RvS with return conditioning, RyS with ESPER conditioning (ours), and CQL [4], strong value-based offline-RL baseline. The maximum possible expected return for policy is 1.0 on all tasks. Since all our tasks are stochastic, return conditioning fails to learn an optimal policy while ESPER performs optimally or near optimally. CQL performs well on the gambling task, but cannot achieve the same level of performance as ESPER on the more complicated tasks. Numbers in parenthesis are standard deviations over 3 seeds.
& by predicting the action from the current state and estimated expected return by minimizing the following loss just in prior RvS works [2, 1]:
I(g) = E BLND log me ( Ri))
(2.5)
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, and transition predictor are all implemented using MLPs. In practice, we found that the dynamics predictor works well even without conditioning on the entire trajectory history. We use a transformer policy as in Decision Transformer [1] (see section 3.2), although we find that using an MLP is sufficient for strong performance in our benchmark tasks (see appendix A.3). Our clustering outputs a discrete value sampled from a categorical distribution and we use gumbel-softmax [21, 22] in order to backpropagate gradients through it. We use normal distributions with unit variances for the transition predictor. Code for our implementation available at https://sites.. google.coom/ve von/view/esppr- paper. See appendix 7 for pscudocode for the adversarial clustering step of ESPER. See appendix A.5 for more information about the implementation, including specific hyperparameters for our environments as well as best practices.
3 Experiments
We designed our experiments to answer the following questions: (i) How well does return eeturn-conditioned RvS work on stochastic problems where returns are correlated with environment stochasticity? (ii) Does ESPER improve performance, both in terms of maximum achievable performance as well as in correlation between the target and achieved returns? (iii) Does the degree to which the learned representations in ESPER independent from environment stochasticity measurably affect the performance of the learned agent?
3.1 Stochastic Benchmark Tasks
Prior offline RL methods including Decision Transformer have been tested primarily deterministic or near-deterministic environments such as continuous control tasks from the D4RL. dataset [13] and Atari games [16]. To get a better idea of how these methods and our approach work on realistic, stochastic environments, we created three new benchmark tasks. Since we are evaluating in an offline setting, each task consists of stochastic environment as well as one or several data collection policies in order to gather the offline dataset.
Gambling. An illustrative gambling environment described in Figure 1 with only 3 actions. The offline dataset for this task is 100k steps collected with a random policy.
Connect Four A game of Connect Four, a tile-based game where players compete to get four-in-a- row, against an opponent that does not place a tile the rightmost column with probability of 20%. The agent gets reward for winning, 0 for draw, and -1 for losing. The offline dataset for this task

---

Connect Four
20%
80%
1.0
Return Coud.
ESPER (Quss)
0.5
CQL
Peffrraace
0.0
-0.5
-1.0
11000
-0.75
-0.50
-0.25
0.00
0.25
0.50
0.75
1.00
Turgel Return
Figure 4: We consider Connect Four with a stochastic opponent that fails to place a piece in the rightmost column with 20% chance. In this environment, the agent gets reward of 1 for winning, 0 for a draw, and -1 for losing. The histogram represents the distribution of returns each method is trained and out-of-distribution regions for ESPER are shaded. On in-distribution returns, ESPER achieves performance (y-axis) closer to the target performance (x-axis) than return conditioning.
IM steps collected using a mixture of an c-greedy policy and policy that always places tiles in the rightmost column.
2048. simplified version of 2048 [23], a sliding puzzle game where alike numbers can combined to create tiles of larger value, where the agent gets reward of by creating 128 tile and gets no reward otherwise. 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].
Detailed information about each environment can be found in Appendix A.4.
3.2 Baselines
We compare our method primarily against Decision Transformer [1], which trains a transformer [25] to predict the next action conditioned a sequence of previous states, actions, and return-to-go targets. Our method uses the same model and training code, except rather than condition on target returns condition expected returns learned using ESPER. While conduct our experiments using transformers due to the popularity of Decision Transformer [1]. we note that our observations still hold when using MLP policies in Emmons et al. [2] while using significantly less compute and getting similar performance (see Appendix A.3). 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.
3.3 Performance in Stochastic Environments
When comparing the performance of our approach, ESPER, with return-conditioned RyS, we found that ESPER consistently achieves stronger alignment between target return and average perfor- mance while also achieving higher maximum level of performance when tuning the target return. As shown in Figure 5a, the conuiincoddiionee baseline cannot achieve the maximum performance of 1 while ESPER learns behaviors corresponding to average performances ranging from -5 to 1. In Connect Four (Figure 4), despite seeing many examples winning trajectories, an RvS agent cannot achieve more than 0.2 average return (corresponding to win-rate of 60%) while ESPER can both win and lose with perfect accuracy depending on the target performance. In 2048 (Figure 5b), RyS 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. ESPER learns many modes of behavior, and can be controlled to win anywhere from 30% to 80% of the time.
Additionally, Figure shows that this lack of performance for the return-condittoned agent is not due to a lack of data. In the figure, going from 5% data to 100% (5 million frames) does not affect the performance of RvS while the performance ESPER gets closer to matching the target return with more data. Finally, we compared the maximum performance of RvS and ESPER with the state-of-the-an value-based offline RI. algorithm, Conservative Q-Learning (CQL.) [4]. As shown in Figure 3, ESPER and CQL achieve perfect performance the gambling task while RyS cannot get

---

Gambling
2048
1.0
Return Cand.
Return Cand.
ESPER (Oum))
ESPER (Ouns)
CQL
0.6
CQL
11
-2-
0.6
-3
miniom
0.4
expert policy
0.2
-6
-10
0.5
Target Return
Target Return
(b)
Figure 5: The histogram represents the distribution of returns each method is trained Left: In the illustrative gambling environment, ESPER can achieve a range performances from -5 to 1, while the performance (y-axis) of the retum-conditioned agent is not aligned with the target return (x-axis). Right: In our modified 2048 task, the agent receives a reward of 1 for creating tile value 128 and 0 otherwise. While eeturn-conditione agent does not achieve high level of performance, the ESPER agent disentangles actions from environment stochasticity and able achieve performance close to the target return. Interestingly, the distribution of expected returns learned by ESPER match the performance-leve of the data-collection policies.
Scaling 2018
Performance va. Dynamies Loes Gambling
Return Cond.
1.0
ESPER (Oum)
0.5
Prohabliry
Performance
0.0
-1.5
12.5
25
50
0.00
0.01
0.02
0.03
0.04
0.05
0.06
0.07
(10*)
Dynamics Losa
Figure 6: Left: While return crurn-conditioned agent does not improve its performance with more data, ESPER achieves performance (y-axis) closer to the target return (x-axis) when trained on more data. 100% data usage 5 million frames. Right: We trained several agents with different hyperparameters on the gambling task. Notably, agents with trajectory statistics that enabled the dynamics model to "cheat" and get a low loss performed worse than those with representations independent from environment stochasticity, empirically confirming Theorem 2.1.
positive reward. In 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.
3.4 Learned Representations and Behaviors
Rather than simply condition on returns, ESPER conditions on learned expected return values. 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. 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. Unlike RvS, when an agent is trained on particular expected return value, the actual average performance of the agent is often close to this value. This is useful for offline settings where tuning the target return an online environment is not feasible, since with ESPER one have confidence in which levels of performance the model will be able to achieve. In contrast, with return-conen condiiiiiii. the performance of the return-conditioned policy may correlate with the returns which it was trained.
Finally, we empirically measured the relationship between the independence of learned trajectory statistics and the performance of the agent. As shown in Figure 6, trained several agents with

---

different hyperparameters on the gambling task and measured the dynamics loss (Equation 2.3) and performance of each agent. Indeed. agents with trajectory statisties that enable the dy- namics model to "cheat" and get a low loss performed worse than those with representations independent from environment stochasticity, empirically confirming Theorem 2.1.
4 Related Work
Offline Reinforcement Learning. Offline reinforcement learning is framework where logged data is used to learn behaviors [3]. 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. The main challenge in the offline setting is that offinn RL agents cannot collect more data reduce their uncertainty and therefore offlin RL approaches often have some type of value pessimism or policy constraints keep the learned policy within the support of the offline dataset. Conservative 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. Other methods rely on ensembles (e.g., [5]) or policy regularization (e.g.. (26-28]). Several standardized benchmarks have emerged for testing offiine RL agents, including D4RL [13] and RL Unplugged [14]. 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 wear-deterministicc [4].
Reinforcement Learning via Supervised Learning. At the core of our work is the idea of reducing reinforcement learning to a prediction problem and solving it via supervised learning. This idea was first proposed by Schmidhuber [6] and Kumar et al. [8], who proposed to learn behaviors by predicting actions conditioned the same approach to solve goal-conditioned tasks and Paster et al. [19] proposed to predict multi-step action sequences, showing connection between sequence modeling and optimal open-loop planning. Decision Transformer [1] proposed to use a transformer [25] to condition on trajectory history when predicting actions and tested the approach on offline RL rather than online, achieving results competi- tive with value-based offline RL approaches. Several followup extensions to Decision Transformer were explored, including using non-return goals [18] and using pretrained transformers [29]. In recent work, Emmons et al. [2] coined the term RvS (reinforcement learning via supervised learning) to describe such approaches. Notably found that a transformer not necessary to perform well on most tasks, showing that two-layer MLP with sufficient regularization can actually outperform transformers at fraction of the computational cost.
desired outcome such as return. Ghosh et al. [7] proposed to use
No prior RvS approach has been thoroughly evaluated in stochastic environments. Paster et al. [19] describes a counter-example where RvS agent in Ghosh al. [7] doesn't converge in stochastic environment, and proposes a solution. 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. Ortega et al. [30] give a high level view of how sequence modeling can be affected by delusions in various problem settings when not treating actions as interventions. Our contribution is a more precise characterization of the problem within the framework of RvS which relates the choice of goal to environmental stochasticity and directly evokes an efficient algorithm for using RvS in stochastic environments without the need for explicit causal inference or intervention by carefully choosing the goals on which the agents conditions.
Adversarial Learning. ESPER uses an adversarial loss to ensure that trajectory statistics are independent from environment stochasticity. Adversarial losses are also widely used in fairness (c.g., [31-33]), where it is used to prevent models from using sensitive attributes, and in generative modeling such as in GANs [34].
5 Conclusion
As models are scaled up in pursuit of better performance [10], it is important to acknowledge the limitations that cannot fixed by more scale and data alone. 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. We show why such errors occur and show that these approaches work if and only if the goals that they trained to achieve are independent from

---

environment stochasticity. We give practical algorithm that transforms problematic goals such as trajectory returns into expected returns that satisfy this condition. 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 maximum performance. 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.
Limitations. There are several limitations and opportunities for future work. First, our approach relies on learning a dynamics model. In environments where decently-strong dynamics model cannot be learned, the adversarial clustering in ESPER may fail. Second, ESPER relies on adversarial learning, which similarly to GANs [34], can have learning dynamics that are sensitive to hyperpa- rameter choices (for example learning rate). Additionally, just like in k-means clustering, there are many possible cluster assignments for a given offline dataset. The final performance of ESPER can be affected by variance in the adversarial clustering phase. Finally, 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.
Societal Impact. 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. However, we acknowledge that powerful automated decision making algorithms such as ESPER have the potential to make harmful decisions when trained with under-specified objectives or on biased data.
Acknowledgements
The authors thank Forest Yang, Harris Chan, Dami Choi, Joyce Yang, Sheng Jia, Nikita Dhawan, Jonah Philion, Michael Zhang, Qizhen Zhang, Lev McKinney, and Yongchao Zhou for insightful discussions and helpful feedback on drafts of the paper. We gratefully acknowledge funding from the Natural Sciences and Engineering Research Council of Canada (NSERC) [JB: 2020-06904), the Canada CIFAR AI Chairs Program, Microsoft Research, the Google Research Scholar Program and the Amazon Research Award. Resources used preparing this research were provided, part, by the Province of Ontario, the Government of Canada through CIFAR, and companies sponsoring the Vector Institute for Artificial Intelligence www.veectorinstiiuuee ai ai/partners).
References
[1] Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Michael Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision Transformer: Rein- forcement Learning via Sequence Modeling, In Mare' Aurelio Ranzato, Alina Beygelz- imer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan, editors, Ad- vances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual, pages 15084-15097, 2021. URL https: http://e/ceeeeeeeeee/eeee/ceseee www.ediiinnnnnnuiiis..eeo neurips. cttccc//ppaa http:/appr/////00 paaper2221/hash/ T148566220aab00077566111666666. Abstract. html.
[2] Scott Emmons, Benjamin Eysenbach, Ilya Kostrikov, and Sergey Levine. RvS: What is Essential for Offline RL via Supervised Learning? CoRR, abs/2112.10751,, 2021. URL
https://arxiv. wwteriiivrrr///oo abtp/aas//212/ 10751. arXiv: 2112.10751.
[3] Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline Reinforcement Learning: Tutorial, Review, and Perspectives on Open Problems. arXiv:2005. 00005006663/1..0 [cs, stat], November 2020. URL http: http///axvvvv./////////e//eaaavr org/ abs 01643. arXiv: 2005.01643.
[4] Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative Q-Learning for Offline Reinforcement Learning. In Hugo Larochelle, Marc' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien Lin, editors, Advances in Neural Information Processing Systems 33: Annual Conference Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual, 2020. URI. https: http:://///oeeegeeeeee newriieurpppp. cc/ paper/ naper/2020// hww.o//aahhaaeeeeeeeaaeeeeaaeeea http://2222201122222222222222212
Abstract. html.
10

---

[5] Rishabh Agarwal. Dale Schuurmans, and Mohammad Norouzi. An Optimistic Perspective on Offline Reinforcement Learning. In Proceedings the 37th International Conference on Machine Learning, ICMI. 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pages 104-114, PMLR, 2020. BtR.htttp//prroeeeddae press/v119/ wwtpw/////ggaaaaaa.aao html.
mlr.
[6] Jürgen Schmidhuber. Reinforcement Learning Upside Down: Don't Predict Rewards Just Map Them to Actions. CoRR, abs/ 1b1//199,.0287, 2019. URL http: //appiv///xii.////a org/ htt//bb///// 191.. 02875.
arXiv: 1912.02875.
[7] Dibya Ghosh, Abhishek Gupta, Ashwin Reddy, Justin Fu, Coline Manon Devin, Benjamin Eysenbach, and Sergey Levine. Learning to Reach Goals via Iterated Supervised Learning. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Aus- tria, May 3-7, 2021. OpenReview.net, 2021. URL https:://ppeereeeeieee http://///ppeeeeeeeeee//errrrm net/forummid= rALAOXo6yNJ.
[8] Aviral Kumar, Xue Bin Peng, and Sergey Levine. Reward- Conditioned Policies. CoRR. abs/1912. 13465, 2019. URL http://arxiv org/ aaaaab//1999 1912. 13465. arXiv: 1912.13463
[9] Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language Models are Unsupervised Multitask Learners. 2019.
[10]
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanic Subbiah, Jared Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert- Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language Models are Few-Shot Learn- ers. In Hugo Larochelle, Marc' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien Lin, editors, Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6- 12, 2020, virtual, 2020. URL https: / h/ppc//ee//cooo.oooo/oooo www.eediiinnnnuuunis...e. www.nnnrri.... 1457c0d6bfcb496744120
(amper/2020//1
Abstract html.
[11] Jordan Hoffmann. Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas. Lisa Anne Hendricks, Johannes Welbi, Aidan Clark, Tom Hennigan, Eric Noland, Katie Millican, George van den Driessche, Bogdan Damoc, Aurelia Guy, Simon Osindero, Karen Simonyan, Erich Elsen, Jack W. Rae, Oriol Vinyals, and Laurent Sifre. Training Compute- Optimal Large Language Models. CoRR, abs/2203. 15556, 2022. doi: arXiv: 2203.15556.
10.485550arXX..2223..555 15556.
[12] Jean-Baptiste Alayrac, Jeff Donahue, Pauline Luc, Antoine Miech, lain Barr, Yana Hasson, Karel Lene, Arthur Mensch, Katic Millican, Malcolm Reynolds, Roman Ring, Eliza Ruther- ford, Serkan Cabi, Tengda Han, Zhitao Gong, Sina Samangooci, Marianne Monteiro, Jacob Menick, Sebastian Borgeaud, Andrew Brock, Aida Nematzadeh, Sahand Sharifzadeh, Miko- laj Binkowski, Ricardo Barreira, Oriol Vinyals, Andrew Zisserman, and Karen Simonyan. Flamingo: Visual Language Model for Few-Shot Learning. CoRR, abs/2204.14198, 2022. doi: 10448550/axXxx,2244444999
arXiv: 1419..
[13] Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4RL: Datasets for Deep Data-Driven Reinforcement Learning. CoRR, abs/2004.07219, 2020. URL https: //aaxiv..rrg orgg abs/ baabe/20044 2004. 07219. arviv: 2004.07219.
[14] Çaglar Gülçehre, Ziyu Wang, Alexander Novikov, Thomas Paine, Sergio Gómez Colmenarejo, Konrad Zolna, Rishabh Agarwal, Josh Merel, Daniel J. Mankowitz, Cosmin Paduraru, Gabriel Dulac-Arnold, Jerry Li, Mohammad Norouzi, Matthew Hoffman, Nicolas Heess, and Nando de Freitas. RL Unplugged: A Collection of Benchmarks for Offline Reinforcement Learn- ing. In Hugo Larochelle, Mare' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien Lin, editors, Advunces Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6- 12, 2020, virtual, 2020. URL. https: https://proceedingss
wwwediiinnnnnneiiinneeeee www.ipsur.... httppscccppppeee2e..// /hash/ htt.d/daabbbbbaee-------aaeteet Abstract. html.
11

---

[15] Emanuel Todorov. Tom Erez, and Yuval Tassa. MuJoCo: A physics engine for model-based control. In 2012 IEEE/RSJ International Conference on Intelligent Robots and Systems, pages 5033, October 2012. doi: 10.1 10110000.20002.663660
ISSN: 215 2133-0866.
[16] Marc G. Bellemare, Yavar Naddaf, Joel Veness, and Michael Bowling, The Arcade Learning Environment: An Evaluation Platform for General Agents, J. Artif. Intell. Res., 47:253-279, 2013. doi: 10. 1/1.i33/.....
[17] Marlos C. Machado, Mare G. Bellemare, Erik Talvitie, Joel Veness, Matthew J. Hausknecht, and Michael Bowling. Revisiting the Arcade Learning Environment: Evaluation Protocols and Open Problems for General Agents. J. Artif. Intell. Res., 61: 61:523-5643 562, 2018. doi: 10. 10.1613/jair.566
[18] Hiroki Furuta, Yutaka Matsuo, and Shixiang Shane Gu. Generalized Decision Transformer for Offline Hindsight Information Matching. arXiv:2111.1036 [cs, stat], February 2022. URL http://arxiv.ooo org/ aabpg///b///a 2111 10364. arXiv: 2111.10364.
[19] Keiran Paster, Sheila A. Mcllraith, and Jimmy Ba. Planning from Pixels using Inverse Dynamies Models. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021, URL https:://pperrvieeeee net/
Barum?id-V6B:Bgkaall
[20] Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9(8): 1735-1780, 1997. Publisher: MIT Press.
[21] Eric Jang, Shixiang Gu, and Ben Poole. Categorical Reparameterization with Gumbel-Softmax. In 5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings. OpenReview.net, 2017. URL https:// openreview. wwt//eeeuoutttee iiruum?ddrEEEy8eee
[22] Chris J. Maddison, Andriy Mnih, and Yee Whye Teh. The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables. In 5th International Conference Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings. OpenReview.net, 2017. URL https: bectps::/openneeee httpp///rrrii.....oooooooree//ee
[23] Gabriele Cirulli. 2048, March 2014. URL http::/tp//p
[24] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal Policy Optimization Algorithms. CoRR, abs/1707.06347, 2017. URL http://arxiv org/ abs/1707 06347. arXiv: 1707.06347.
[25] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is All you Need. In Advances in Neural Information Processing Systems, volume 30. Curran Asso- ciates, Inc., 2017. URL https: httpp://prooeeedngg www.sne.urisscco.e wwtp:ccaapape//1///aaaaaa Awtaaa-bbtttttttteeeeaaeee html.
[26] Yifan Wu, George Tucker, and Ofir Nachum. Behavior Regularized Offline Reinforcement Learning. CoRR, abs/1911.11361, 2019. URL http: http:/awwwv....rraaaaaaaa/ wwtp://////////aeeeeeeee
arXiv:
1911.11361.
[27] Aviral Kumar, Justin Fu, Matthew Soh, George Tucker, and Sergey Levine. Stabilizing Off- Policy Q-Learning via Bootstrapping Error Reduction. In Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer, Florence Alché Buc, Emily Fox, and Roman Garnett, editors, Advances in Neural Information Processing Systems 32: Annual Conference
Neural Information Processing Systems 2019, NeurlPS 2019, December 8-14, 2019, Vancouver, BC, Canada, pages 11761-11771, 2019. URL https: /ttttps::/prooee h/tp:///eio......o/////a/ neurips...... cc/ (paper/2019/has/) 777755f4777533774448851111118882 htwp://bboooosaaaaaaaatt..c..cc Abstract.
[28] Scott Fujimoto and Shixiang Shane Gu. Minimalist Approach to Offline Reinforcement Learning. In Marc' Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan, editors, Advances in Neural Information Processing Systems 34: Annual Conference on Newral Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual, pages 20132-20145, 2021. URL https: http////rooceeneseeeeeeee/eessee neur suurips.coo cc/ www.:/eeee--bbaaa..e.ea Abstract. html.
paper/ http://hhhhheeeaaaeaeeee
12

---

[29] Machel Reid, Yutaro Yamada, and Shixiang Shane Gu. Can Wikipedia Help Offline Reinforce- ment Learning? CoRR, abs/2201.1212, 2022. URL httpp:://arsa org/ abs/ /200/ 12122.
arXiv: 2201.12122.
[30] Pedro A. Ortega, Markus Kunesch, Grégoire Delétang, Tim Genewein, Jordi Grau-Moya, Joel Veness, Jonas Buchli, Jonas Degrave, Bilal Piot, Julien Pérolat, Tom Everitt, Corentin Tallec, Emilio Parisotto, Tom Erez, Yutian Chen, Scott E. Reed, Marcus Hutter, Nando de Freitas, and Shane Legg. Shaking the foundations: delusions in sequence models for interaction and control. CoRR, abs/2110.10819, 2021. URL https://arrs http:/oog//bbbb/ba/eeeeee 10819. arXiv: 2110.10819.
[31] Alex Beutel, Jilin Chen, Zhe Zhao, and Ed H. Chi. Data Decisions and Theoretical Implications when Adversarially Learning Fair Representations. CoRR, abs/1707.00077 .0075, 2017. URL
http://aaxiv..rrgabbsee 170. 00075. arXiv: 1707.00075.
[32] Yaroslav Ganin, Evgeniya Ustinova, Hana Ajakan, Pascal Germain, Hugo Larochelle, François Laviolette, Mario Marchand, and Victor S. Lempitsky. Domain-Adversari Training of Neural Networks. Gabriela Csurka, editor, Domain Adaptation in Computer Vision Applications, Advances in Computer Vision and Pattern Recognition, pages 189-209. Springer, 2017. doi: URL pesttps://dol: 1000//1000079978--33119
10. 00.1007/978-33119-584477. 10.
[33] Brian Hu Zhang, Blake Lemoine, and Margaret Mitchell. Mitigating Unwanted Biases with Adversarial Learning. In Jason Furman. Gary E. Marchant, Huw Price, and Francesca Rossi, editors, Proceedings of the 2018 AAAI/AC Conference on AI, Ethics, and Society, AIES 2018, New Orleans, LA, USA, February 02-03, 2018, pages 335-340. ACM, 2018. doi: 10.1 1411437787213327777777
[34] Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron C. Courville, and Yoshua Bengio. Generative Adversarial Networks. CoRR, abs/1406.2661, 2014. URL httLt//://aat wwwrriiiv..gg/ee hts///b//1/// /eee//14..../ 2661. arXiv: 1406.2661.
[35] Mengjiao Yang, Dale Schuurmans, Pieter Abbeel, and Ofir Nachum. Dichotomy of Control: Separating What You Can Control from What You Cannot. CoRR, abs/2210.13435, 2022. doi:
00448550/rXxw.222000 13435. arXiv: 2210.13435.
[36] loannis Antonoglou, Julian Schrittwieser, Sherjil Ozair, Thomas K. Hubert, and David Silver. Planning in Stochastic Environments with Learned Model. September 2021. URL https:
/ /openreviev. wwwtteeeutdid/eeeeeeeee (tt.rrmmmm..xX
[37] Antonin Raffin, Ashley Hill, Adam Gleave, Anssi Kanervisto, Maximilian Ernestus, and Noah Dormann. Ktable-Baselinest Reliable Reinforcement Learning Implementations. Journal of Machine Learning Research, 22(268): 1-8, 2021. URL http://jmmle org httg//pprr//eeeee /22/
20 1364. html.
[38] Ilya Loshehilov and Frank Hutter. Decoupled Weight Decay Regularization. In 7th International Conference Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net, 2019. URL
[39] Sergey loffe and Christian Szegedy. Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift. In Francis R. Bach and David M. Blei, editors, Proceed- ings of the 32nd International Conference Machine Learning, ICML 2015, Lille, France, 6-11 July 2015, volume 37 of JMLR Workshop and Conference Proceedings, pages 448-456. JMrR. 2015. URL http: pttp///ecccee....//oo// mlr www.m.r..eeeeeeee war.spress/vv v/ress////// h3t/vvfff//f5//// http//fffff5...11 html.
[40] Adam Paszke. Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang. Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang. Junjie Bai, and Soumith Chintala. PyTorch: An Imperative Style, High- Performance Deep Learning Library. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d' Alché- Buc, E. Fox, and Garnett, editors, Advances in Neural Information Processing Systems 32, pages 8024-8035. Curran Associates, Inc., 2019. URL http:: neurips cttpspprrccc hi high dwwpeerrraaaeeeeeeeeeeeeeeeeeeee 2015-pytorche an- aan-imperative-see att.pppreetttttt.ieeeeeaeaeaaeoe learning- library. pdf.
13

---

[41] Lukas Biewald. Experiment Tracking with Weights and Biases, 2020. URL https://www wandb.com com/.
[42] Takuma Seno and Michita Imai. d3rlpy: An Offline Deep Reinforcement Learning Li- brary. CoRR, abs/a111.037/6 2021. URL wrt. 03788. arXiv: 2111.03788.
Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? [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] See section 5.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [Yes]
(b) Did you include complete proofs of all theoretical results? [Yes] See the appendix in the supplemental materials.
3. If you ran experiments...
(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 available in the supplemental material.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they chosen)? [Yes] See appendix.
(c) Did you report error bars (e.g., with respect to the random seed after running exper- iments multiple times)? [Yes] Experiments are run on multiple seeds with standard deviation error bars.
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, cloud provider)? [Yes] See appendix.
4. If you using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators? [N/A]
(b) Did you mention the license of the assets? [N/A]
(c) Did you include any assets either in the supplemental material as URL? [N/A]
(d) Did you discuss whether and how consent was obtained from people whose data you're 're using/curating? [N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable information offensive content? [N/A]
5. If you used crowdsourcing conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, applicable? [N/A]
(b) Did you describe any potential participant risks, with links Institutional Review Board (IRB) approvals, if applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]
14


Reviews:
[{"review_id": "V7M77NsRUcd", "reviewer": "Reviewer_bsi4", "summary": "This paper first identify that the existing offline \"reinforcement learning via supervised learning\" (RvS) methods that conditioned on returns/outcomes (e.g. Decision Transformers) can only reliably achieve outcomes that are independent from environment stochasticity. In other words, they do not consider whether the offline demonstration results are achieved by pure luck. \n\nBased on this observation, the authors proposed learning learn environment-stochasticity-independent representations of trajectories, by adversarially clustering offline experiences by data collection policy, to disentangle actions from environment stochasticity. Trajectories in each cluster can be annotated with the average return for trajectories. These pieces together give ESPER, environment-stochasticity-independent representation. RvS with ESPER achieves substantial improvements on several simple stochastic environment benchmarks: Gambling (an illustrative single-step stochastic environment), Connect Four with a stochastic opponent, and 2048.", "questions": "As mentioned in the above section, I wonder if we can have more discussion around potential limitations and future challenges of the proposed approach.\n\nnitpicks\n* In the Figure 2 caption, it's unclear what \"... such that the dynamics within each cluster matches the environments dynamics ...\" means. I suggest rephrasing it.\n* I recommend spelling out \"ESPER\" at least in introduce. It doesn't appear until Sec 2.3, which makes it difficult to follow.\n", "limitations": "The potential negative societal impact is addressed. However, the limitation is not sufficiently discussed, which I think is also my main concern about his paper and discussed in the \"Strengths And Weaknesses\" section. ", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n* I appreciate the insightful observation that existing RvS can only reliably achieve outcomes that are independent from environment stochasticity. This paper presents a clear illustration to explain why this could be an issue. I believe future RvS research could potentially benefit a lot from this work.\n* The proposed solution is theoretically sound. Empirically, it shows decent improvements, compared to return conditioning, on several simple benchmarks, including Gambling, Connect Four, and 2048. I think these are enough proof that the issue does exist, and that the proposed method ESPER achieve some initial successes.\n\nWeaknesses:\nMy concerns are about whether this paper has discussed enough limitations of the proposed solution (i.e. clustering by policy) and enough future challenges. Besides high-dimensional observation space, I think clustering by behavior/policy implies that, when the data is collected with many behavioral modes, the number of clusters need to be able to scale with that. For that reason, I think we will perhaps see that it struggles with high-dimensional and continuous action space. Potentially, (state-dependent) no-ops and near-no-ops could also lead to unnecessary complexity. Furthermore, my haunch is that clustering by behavior will go against behavioral compositionality, making it difficult to interpolate desired outcomes, as opposed to return-conditioning, and thus probably go against the authors' desire to \"pave the way for creation of more general agent\" as stated in the conclusion section. \n\nGiven the early position of this paper (which focus on a limited set of empirical evaluations), I would like to see much more discussion about the potential limitations and future challenges of the proposed approach, to help future works to extend this contribution. If these could be sufficiently addressed or clarified during rebuttal, which I think would not be very difficult, I will increase my rating to 6 or 7 to recommend acceptance. ", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "Kh7-lIc0DeH", "reviewer": "Reviewer_mkK7", "summary": "This paper proposes an approach to handling the challenge of applying RL via supervision methods to stochastic environments. The method is empirically demonstrated to outperform two recent baselines across three environments.", "questions": "1) Please clarify how the observation regarding the fundamental problem of applying RL via supervision agents to stochastic environments differs from or compliments the findings of \"Shaking the foundations: delusions in sequential models for interaction and control\" Ortega et al. [Arxiv 2020]. This could improve my rating of the significance of contribution (1).\n\n2) Why was CQL not included in Figures 4-6? Does including it provide further insight into the relative strengths and weaknesses of these approaches? This could improve my rating of the significance of contribution (2).\n\n3) On lines 243-244, the paper notes \"ESPER can both win and lose with perfect accuracy depending on the target performance\" but misses an opportunity to discuss why it doesn't do well at drawing (0 target return). Is this because this return is not included in the training distribution due to the policies used to generate the offline dataset of trajectories? This question would help clarify my current understanding of the empirical results.", "limitations": "The authors have adequately addressed the limitations and potential negative societal impact of their work. I would be very interested in seeing the follow up study proposed regarding testing the performance of ESPER on other types of goals as I agree this may require further innovation.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "At the bottom of page 2, the paper claims three contributions:\n1) the observation that RL via supervision agents \"can only reliably achieve outcomes that are independent from environment stochasticity.\"\n2) the proposed method ESPER\n3) the stochastic offline RL benchmark tasks used to evaluate ESPER\n\nHowever:\n- observation (1) was previously made in \"Shaking the foundations: delusions in sequential models for interaction and control\" Ortega et al. [Arxiv 2020]\n- the environments (3) contribution cannot be fully evaluated as code is only available for the simplest of the three environments. Open sourcing all environments would improve the significance of this contribution.\n\nESPER (the proposed method) remains a novel contribution that significantly outperforms two recent related baselines. The empirical evaluation of this method could be improved by including CQL in figures 4-6 and quantifying the correlation between performance and target return currently visualized in Figure 6. I agree with the conclusion the authors reach from this plot, but in this case believe the raw value would be more informative than the plots currently used. The space saved could then be used to include the more insightful results currently in Appendix A.3 on the representations learnt.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "GnLbbwxLR37", "reviewer": "Reviewer_Hn8H", "summary": "Offline Decision Transformers trained with supervised learning do not account for the fact that some favorable outcomes are due to stochastic dynamics and not to the quality of the policy. \n\nTherefore, by design, some probability mass may be distributed to sub-optimal actions. The solution proposed is to focus on trajectory statistics that are independent of uncontrollable randomness in the environment. From cluster assignments based on these statistics, useful conditioning quantities such as expected returns can be learnt.\n\nThree predictors are optimized alternatively with two losses:\n\n- On the one hand, a transition predictor is trained to predict the next state conditioned on the previous state, action, and the trajectory cluster assignment.\n- On the other hand, the cluster assignment predictor is trained to maximize the loss of the transition predictor while maximizing the likelihood of a policy predictor trained to predict the action conditioned on the current state and cluster assignment.\n\nFinally, another predictor is trained to predict the return conditioned on a cluster assignment. It is then used to label trajectories for policy learning conditioned on expected returns.\n\nThe authors propose a benchmark of three stochastic environments and build offline datasets with various policies ranging from random policies to experts. Empirically, their approach outperforms naive return-conditioning and a Conservative Q-Learning baseline.", "questions": "On a simple Atari environment considered in the original [Decision Transformer](https://arxiv.org/pdf/2106.01345.pdf) paper (e.g. Pong or Breakout), would it be possible to run your set of experiments with an actual Decision Transformer instead of an MLP? Their [code](https://github.com/kzl/decision-transformer) and dataset of offline trajectories are available.\n\n", "limitations": "The main limitation of this work is the lack of empirical evidence beyond toy environments and multi-layer perceptrons. While I agree that MLPs can be used instead of Transformers in the MDPs considered, the question of whether Decision Transformers fail in complex stochastic environments is not answered.", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "# Strengths\n\n- To the best of my knowledge, this method is new. It is simple to understand, well motivated, and easy to implement.\n- Claims made at the end of the introduction are supported by experimental results.\n- The paper is mostly easy to follow and understand.\n- The authors highlight a potential key limitation of Decision Transformers trained with offline trajectories in stochastic environments. As Transformers become more popular as agents, such insights are beneficial to the community.\n\n# Weaknesses\n\nThe paper is titled “Why Decision Transformers Fail in Stochastic Environments” but not a single Decision Transformer is trained nor evaluated in this work. While the authors provide evidence that MLP-based behavioral cloning agents conditioned on returns fail in toy stochastic environments, they do not illustrate this phenomenom in more complex environments with recent architectures. For instance, Decision Transformers in [Atari with sticky actions](https://arxiv.org/abs/1709.06009) or [Crafter](https://arxiv.org/abs/2109.06780) could be considered.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept