Dichotomy of Control: Separating What You Can Control from What You Cannot
Authors: Sherry Yang
Dale Schuurmans
Pieter Abbeel
Ofir Nachum
Venue: ICLR 2023 notable top 5% (2023)
Keywords: Offline reinforcement learning
return-conditioned supervised learning
stochastic environments
decision transformer

Abstract: Future- or return-conditioned supervised learning is an emerging paradigm for offline reinforcement learning (RL), in which the future outcome (i.e., return) associated with a sequence of actions in an offline dataset is used as input to a policy trained to imitate those same actions. While return-conditioning is at the heart of popular algorithms such as decision transformer (DT), these methods tend to perform poorly in highly stochastic environments, where an occasional high return associated with a sequence of actions may be due more to the randomness of the environment than to the actions themselves. Such situations can lead to a learned policy that is inconsistent with its conditioning inputs; i.e., using the policy – while conditioned on a specific desired return – to act in the environment can lead to a distribution of real returns that is wildly different than desired. In this work, we propose the dichotomy of control (DoC), a future-conditioned supervised learning framework that separates mechanisms within a policy’s control (actions) from those outside of a policy’s control (environment stochasticity). We achieve this by conditioning the policy on a latent variable representation of the future and designing a mutual information constraint that removes any future information from the latent variable that is only due to randomness of the environment. Theoretically, we show that DoC yields policies that are consistent with their conditioning inputs, ensuring that conditioning a learned policy on a desired high-return future outcome will correctly induce high-return behavior. Empirically, we show that DoC is able to achieve significantly better performance than DT on environments with highly stochastic rewards (e.g., Bandit) and transitions (e.g., FrozenLake).

Published as a conference paper at ICLR 2023
DICHOTOMY OF CONTROL: SEPARATING WHAT YOU CAN CONTROL FROM WHAT YOU CANNOT
Mengjiao Yang University of California, Berkeley Google Research, Brain Team sherryy@berkeley..dd
Dale Schuurmans
University of Alberta
Google Research, Brain Team
Pieter Abbeel
Ofir Nachum
University of California, Berkeley
Google Research, Brain Team
ABSTRACT
Future- or return-conditioned supervised learning is an emerging paradigm for offline reinforcement learning (RL), where the future outcome (i.e., return) asso- ciated with an observed action sequence is used as input to a policy trained to imitate those same actions. While return-condinn conditionnn is at the heart of popular al- gorithms such as decision transformer (DT), these methods tend to perform poorly highly stochastic environments, where an occasional high return can arise from randomness in the environment rather than the actions themselves. Such situations can lead to a learned policy that is inconsistent with conditioning inputs; i.e., specific desired
using the policy to act in the environment, when conditioning return, leads to a distribution of real returns that is wildly different than desired. In this work, we propose the dichotomy of control (DoC), cature-conditioned su- pervised learning framework that separates mechanisms within a policy's 's control (actions) from those beyond a policy's 's control (environment stochasticity). We achieve this separation by conditioning the policy on latent variable represen- tation of the future, and designing mutual information constraint that removes any information from the latent variable associated with randomness in the envi- ronment. Theoretically, we show that DoC yields policies that are consistent with their conditioning inputs, ensuring that conditioning learned policy on a desired high-return future outcome will correctly induce high-return behavior. Empiri- cally, we show that DoC is able to achieve significantly better performance than DT on environments that have highly stochastic rewards and transitions.'
INTRODUCTION
Offline reinforcement learning (RL) aims to extract an optimal policy solely from an existing dataset of previous interactions (Fujimoto al., 2019; Wu et al., 2019; Kumar al., 2020). As researchers begin to scale offline RL to large image, text, and video datasets (Agarwal et al., 2020; Fan et al., 2022: Baker et al., 2022; Reed al., 2022; Reid et al., 2022), a family of methods known as return- conditioned supervised learning (RCSL), including Decision Transformer (DT) (Chen et al., 2021; Lee al., 2022) and RI. via Supervised Learning (RvS) (Emmons et al., 2021), 1), have gained pop- ularity due to their algorithmic simplicity and case of scaling. At the heart of RCSL is the idea of conditioning a policy on a specific future outcome, often a return (Srivastava et al., 2019; Kumar et al., 2019; Chen et al., 2021) but also sometimes goal state generie future event (Codevilla et al., 2018; Ghosh et al., 2019; Lynch et al., 2020). RCSL trains a policy to imitate actions as- sociated with conditioning input via supervised learning. During inference (i.e., at evaluation), the policy is conditioned on a desirable high-return or future outcome, with the hope of inducing behavior that can achieve this desirable outcome.
'Code available https: //gi thub. goog1 1e gessearrh/googleeeeessssse

---

Published as a conference paper at ICLR 2023
RCSL/ Decision Transformer
Dichotomy of Control
a
r = 100
E
a
100
T = 0.01
S'
T 0.01
S
a
S'
2
S'
= 10
2
10
T=1
T=1
Figure 1: Illustration of DT (RCSL) and DoC. Circles and squares denote states and actions. Solid arrows denote policy decisions. Dotted arrows denote (stochastic) environment transitions. All arrows and nodes are present in the dataset, i.e., there are trajectories, which achieve reward. DT maximizes returns across entire trajectory, leading to suboptimal policies when a large return (r = 100) is achieved only due to very low cow-probayility environment transitions (T' = 0.01). DoC separates policy stochasticity from that of the environment and only tries to control action decisions (solid arrows), achieving optimal control through maximizing expected returns at each timestep.
Despite the empirical advantages that come with supervised training (Emmons al., 2021; Kumar et al., 2021), 1), RCSL can be highly suboptimal in stochastic environments (Paster et al., 2022; Brand- fonbrener et al., 2022), where the future RCSL policy conditions on (e.g., return) can be primarily determined by randomness in the environment rather than the data collecting policy itself. Figure (left) illustrates an example, where conditioning an RCSL policy on the highest return observed in the dataset (r = 100) leads to a policy (az) that relies on a stochastic transition of very low proba- bility (T 0.01) to achieve the desired return of r 100; by comparison the choice of a2 is much better in terms of average return, as it surely achieves = 10. The crux of the issue is that the RCSI, policy is inconsistent with its conditioning input. Conditioning the policy on a desired return (i.e., 100) to act in the environment leads to distribution of real returns (i.e., 0.01 100) that is wildly different from the return value being conditioned on. This issue would not have occurred the policy could also maximize the transition probability that led to the high-return state, but this is not possible as transition probabilities are a part of the environment and not subject to the policy's control.
A number of works propose a generalization of RCSL, known as coture-conditioned supervised learning methods. These techniques have been shown be effective in imitation learning (Singh et al., 2020; Pertsch et al., 2020), offiin Q-learning (Ajay et al., 2020), and online policy gradi- ent (Venuto et al., 2021). It is common in future-conditioned supervised learning to apply a KL divergence regularizer on the latent variable- - inspired by variational auto-encoders (VAE) (Kingma & Welling, 2013) and measured with respect learned prior conditioned only on past informa- tion to limit the amount of future information captured in the latent variable. It is natural to ask whether this regularizer could remedy the insconsistency of RCSL. Unfortunately, as the KL regu- larizer makes no distinction between future information that is controllable versus that which is not, such an approach will still exhibit inconsistency, the sense that the latent variable representation may contain information about the future that is due only to environment stochasticity.
is clear that the major issue with both RCSL and naive variational methods is that they make no distinction between stochasticity of the policy (controllable) and stochasticity of the environment (uncontrollable) (Paster et al., 2020; Strupl et al., 2022). An optimal policy should maximize over the controllable (actions) and take expectations over uncontrollable (e.g., transitions) as shown in Figure 1 (right). This implics that, under variational approach, the latent variable representation that a policy conditions on should not incorporate any information that is solely due to randomness in the environment. In other words, while the latent representation can and should include informa- tion about future behavior (i.e., actions), it should not reveal any information about the rewards or transitions associated with this behavior.
To this end, we propose a future-conditionea supervised learning framework termed dichotomy of control (DoC), which, in Stoie terms (Shapiro, 2014), has "the serenity to accept the things it cannot
2

---

Published as a conference paper at ICLR 2023
change, courage to change the things it can, and wisdom to know the difference." DoC separates mechanisms within a policy's 's control (actions) from those beyond a policy's ' control (environment stochasticity). To achieve this separation, we condition the policy a latent variable representa- tion of the future while minimizing the mutual information between the latent variable and future stochastic rewards and transitions in the environment. DoC only captures information from the con- trollable actions and avoids capturing information from the uncontrollable environment transitions in the latent variable so that maximization only happens with respect the controllable actions. Theoretically, we show that DoC policies are consistent with their conditioning inputs, ensuring that conditioning on a high-return future will correctly induce high-return behavior. Empirically, we show that DoC can outperform both RCSL and naive variational methods on highly stochastic environments.
2 RELATED WORK
Return-Conditione Supervised Learning. Since offline RL algorithms (Fujimoto et al., 2019;
Wu et al., 2019; Kumar et al., 2020) can be sensitive to hyper-parameters and difficult apply in practice (Emmons et al., 2021; Kumar et al., 2021), return-conditioned supervised learning (RCSL) has become a popular alternative, particularly when the environment deterministic and near-expert demonstrations are available (Brandfonbrener et al., 2022). RCSL learns to predict behaviors (ac- tions) by conditioning on desired returns (Schmidhuber, 2019; Kumar et al., 2019) using an MLP policy (Emmons et al., 2021) or a transformer based policy that encapsulates history (Chen et al., 2021). Richer information other than returns, such as goals (Codevilla et al., 2018; Ghosh et al., 2019) or trajectory-level aggregates (Furuta et al., 2021), have also been used as inputs to condi- tional policy in practice. Our work also conditions policies on richer trajectory-level information in the form of latent variable representation of the future, with additional theoretical justifications of such conditioning in stochastic environments.
RCSL Failures in Stochastic Environments. Despite the empirical success of RCSL achieved by DT and RvS, recent work has noted the failure modes in stochastic environments. Paster et al. (2020) and Strupl et al. (2022) presented counter-examplee vironments. Paster al. (2022) first identified the failure of retur-conditioned supervised learning with stochastic transitions and proposed to cluster offline trajectories and condition the policy on the average cluster returns. While conditioning on expected as opposed to maximum returns is more reasonable, this approach has technical limitations (see Appendix C), and can lead to undesirable policy-averaging, i.c., single policy covering two very different behaviors (clusters) that happen to have the same return. Brandfonbrener et al. (2022) also identified near-determinism as a nec- essary condition for RCSL to achieve optimality guarantees similar to other offline RL algorithms but did not propose a solution for RCSL in stochastic settings. Villaflor et al. (2022) also identi- fies overly optimistic behavior of DT and proposes to use discrete 3-VAE to induce diverse future predictions policy can condition on. This approach only differs the issue with stochastic environ- ments to stochastic latent variables, i.c., the latent variables will still contain stochastic environment information that the policy cannot reliably reproduce.
where online RvS can diverge in stochastic en-
Learning Latent Variables from Offline Data. Various works have explored learning latent variable representation of the future (or past) transitions in offline data via maximum likelihood and use the latent variable to assist planning (Lynch et al., 2020), imitation learning (Kipf et al., 2019; Ajay et al., 2020; Hakhamaneshi et al., 202 1), offline RL (Ajay et al., 2020; Zhou et al., 2020), or online RL (Fox et al., 2017; Krishnan et al., 2017; Goyal al., 2019; Shankar & Gupta, 2020: Singh et al., 2020; Wang et al., 2021; Venuto et al., 2021). These works generally focus on the benefit of increased temporal abstraction afforded by using latent variables as higher-level actions in a hi- erarchical policy. Villaflor et al. (2022) has introduced latent variable models into RCSL, which is one of the essential tools that enables our method, but they did not incoporate the appropriate con- straints which can allow RCSL to effectively combat environment stochasticity, as we will see in our work. In general, existing work in latent variable models for future-conditioned supervised learn- ing provides no theoretical guarantees in the literature, whereas our approach provides consistency guarantees. Lastly, Dietterich et al. (2018) used mutual information constraints to separate control- lable from uncontrollable aspects of an MDP with the goal of accelerating reinforcement learning, whereas we study the dichotomy of control under the context of return-and and-ffuure conditioned su- pervised learning.
3

---

empty sub-episode. The return associated with an Published as a conference paper at ICLR 2023
3 PRELIMINARIES
Environment Notation We consider the problem of learning decision-making agent interact with a sequential, finite-horizon environment. At time O, the agent observes an initial state so determined by the environment. After observing s, at timestep 0 S t S H. the agent chooses an action dt. After the action is applied the environment yields an immediate scalar reward and, if t < H, a next state 8-+1- We use == (st, 0,, Pr) to denote a generic episode generated from interactions with the environment, and use Ts: == (s,, az, TT) to denote a generic sub-episode, with the understanding that 70:-1 refers episode T is defined as R(r) = Σ Tb.
We will use M to denote the environment. We assume that M is determined by stochastic reward function R, stochastic transition function T. and unique initial state sp. that r, RTo=:-11 1, St. ae) and 81+1 (T 8t, ax) during interactions with the environment. Note that these definitions specify bistory-dependent environment, as opposed to less general Markovian environment.
Learning Policy in RCSL In future- or return-conditioned supervised learning, one uses fixed training data distribution D of episodes (collected by unknown and potentially multiple agents) to learn policy x, where æ is trained predict a, conditioned on the history To:-1, the observation st, and an additional conditioning variable z that may depend on both the past and future of the episode. For example, return- return-conditioned supervised learning, policy training minimizes the following objective over
CRCSL(T) = EDND
Σ
log ग (aa -o--11 1, St, 2(7)
where z(7) is the return R(T).
Inconsistency of RCSL To apply RCSL-trained policy during inference i.e., interacting online with the environment - one must first choose a specific 2.² For example, one might set be the maximal return observed in the dataset, in the hopes of inducing a behavior policy which achieves this high return. Using as shorthand to denote the policy π conditioned on a specific 2, we define the expected return VM (T=) of n; in M as,
(M((T)) E. [R(T)]-
Ideally the expected return induced by n. close 2, i.e. VM((T:), so that acting according to conditioned on high return induces behavior which actually achieves high return. However, RCSL training according to Equation will generally yield policies that are highly inconsistent in stochastic environments, meaning that the achieved returns may be significantly different than (i.e., (T,) * 2). This has been highlighted in various previous works (Brandfonbrener et al., 2022; Paster et al., 2022: Strupl et al., 2022; Eysenbach et al., 2022; Villaflor et al., 2022), and we provided our own example in Figure 1.
Approaches to Mitigating Inconsistency A number of future-conditioned supervised learning approaches propose to learn a stochastic latent variable embedding of the future, q(=/7), while reg- ularizing with KL--divergence from learnable prior conditioned only on the past p(z]so (Ajay et al., 2020; Venuto al., 2021; Lynch et al., 2020), thereby minimizing:
H
LVVE(T, = E, EIND,
Σ
log π (arl S1, =) + B ETAD (Dr. (Daaa((())))ll](sssa))
(3)
One could consider adopting such a future-conditooned objective RCSL. However, since the KI. regularizer makes no distinction between observations the agent control (actions) from those cannot (environment stochasticity), the choice of coefficient 3 applied to the regularizer introduces
Forr simplicitly, we assume is chosen at timestep t 0 and held constant throughout an entire episode. As noted in Brandfonhrener et al. (2022), this protocol also encompasses instances like DT (Chen al., 2021) in which timestep the (desired) return summed starting f.
4

---

Published as a conference paper at ICLR 2023
'lose-lose" trade-off. Namely, as noted in Ajay et al. (2020), if the regularization coefficient is too large (3 2 1), the policy will not learn diverse behavior (since the KL limits how much information of the future actions is contained in 2); while if the coefficient is too small (8 < 1), the policy's learned behavior will be inconsistent (in the sense that 2 will contain information of environment stochasticity that the policy cannot reliably reproduce). The discrete /- B-VAE incoporated by Villaflor et al. (2022) with 3 < corresponds to this second failure mode.
DICHOTOMY OF CONTROL.
In this section, we first propose the DoC objective for learning future-conditioned policies that are guaranteed to be consistent. We then present practical framework for optimizing DoC's con- strained objective in practice and an inference scheme to enable better-than-datan behavior via learned value function and prior.
4.1 DICHOTOMY OF CONTROL VIA MUTUAL INFORMATION MINIMIZATION
As elaborated in the prevous section, whether =(7) is the return R((T) more generally stochastic latent variable with distribution q(=lT), existing RCSL methods fail to satisfy consistency because they insufficiently enforce the type of future information can contain. key observation that should not include any information due to environment stochasticity. i.e., any information about a future rt, 81+1 that is not already known given the previous history up to that point To:t-1, 8t, az- (A similar observation was made in (Paster al., 2022) under the more restrictive assumption that z is a cluster index, which we do not require here.) To address the independence requirement general and sound way, modify the RCSL objective from Equation to incorporate conditional mutual information constraint between 2 and each pair rı, S141 in the future:
LDoC(TT 9) := BrnDDe-g(EP))
log π To:e (4,
s.t. MI(ri; 10 ax) 0, 0,M((5++ii2 T0, ar) =0, Y TD:t 1, St, de and 0 < H,
(4)
where MI(r:: 2 To:t-1, St, ax) denotes the mutual information between F₃ and given T0:-11 8t: at when measured under samples of r,, from D, and analogously for MI 2/70 Sr, az).
The first part of the DoC objective conditions the policy latent variable representation of the future, similar to the first part of the future-conditioned VAE objective Equation 3. However, unlike Equation 3, the DoC objective enforces a much more precise constraint g. given by the MI constraints in Equation 4.
4.2 DICHOTOMY OF CONTROL IN PRACTICE
Contrastive Learning of DoC Constraints. To satisfy the mutual information constraints in Equation 4 transform the MI contrastive learning objective. Specifically, for the con- straint on r and 2 (and similarly on 8141 and 2) one can derive,
MI(re: To. St, a,)
DKL (Pr(re,2/Tou St, ael
Pr[r To:--1 St, ai Pr a.)
Pr[relz, Tu:-1, ar)
Epe(rial mi
log
Pr(ra To. St, ar)
Epr(r roir (1. log Pr(rxle, T0: 8c, 0r]
Eprir roc-11 log Pr|rz /0a-1, St, ar). (6)
The second expectation above is a constant with respect to and so can be ignored during learning. We further introduce conditional distribution wrr, To:4-1 St, ax) parametrized by an energy- based function f Ω + R:
w i,, T0.-1, St, 0x) & p(r:) exp (f T0 St, ax)}
where p is some fixed sampling distribution of rewards. In practice, we set p to be the marginal dis- tribution of rewards in the dataset. Hence express the first term of Equation 6 via an optimization
5

---

Published as a conference paper at ICLR 2023
over w. i.e,
max e
at)]
~
= max pe [Tr, 10 81, ae) - log Ep(r) [expL_ To St,a,)) az)
Combining this (together with the analogous derivation for (5++1:1; 70: 8t, a+)) with Equa- tion 4 via the Lagrangian, we can learn π and ((=lT) by minimizing the final DoC objective:
H
CDo 9) = max B D,
log π (a. To,-,
+
B Br-D,zng(=1t) [f 8641, T0.-1, a4) - log Bp(r,)) [exp(f 20 Toa-1, a]H
t=
(8)
Algorithm Inference with Dichotomy of Control
Inputs Policy x-1l:1,)), prior p(-), value function V(),, initial state So, number of samples hy- perparameter K.
Initialize 2"; V*
Track the best latent and its value.
for = to K do
Sample 2 p(2]ss) if V (zx) (zz)) > V then
Sample latent from the learned prior.
Z45 V* =
Set best latent to the one with the highest value. D Policy conditioned on the best =".
return 2-1+:,,,")
DoC Inference. As is standard in RCSL approaches, the policy learned by DoC requires an appro- priate conditioning input be chosen during inference. To choose a desirable 2 associated with high return, propose to (I) enumerate or sample large number of potential values of z, (2) es- timate the expected return for each of these values of z. (3) choose the = with the highest associated expected return to feed into the policy. To enable such inference-ting time procedure, two more components to the method formulation: First, a prior distribution p(2[so) from which we will sample large number of values of 2; second, value function V(z) with which we will rank the potential values of z. These components are learned by minimizing the following objective:
need to add
Laux (V,p) E..D.
(V(=)-R(7)) + DKI. (en..sooprrd(((t))))llobb))))))
Note that we apply trop-ggraient& to (((TT) when learning p so to avoid regularizing q with the prior. This is unlike the VAE approach, which by contrast advocates for regularizing via the prior. See Algorithm I for inference pseudocode (and Appendix D for training pseudocode). Since DoC requires sampling multiple latents, it incurs additional computational overhead compared to vanilla DT. We differ improving the computational efficiency of DoC to future work.
CONSISTENCY GUARANTEES FOR DICHOTOMY OF CONTROL
We provide a theoretical justification of the proposed learning objectives LDoC and Lunx, showing that, if they are minimized, the resulting inference- DoC will learn V and such that the true value of т. in the environment M is equal (z)). More specifically we define the following notion of consistency:
procedure will be sound, in the sense that
Definition (Consistency). A future-conditions policy a and value function V are consistent for specific conditioning input if the expected return of predicted by V is equal to the true expected return of nz in the environment: V(z) (M(
To guarantee consistency of V, we will make the following two assumptions:
Assumption 2 (Data and environment agreement). The per-step reward and next-state transi- tions observed in the data distribution are the same as those of the environment. In other words,

---

Published as a conference paper at ICLR 2023
for any T0:1-1, 81,4, az with Pr(ro:-11, Sc, aD 0. have Prffe relTode-1 St, at DJ R(fe To:--1, at) and Pr(š1+1 = To.1 St, at, DJ ( St, az) for all ft. 81+1- Assumption (No optimization or approximation errors). DoC yields policy and value function V that are Bayes-optimal with respect to the training data distribution and q. In other words, V(z) = E..t [R(T]] and ((ū 10:-11, S1,
To. St, 2, D..
Given these two assumptions,
then establish the following consistency guarantee for DoC.
Theorem 4. Suppose DoC yields т, V, with q satisfying the MI constraints:
MI(ri To 0+) Mi($++1; To 1t. Q4)
(10)
for all To:-1; St, with Pr(fo: at,a, > 0. Then under Assumptions 2 and 3, V and a are consistent for any a with Pr[z/g, DJ >
For proof, see Appendix A.
Consistency in Markovian environments. While the results above are focused environments and policies that are non-Markovian, one can extend Theorem to Markovian environments and policies. This result is somewhat surprising, the assignments of to episodes induced by are necessarily history-dependent policy would seemingly lose important information. However, Markovian sumption on the rewards and transitions of the environment is sufficient to ensure that no 'important' information will be lost, least terms of the satisfying requirements for consistency in Defini- tion 1. Alternative notions of consistency are not as generally applicable; see Appendix C.
and projecting the actions appearing in these clusters to a non-
We begin by stating our assumptions.
Assumption 5 (Markov environment). The rewards and transitions of M are Markovian; i.e., St, az) and ((To-=1, S1, ae) = - 8t, at) for all T, St, az We use the shorthand R(s,,a)), T(s,,a) ) for these history- iintror-iideeeeeenttte functions.
RT St, (t) = R(Too St,
function.n
Assumption (Markov policy, without optimization or approximation errors). The policy learned by DoC Markov. This policy as well as its corresponding learned value function are Bayes-optimal with respect to the training data distribution and 9. In other words, V(2) = Ex.PP(HE) [R(T)] and n[ 81, Pr [a at 8t; 2, D).
With these two assumptions, can then establish the analogue to Theorem 4, which relaxes the dependency on history for both the policy æ and the MI constraints:
Theorem 7, Suppose DoC yields n, V.g with q satisfying the MI constraints:
MI(re; 2/82, 04) = MT(s++1; 2|st, ae) = =0,
(11)
for all 8t, a, with Pr[se, az|D) [D] > 0. Then under Assumptions 2, 5, and 6, and a are consistent for any with Pr[z|g, D] 0.
For proof, see Appendix B.
EXPERIMENTS
We conducted an empirical evaluation to ascertain the effectiveness of DoC. For this evaluation, considered three settings: (1) Bernoulli bandit problem with stochastic rewards, based on canonical 'worst-case scenario' for RCSL (Brandfonbrener et al., 2022); (2) the FrozenLake do- main from (Brockman al., 2016), where the future VAE approach proves ineffective; and finally (3) a modified set of OpenAl Gym (Brockman et al., 2016) environments where we introduced en- vironment stochasticity. In these studies, we found that DoC exhibits a significant advantage over RCSL/DT, and outperforms future VAE when the analogous to "one-step" RL is insufficient. We consider vanilla DT (Chen et al., 2021), which trains a "returns-to-go. state, action" transformer. For DoC and the VAE baseline, we adopt similar inference scheme as (Lee et al., 2022), where multiple latent is sampled to determine the best next action. For DT, we use the same implementation and hyperparameters as Chen al. (2021). Both VAE and DoC are built upon the DT implementation and additionally learn a Gaussian latent variable over succeeding 20 future steps. See experiment details in Appendix E and additional results in Appendix F.

---

Published as a conference paper at ICLR 2023
6.1 EVALUATING STOCHASTIC REWARDS IN BERNOULLI BANDIT
Bernoulli Bandit. Consider two-armed bandit as shown in Figure 2 (left). The two arms, 01; d2. have stochastic rewards drawn from Bernoulli distributions of Bern(1 - p) and Bern(p), re- spectively. In the offline dataset, the a, arm with reward Bern(1 - p) is pulled with probability (a(a)) When p is small, this corresponds to the better arm only being pulled occasionally. Under this setup, TRCSI (a = 1) = TRCSL (a = 1) = 0.5, which is highly suboptimal compared to always pulling the optimal arm d1 with reward Bern(1 p) for p 0.5.
Results. We train tabular DoC and base- lines on 1000 samples where the superior arm with r Bern(1 p) is pulled with probability for p {0.1, ..... 0.5). Fig- ure 2 (right) shows that RCSL and percent- age BC (filtered by in policies that are indifferent in the arms, whereas DoC is able to recover the Bayes- optimal performance (dotted line) for all p values considered. Future VAE performs similarly to DoC for small values, but sensitive to the KL regularization coef- ficient when p is close to 0.5.
BC
RCSL/ DT
DoC
%BC
VAE
Bern(1 p)
xppaa) =p
performaace
6.8
1) always result
(nlaa) = -p
0.6
a2
r- Bern(p)
0.2
0.1
0.2
0.3 0.4 0.5
Bernoulli(
p)
Figure [Left] Bernoulli bandit where the better arm with reward Bern(1 p) for p < 0.5 is pulled with probability (a = p in the offline data. [Right] Av- erage rewards achieved by DoC and baselines across 5 environment seeds. RCSL is highly suboptimal when is small, whereas DoC achieves close to Bayes-optimal performance (dotted line) for all values of p.
6.2 EVALUATING STOCHASTIC TRANSITIONS IN FROZENLAKE
FrozenLake. Next, we consider the FrozenLake environment with stochastic transitions where the agent taking an ac- tion has probability p of moving in the in- tended direction, and probability 0.5 (1 - p) of slipping to either of the two sides of the intended direction. We collect 100 trajectories of length 100 using a DQN policy trained in the original en- vironment (p }t) which achieves an average return of 0.7, and vary p during data collection and evaluation to test different levels of stochasticity. We also include uniform actions with probability c to lower the performance of the offline data so that BC is highly suboptimal.
Results. Figure presents the visualization (left) and results (right) for this task. When the offline data is closer to being expert (e 0.3), DT, future VAE, and DoC perform similarly with better performance in more deterministic environments. As the offline dataset becomes more suboptimal (e 0.5). DoC starts to dominate across all levels of transition stochasticity. When the offline data is highly suboptimal (e = 0.7), DT and future VAE has little advantage over BC, whereas DoC continues to learn policies with reasonable performance.
RC
NBC
DDT
DoC random in / 0.7
mendam in D 0.3
mandaam in D 0.5
Start
Frozen
Holp
performane
0.6
R
0.4
0.4
0.2
0.2
0.2
0,4
0.4
moving intanded dir
Prob.
intended
Figure 3: [Left] Visualization of the stochastic Frozenlaake task. The agent has probability p of moving in the intended direction and - p of slipping to either sides. [Right] Average performance FrozenLake with different levels of stochasticity (p) and
(across 5 seeds) of DoC and baselines offline dataset quality (e). DoC outperforms DT and future VAE, where the gain is more salient when the offline data is less optimal (c 0.5 and = 0.7).
8

---

Published as a conference paper at ICLR 2023
Beacheev2
Hopperow2
HellChheetabb222
untmare-umaan diverse-10
Icees
5003
2000
500
2.2
226
Figure 4: Average performance (across seeds) of DoC and baselines on modified stochastic Gym MuloCo and AntMaze tasks. DoC and future VAE generally provide benefits over DT, where DoC provide more benefits on harder tasks such as Humanoid. Future VAE can be sensitive to the KL coefficient 3, which can result in the failure mode shown in Reacher-v2 if not tuned properly.
6.3 EVALUATING STOCHASTIC GYM MuJoCo
Environments. We now consider a set of Gym MuJoCo environments including Reacher, Hop- per, HalfCheetah, and Humanoid. We additionally consider AntMaze from D4RL (Fu et al., 2020). These environments are deterministic by default, which we modify by introducing time-correlated Gaussian noise to the actions before inputing the action into the physics simulator during data collec- tion and evaluation for all but AntMaze environments. Specifically, the Gaussian noise we introduce to the actions has O mean and standard deviation of the form (1 --0.0--f) sin(t) where í is the step number and a E [o, 1]. For AntMaze where the dataset has already been collected in the de- terministic environment by D4RL, add gaussian noise with 0.1 standard deviation to the reward uniformly with probability 0.1 (both to the dataset and during evaluation).
Results. Figure 4 shows the average performance (across 5 seeds) of DI,, future VAF, and DoC on these stochastic environments. Both future VAE and DoC generally provide benefits over DT. where the benefit of DoC is more salient in harder environments such as HalfCheetah and Humanoid. We found future VAE to be sensitive to the 3 hyperparameter, and simply using 3 = 1 can result in the falure shown in Reacher-v2.
7 CONCLUSION
Despite the empirical promise of return- or future-conditioned supervised learning (RCSL) with large transformer architectures, environment stochasticity hampers the application of supervised learning to sequential decision making. To address this issue, we proposed to augment supervised learning with the dichotomy of control principle (DoC), guiding supervised policy only con- trol the controllable (actions). Theoretically. DoC learns consistent policies, guaranteeing that they achieve the future return they are conditioned on. Empirically, DoC outperforms RCSL in highly stochastic environments. While DoC still falls short in addressing other RL challenges such as "stitching' (i.e., composing sub-optimal trajectories), we hope that dichotomy control serves as stepping stone in solving sequential decision making with large-scat scale supervised learning.
ACKNOWLEDGMENTS
Thanks to George Tucker for reviewing draft versions of this manuscript. Thanks to Mare Bellemare for help with derivations. Thanks to David Brandfonbrener and Keiran Paster for discussions around stochastic environments. We gratefully acknowledges the support of a Canada CIFAR AI Chair, NSERC and Amii, and support from Berkeley BAIR industrial consortion.
REFERENCES
Rishabh Agarwal, Dale Schuurmans, and Mohammad Norouzi. An optimistic perspective on offline reinforcement learning. In International Conference Machine Learning, pp. 104-114, PMLR, 2020.
Anurag Ajay. Aviral Kumar, Pulkit Agrawal, Sergey Levine, and Ofir Nachum. Opal: Offline prim- itive discovery for accelerating offline reinforcement learning. arXiv preprint arXiv:2010. 13611, 2020.

---

Published as a conference paper at ICLR 2023
Bowen Baker, lige Akkaya, Peter Zhokhov, Joost Huizinga, Jie Tang, Adrien Ecoffet, Brandon Houghton, Raul Sampedro, and Jeff Clune. Video pretraining (vpt): Learning to act by watching unlabeled online videos. arXiv preprint arXiv:2206.11799, 2022.
David Brandfonbrener, Alberto Bietti, Jacob Buckman, Romain Laroche, and Joan Bruna. When does return-oonditonedd supervised learning work for offline reinforcement learning? arXiv preprint arXiv:2206.01009, 2022.
Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba. Openai gym. arXiv preprint arXiv:1606.01540 2016.
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. Advances in neural information processing systems, 34:1 15044- 2021.
Felipe Codevilla, Matthias Müller, Antonio López, Vladlen Koltun, and Alexey Dosovitskiy. End- to-end driving via conditional imitation learning. In 2018 IEEE international conference on robotics and automation (ICRA), pp. 4693-4700. IEEE, 2018.
Thomas Dietterich, George Trimponias, and Zhitang Chen. Discovering and removing exogenous state variables and rewards for reinforcement learning. In International Conference Machine Learning, pp. 1262-1270. PMLR, 2018.
Scott Emmons, Benjamin Eysenbach, Ilya Kostrikov, and Sergey Levine. Rvs: What is essential for offline via supervised learning? arXiv preprint arXiv:21 2021.
Benjamin Eysenbach, Soumith Udatha, Sergey Levine, and Ruslan Salakhutdinov. Imitating past successes can be very suboptimal. arXiv preprint arXiv:2206.03378, 2022.
Linxi Fan, Guanzhi Wang. Yunfan Jiang, Ajay Mandlekar, Yuncong Yang. Haoyi Zhu, Andrew Tang, De-An Huang. Yuke Zhu, and Anima Anandkumar. Minedojo: Building open-ended embodied agents with snternet-scale knowledge. arXiv preprint arXiv:2206.08855 2022.
Roy Fox, Sanjay Krishnan, Ion Stoica, and Ken Goldberg. Multi-level discovery of deep options. arXiv preprint arXiv: 070300824, 2017.
Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data-driven reinforcement learning. arXiv preprint arXiv:2004.07219, 2020.
Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration. In International conference on machine learning, pp. 2052-2062. PMLR, 2019.
Hiroki Furuta, Yutaka Matsuo, and Shixiang Shane Gu. Generalized decision transformer for offline hindsight information matching. arXiv preprint arXiv:2111.10330 2021.
Dibya Ghosh, Abhishek Gupta, Ashwin Reddy, Justin Fu, Coline Devin, Benjamin Eysenbach, and Sergey Levine. Learning to reach goals via iterated supervised learning. arXiv preprint 19Xiv: 2019,
Anirudh Goyal, Alex Lamb, Jordan Hoffmann, Shagun Sodhani, Sergey Levine, Yoshua Bengio, and Bernhard Schölkopf. Recurrent independent mechanisms. arXiv preprint arXiv: 1909. 190.::9999.1.9....... 2019.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with stochastic actor. In International confer- ence on machine learming, pp. 1861-1870. PMLR, 2018.
Kourosh Hakhamaneshi, Ruihan Zhao, Albert Zhan, Pieter Abbeel, and Michael Laskin. Hierarchi- cal few-shot imitation with skill transition models. arXiv preprint 1rXiv:2107.08900 2021.
Diederik Kingma and Max Welling. 1312 2013.
Auto-encoding variational bayes.
arXiv preprint
10

---

Published as a conference paper at ICLR 2023
Thomas Kipf. Yujia Li, Hanjun Dai, Vinicius Zambaldi, Alvaro Manchez-Gonzalee Edward Grefen- stette, Pushmeet Kohli, and Peter Battaglia. Compile: Compositional imitation learning and execution. In International Conference Machine Learning, pp. 3418-34188 3428. PMLR. 2019.
Sanjay Krishnan, Roy Fox, Ion Stoica, and Ken Goldberg. Ddco: Discovery of deep continuous options for robot learning from demonstrations. In Conference robot learning, pp- 418-437. PMLR, 2017.
Aviral Kumar, Xue Bin Peng, and Sergey Levine. Reward-condiioned policies. arXiv preprint MiXiv:192..33444 2019.
Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. Advances in Neural Information Processing Systems, 1911179919911 2020.
Aviral Kumar, Joey Hong, Anikait Singh, and Sergey Levine. Should run offline reinforcement learning or behavioral cloning? In International Conference on Learning Representations, 2021.
Kuang- Huei Lee, Ofir Nachum, Mengjiao Yang. Lisa Lee, Daniel Freeman, Winnie Xu, Sergio Guadarrama, lan Fischer, Eric Jang, Henryk Michalewski, et al. Multi-game decision transform- ers. arXiv preprint arXiv:2205. 15241, 2022.
Corey Lynch, Mohi Khansari, Ted Xiao, Vikash Kumar, Jonathan Tompson, Sergey Levine, and Pierre Sermanet. Learning latent plans from play. In Conference Robot Learning, pp. 1113- 1132. PMLR, 2020.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wier- stra, and Martin Riedmiller. Playing atari with deep reinforcement learning. arXiv preprint 1312v: 2013.
Keiran Paster, Sheila Mcllraith, and Jimmy Ba. Planning from pixels using inverse dynamics models. arXiv preprint arXhv:2012.02419 2020.
Keiran Paster, Sheila Mcllraith, and Jimmy Ba. You can't count on luck: Why decision transformers fail in stochastic environments. arXiv preprint arXiv:2205. 2022.
Karl Pertsch, Youngwoon Lee, and Joseph Lim. Accelerating reinforcement learning with learned skill priors. arXiv preprint arXiv:2010.11944, 2020.
Martin L Puterman. Markov decision processes: discrete stochastic dynamic programming. John
Wiley & Sons, 2014.
Scott Reed, Konrad ZoIna, Emilio Parisotto, Sergio Gomez Colmenarejo, Alexander Novikov, Gabriel Barth-Maron, Mai Gimenez, Yury Sulsky, Jackie Kay, Jost Tobias Springenberg, et al. A generalist agent. arXiv preprint arXiv:2205.06/75. 2022.
Machel Reid, Yutaro Yamada, and Shixiang Shane Gu. Can wikipedia help offline reinforcement learning? arXiv preprint 220.:22001.11222. 2022.
Juergen Schmidhuber. Reinforcement learning upside down: Don't predict rewards-jus map them to actions. arXiv preprint arXiv:1912.02875, 2019.
Tanmay Shankar and Abhinav Gupta. Learning robot skills with temporal variational inference. In International Conference on Machine Learning, pp. 8624-8633, PMLR, 2020.
Fred R Shapiro. Who wrote the serenity prayer? The Chronicle Review, 28, 2014.
Avi Singh, Huihan Liu, Gaoyue Zhou, Albert Yu, Nicholas Rhinehart, and Sergey Levine. Parrot: Data-driven behavioral priors for reinforcement learning. arXiv preprint arXiv:2011. 2020.
Rupesh Kumar Srivastava, Pranav Shyam, Filipe Mutz, Wojeiech Jaskowski. and Jürgen Schmidhu- ber. Training agents using upside-down reinforcement learning. arXiv preprint 19120719,, 2019.
11

---

Published as a conference paper at ICLR 2023
Miroslav Strupl, Francesco Faccio, Dylan R Ashley, Jürgen Schmidhuber, and Rupesh Kumar Srivastava. Upside-down reinforcement learning can diverge stochastic environments with episodic resets. arXiv preprint 2rXhv:2205.06595, 2022.
David Venuto, Elaine Lau, Doina Precup, and Ofir Nachum. Policy gradients incorporating the future. arXiv preprini arXiv: 108. 220000,,,,,,006, 2021.
Adam R Villaflor, Zhe Huang, Swapnil Pande, John M Dolan, and Jeff Schneider. Addressing optimism bias in sequence modeling for reinforcement learning. International Conference on Machine Learning, pp. 22270-22235 PMLR, 2022.
Xiaofei Wang. Kimin Lee, Kourosh Hakhamaneshi, Pieter Abbeel, and Michael Laskin. Skill pref- erences: Learning to extract and execute robotic skills from human feedback. arXiv preprint 0553: 2021.
Yifan Wu, George Tucker, and Ofir Nachum. Behavior regularized offline reinforcement learning. arXiv preprint arXiv:1911.11366 2019.
Wenxuan Zhou, Sujay Bajracharya, and David Held. Plas: Latent action space for offline reinforce- ment learning. arXiv preprint 1rXiv:2011.02/1 2020.
12

---

Published as a conference paper at ICLR 2023
Appendix
A PROOF OF THEOREM 4
The proof relies on the following lemma, showing that the MI constraints ensure that the observed rewards and dynamics conditioned on in the training data equal the rewards and dynamics of the environment.
Lemma 8. Suppose DoC yields 9 satisfying the MI constraints:
MT(ri =/T0:-1,00 04) MT($++1; 2/ St. 04) O,
(12)
for all T04-11 8c, ae with Pr[To:-11,&, 81,
0. Then under Assumption
Pr ff = Tt To:--1, 86s at 2, DJ - R(e Tu
(13)
Pr Pr14111 = Se+1 To a2 D] T( 1t, az);
(14)
for all St, ar, and f, S141, as long as Pr[to:t-1, St, azs 2 2
Proof. We show the derivations relevant to reward, with those for next-state being analogous. We start with the definition of mutual information:
Pr [r' Toa-1, Gz. 2, D]
MI(ra 2 a4) Ptr.a))-P(f log
(15)
Pr [rx T0: 1, St, or: D]
Ee [KL (Xr[[[]To: St, a,, 2, 2,, Pr a, D))
(16)
The KL divergence is nonnegative quantity, and is zero only when the two input distributions are equal. Thus, the constraint MI(r:: z/Tu.-1,, ae) 0 implies,
Pr [r|/oo--1,11 0t, D = Po at, D),
(17)
for all 70:-1. 1: St, dz. with Pr[z|mos-1,88 oz. D 0. From Assumption 2 we know
Prfr|7o--11.8 86 0e. D] 86
(18)
and so we immediately have the desired result.
We will further employ the following lemma, which takes most of the way to proving Theorem 4:
Lemma 9, Suppose DoC yields n,q with satisfying the MI constraints:
Mi(ri 2|70:-11 8c: 0x) MI($4+15 2/Toa-1, 8t, (4) 0,
(19)
for all T1 St, ag with Pr[ a, [D] 0. Then under Assumptions and 3, we have
[r D] = Pr M),
(20)
for all 7 and all with Pr(s|g, 2) 0.
Proof. We may write the probability Pr [T 2, D] as,
Pr(r 2, D]
Pr (ar To:-11, 81,
1=
Pr [re Tox- 8t, as 2, D)
H
Pr [8+41 70a-1, 8t, 0r, 0,,
x, D). (21)
13

---

Published as a conference paper at ICLR 2023
Case 1:
We begin by considering the case of satisfying Pr [r[r 12, DJ Assumption we may write the first probability above as
0. For such by
Pr [a 10:--1; 8t, 2, D) (a To:- 8t).
(22)
Moreover, by Lemma 8 may write the second and third probabilities as
Pr [rx To ar, 2, D (0:-1100
(23)
Pr [8(411 T0 2, D T (#141 81, ae)
(24)
Therefore, for any with Pr (T 2, D] we have,
H
H
H-1 II t=0
Pr[T 2, D] = II nx (ad Toxc-1
II 2=0
R(ndToo-1. 11, ar) 78 81,
t=0
= Pr Tz, M).
(25)
Case 2: To handle the case of Pr [r 2, 2, = will show that Pr (roze D] 0 implies
Pr [Tore M] = 0 by induction on t. The base case of t -1 is trivial. For t > -1, we may write,
Pr [7: 2, D2 [s( To:- 81-1, DJ Pr(a; [a, 1t, D) Pr [oor 2,, D)
Pr T0:-11, 0c, D). (26)
81 81-1, a-11) T Pr [Fod M = Pr [Toa-1 T2. 2
st).
R(r, To: az). (27)
Suppose, for the sake of contradiction, that Pr (Toa 2, DJ = 00 while Pr [Tue Tz > 0. By the inductive hypothesis, Pr [To:-- z, D > Thus, by Lemma we must have
Pr [s: To:-2, St- 1,-1, 2, D] ((s
(28)
and T(sc (o: St-1, a)-1) implies that Pr (st Assumption we must have
0. Thus, by
Pr [a To:-11, St, DJ n. (a; 11
(29)
and ==(aa//0--11 8t) > 0 implies that Pr (ar Th.t1, St, z, D 0. Lastly, by Lemma 8 must have
Pr [re. T0.t-1, ae 2, DJ R(r 80 ar);
(30)
and Rrr (o:r//7:t11,11 St, 0x) 0 implies that Pr [r+ Tp:t-1; 0t, 2, > 0. Altogether, find
that each of the three terms on the RHS of Equation 26 is strictly positive and so Pr (roa 2, D] > 0; contradiction.
A.1 THEOREM PROOF
We are now prepared to prove Theorem 4.
Using Assumption we can express V(z) as,
/ Pr = 22 R(f) de.
V(x)
(31)
By Lemma we have,
v(2) / Pr = 2 R(f) di
(32)
-Pr(ft
/ Pr[f = M) R(f) di
(33)
VM((TT).
(34)
desired.
14

---

Published as a conference paper at ICLR 2023
B PROOF OF THEOREM 7
We begin by proving result under stricter conditions, namely, when the MI constraints retain the conditioning on history.
Lemma 10. Suppose DoC yields π, V,q with q satisfying the MI constraints:
MT(r;; = 8c; 0+) M1(8++1; 2/ St; 0e) O,
(35)
for all 1o St, a2 with Pr[vo.-11 1, Sc. arl 0. Then under Assumptions 2, 5, and 6, V and are consistent for any z with Pr|z|g, D] >
Proof. Let
hist (a 1: Pr [ā- 0z 11, DJ. ==
(36)
By Lemma and Theorem 4 we have
Pr[r 2, 2, Pr [T x== bist MM.
(37)
for all 7 and
V(=)= = VVM(T hist)
(38)
for all 2 with Pr[2 q, DJ > O.
It is left to show that (m!iitt (=((T)). To do so, we invoke Theorem 5.5.1 in Puterman policy, there exists a Markov policy such that
(2014), which states that, for any history-depende the state-action visitation occupancies of the two policies are equal (and, accordingly, their values are equal). In other words, there exists a Markov policy # such that
Pr [$ = 82, ā = bibe M Pr (a 81, â at 1. M]:
(39)
for all 8,8,6 à, and
VT hist
(40)
To complete the proof, show that मे n. By Equation 37 we have
Pr [s St, at x bist M
Pr [$ st, â at D)
(41)
Thus, for any 8, à have
Pr ($
st, â 0t i. MM|
f ag $ si)
(42)
Pr = St | M)
Pr $
S1, â 04 Thist M
(43)
Pr = 8t hist .M]
Pr ($ 8t, 0x 2
(44)
Pr [$ 8t 2, DI
= (a a, $ st),
(45)
where the first equality is Bayes' rule, the second equality is due to Equation 39, the third equality is due to Equation 41, and last equality is by definition of т. (Assumption 6).
Before continuing to the main proof, we present the following analogue to Lemma 8:
Lemma 11. Suppose DoC yields q satisfying the MI constraints:
MI(re 2/8t, 0x) MT(s++1; 2|8t, ae) = O,
(46)
for all 8t, at with Pr[se, ag [D] > O. Then under Assumptions 2 and 5,
rt St, az z, D) R(f St, az), Pr [8 = St-1 St, at 2, D] = T ($ Scc 0t),
Pr lf
(47)
(48)
for all St, at, 2 and f, Stt1, as long as Pr[st, 0z, 2 [D] > 0.
Proof. The proof is analogous to the proof of Lemma 8.
15

---

Published as a conference paper at ICLR 2023
B.1 THEOREM PROOF
We can now tackle the proof of Theorem 7. To do so, we start by interpreting the episodes in the training data D as coming from a modified Markovian environment MM. Specifically, we define M as environment with the same state space as M but with action space consisting of tuples (a, r, s"), where a is an action from the action space of M. r a scalar, and is a state from the state that the
space of M. We define the reward and transition functions of M to be deterministic, reward and next state associated with (a, r, 8') is r and s', respectively. This way, we may interpret any episode (5t, de, To) in M an episode
r = (St, (at, Te, (2)
(49)
in the modified environment Mt. Denoting Dt as the training data distribution when interpreted in this way, we note that the MI constraints of Lemma 10 hold, since rewards and transitions are deterministie. Thus, the policy nt defined as
π (()),, Pr[(â, f, 8') = (ar, stt1)ss,,,
(50)
satisfies
V(z) Vin
(51)
left to show that (=!) = (T,). To do so, consider an episode -1 Pr|-Ixx, For any single-step transition in this episode,
(st, 814 re, 8141,
(52)
we have, by definition of -
Pr(a aclse, m!] Pr(a- a, 2, =āālss
(53)
In similar vein, by definition of कł and Lemma 11 we have,
Pr(f re|st, az π Pr[f = ri|st, a,, 2, Dij RR(( at);
(54)
Pr($141 8441/8,, o2 π Pr(s 81-1/88, a2 2, DI] T (8441|s,,0)
(55)
Thus, any 7' (st, (a,,T,, Tz,, st+1) = sampled from Mi can be mapped back to
(St, ae, Tt) T in the original environment M, where Pr[r [T]] Pr|T It clear that R(r!)- R((T), and so we immediately have
Var (m!) = VM((ss,
(56)
as desired.
C INVALIDITY OF ALTERNATIVE CONSISTENCY FRAMEWORKS
Paster et al. (2022) propose similar but distinct notion of consistency compared ours (i.e., Def- inition 1), and claim that it can be achieved with stationary policies in Markovian environments. In this section, we show that this is, in fact, false, supporting the benefits of our framework. We begin by rephrasing Theorem 2.1 of Paster et al. (2022) using our own notation:
(Incorrect) Theorem 2.1 of Paster et al. (2022). Suppose M is Markovian and D. are given such that
Pr[še+1 = 8t 8t, ar, x, D] = Pr[še+1 = St 8t, Oz, D]. for all St, a, 2, 8:41 with Pr[se, ar, 2/g. D and define a Markov policy π as
(57)
(ā]ss,, 2) Pr(a a, Se|St.2 2, D).
(58)
Then for any with Pr[z|q, D) > 0 and any т,
Pr[r T2, M] > 0 if and only if Pr(r 2, D) > 0.
(59)
16

---

Published as a conference paper at ICLR 2023
Figure 5: Deterministic environment used in the counter-exanee described in Appendix C. Circles represent states and squares represent actions; solid arrows represent choice actions and dashed arrows represent environment dynamics.
Counter-example. simple pounter-example may be constructed by considering the Markovian environment displayed in Figure The environment has three states. The first state gives choice of two actions (ou E {0,1)), and each action deterministically transitions to the same second state. The second state again provides choice of two actions (a1 E {0, 1}). and each of these again deter- ministically transitions to the same terminal state. Thus, episodes in this environment are uniquely determined by choice of ag; 01 There are four unique episodes:
To {ao 0, ai 0),
(60)
(ag = 1, 1),
(61)
T2 (ag 0, 01 1),
(62)
(ao = 1,ai 01 0)-
(63)
We now construct q as a deterministic function, clustering these four trajectories into two distinct 8:
Zo =
(64)
((2)) = (73).
(65)
Suppose includes To: T1, T2, Ty with equal probability. Since the environment is deterministic, the conditions of Theorem 2.1 in Paster et al. (2022) are trivially satisfied. Learning policy with respect to 2n yields
<-1150, 20) = |0..5, 0.5],
(66)
x(1181, =0) = [0.5, 0.5].
(67)
However, it is clear that interacting with क (-1,, =0) in the environment will lead to with non- zero probability, while 72, T3 are never associated with 2y in the data D. Contradiction.
D PSEUDOCODE FOR DoC TRAINING
Algorithm 2 Training with Dichotomy of Control
Inputs Offline dataset D {m(m (
Jm=1 where (m)) (m) (m) ((1)) ) with initial states
dg.
1mm1 and initial return-to-go values {1((m) L. Jmmis a parametrized distribution (-), a policy To, (s)), value function Va, (-), prior (c), energy function f.r(),, a fixed distribution
(80 (m) LM
p(r, s'), J.larning rates 7). and training batch size B.
while training has not converged do
Sample batch {(r (8t, 0e, rt) ) (BB) ==
from D, for m 1...., B. Sample z from (T) with reparametrization.
Compute LDOC + Laux according to Equation 8 and Equation 9.
Update o o V. + n wstopgrad Ø), 1/ C. 01 + 01 V,
82 02 1) Vo, C. return Te, (,, VØ, (.), P.p
17

---

1.0 for DoC, Best of 0.1, 1.0, 10 for VAF Published as a conference paper at ICLR 2023
E EXPERIMENT DETAILS
E. HYPERPARAMETERS.
We use the same hyperparameters as the publically available Decision Transformer (Chen et al., 2021) implementation. For VAE, we additionally learn future and a prior both parametrized the same as the policy using transformers with context length 20. All models are trained on NVIDIA GPU P100.
Table 1: Hyperparameters of Decision Transformer, future-conditooned VAE, and Dichotomy of Control.
Hyperparameter
Value
Number of layers
3
Numher of attention heads
1
Embedding dimension
128
Latent future dimension
128
Nonlinearity function
ReLU
Batch size Context length K
64
20 FrozenLake, HalfChectah, Hopper, Humanoid, AntMaze
5 Reacher
Future length Ky Return-to-go conditioning for DT
Same as context length K
1 FrozenLake
6000 HalfCheetah
3600 Hopper
5000 Humanoid
50 Reacher
1 AntMaze
Dropout
0.1
Learning rate
10
Grad norm clip
0.25
Weight decay
10
Learning rate decay
Linear warmup for first 10 training steps
B coefficient
E.2 DETAILS OF THE OFFLINE DATASETS
FrozenLake. We train a DQN (Mnih et al., 2013) policy for 100k steps the original 4x4 Frozen- Lake Gym environment with stochasticity level p = We then modify to simulate environments of different stochasticity levels, while collecting 100 trajectories level using the trained DQN agent with probability of selecting random action opposed the action output by the DQN agent to emulate offline data with different quality.
maximum length 100 at each
Gym MuJoCo. We train SAC (Haarnoja et al., 2018) policies on the original set of Gym MuJoCo environments for 100M steps. To simulate stochasticity in these environments, modify the origi- nal Gym MuJoCo environments by introducing noise to the actions before inputting the action to the physics simulator to compute rewards and next states. The noise has 0 mean and standard deviation of the form (1 - - -0.01-t) sin(t) a where is the step number and a E [0. 1]. We then collect 1000 trajectories of 1000 steps each for all environments except for Reacher (which has 50 steps in each trajectory) in the stochastic version of the environment using the SAC policy to acquire the offline dataset for training.
AntMaze. For the AntMaze task, use the AntMaze dataset from D4RL (Fu et al., 2020), which contains 1000 trajectories of 1000 steps each. We add gaussian noise with standard deviation 0. the rewards in the dataset uniformly with probability 0.1 to both the offline dataset and during environment evaluation to simulate stochastic rewards from the environment.
18

---

Published as a conference paper at ICLR 2023
F ADDITIONAL RESULTS
F.1 FROZENLAKE WITH DIFFERENT OFFLINE DATASET QUALITY.
BC
%BC
DT
VAE
DoC
random D 0.1
random 0.2
e rendom in =
0.7
pesformance
performance
0.6
performance
0.5
0.4
0.4
0.3
0.2
0.2
Prob, moving in Intended
Prob. moving In Intended
Prob. moving in intended
random D 0.4
random -00.5
random D 0.8
0.6
performance
0.6
performance
performance
0.5
0.4
0.2
0.2
0.1
Prob. moving in intended
Prob. moving in Intended
Prob. moving in intended
0.10
performance
0.5
performance
0.15
0.2
0.10
0.04
0.05
0.3
0.4
0.A
Prob. moving in intended dir
Prob.
moving intended dir
Proh. moving intended dir
Figure 6: Average performance (across 5 seeds) of DoC and baselines Frozenl akk with different levels of stochasticity (p) and offline dataset quality (e). DoC outperforms DT and future VAE with bigger gains the offline data is less optimal.
F.2 IMPROVEMENT OF DoC OVER RvS
To test the effect of applying the MI constraint to other future-conditioned supervised learning base- lines, evaluate RyS parametrized by MLPP policies (Emmons et al., 2021) with VAE and DoC modifications. general, MI.P parametrization performs worse than transformer parametrization, but DoC is still able provide significant benefit over vanilla RvS.
Rv6
RyS+VAE
RvS+DuC
rerdam in DD= 0.3
rendtam in 10:cth 0.5
random in D 0.7
start
Frozen
0.200
Toll
0.175
R
0.150
0.08
0.2
0.06
0.100
Cou
lntended
Intendidd
0.4
2ntended
Figure 7: Average performance (across 5 seeds) of DoC and baselines FrozenLake with different levels of stochasticity (p) and offline dataset quality (e). DoC outperforms RvS and future VAE with bigger gains the offline data is less optimal.
19

---

Published as a conference paper at ICLR 2023
G ADDITIONAL ABLATIONS
G.1 VAE WITH STOP GRADIENT
One difference between DoC and VAE is whether there stop gradient operation on the posterior q(=/7) when minimizing the KL-divergence between q(=lr) and the prior p(z[so). We conduct the ablation below Figure 8 where we also apply stop gradient to VAE, and observe that VAE's performance drops significantly.
VAE
VAE stopgrad
DoC
0.6
performance
0.4
0.2
0.2
0.3
0.4
0.5
0.6
Prob. moving in intended dir
Figure 8: Average performance (across 5 seeds) of DoC and baselines on FrozenLake with different levels of stochasticity (p) and offline dataset quality (e). DoC outperforms RvS and future VAE with bigger gains the offline data is less optimal.
20

---

Published as a conference paper at ICLR 2023
G.2
DoC WITH DIFFERENT REGULARIZATION STRENGTH (8)
DoC, beta=0.01
DoC, beta=0.1
DoC, beta=1
DoC, beta=
0.4
performance
0.3
0.2
0.1
0.2
0.3 0.4
0.5
0.6
Prob. moving in intended dir
Figure 9: Average performance (across 5 seeds) of DoC with different regularization strength (3). The effect of B is more pronounced when the dataset is highly optimal (e.g., random in D = 0.7), for which we found smaller 5 (e.g., 0.1) generally perform better.
G.3 DoC WITH DIFFERENT NUMBER OF FUTURE SAMPLES (K)
DoC, K=10
DoC, K=50
DoC, K=100
DoC, K=500
0.5
performance
0.4
0.3
0.2
0.3
0.4
0.5
0.6
Prob. moving in intended dir
Figure 10: Average performance (across 5 seeds) of DoC. with different number of samples during
inference (K). We found that higher number of samples leads better performance as we expect, and the gain beyound 100 samples is negligible.
21


Reviews:
[{"review_id": "taiH6272Mp", "reviewer": "Reviewer_aa3o", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "This paper proposes a new regularizer for future-conditioned supervised learning that specifically prevents the learning procedure from being biased by lucky rollouts (in terms of environment stochasticity). On highly stochastic environments, the proposed method achieves higher performance compared to the baselines.", "summary_of_the_review": "The paper claims that current future-conditioned supervised RL methods lack the ability to distinguish environment randomness from policy randomness and introduces a new regularizer to address this. While the proposed regularizer seems to help boost performance on highly stochastic environments, more experiments and clarifications of the proposed method is required.", "strength_and_weaknesses": "Strength\n- Empirical performance on environments with high stochasticity is strong compared to other future-conditioned supervised learning baselines.\n- The paper presents various examples and experiments to backup their claims.\n\nWeaknesses\n- DoC inference requires sampling latents from the learned prior and choosing the latent with the highest (estimated) value. I worry this process adds randomness to the policy's performance and thus is unstable. Also, if there exists only few expert quality trajectories on the dataset and other trajectories are low quality, much more sampling will be required to include a latent with good returns. In addition, the paper does not discuss the protocol for setting K (number of samples hyperparameter).\n- The authors did not provide the source code, which makes the credibility of the paper questionable.\n- The paper lacks results on hyperparameter sensitivity. For example, how does the return curve change if we change \\beta and K?\n\nQuestions\n- On Gym MuJoCo, why is time-correlated noise used instead of simple Gaussian noise?\n- Can the authors provide results on applying the proposed regularizer to other future-conditioned supervised learning baselines, for example RvS [1]?\n\n[1] Emmons et al., RvS: What is Essential for Offline RL via Supervised Learning?, ICLR 2022.\n\n######## Post-rebuttal comment ########\n\nThe updated manuscript addresses my concerns mentioned above. Thus, I am raising my score.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "jmUbCGlarNM", "reviewer": "Reviewer_LLnY", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "After reviewing the limitations of the current RCSL methods in capturing reward stochasticity, this paper has proposed a new Learning algorithm DoC that solves the control problem in a stochastic environment by separating what is controllable from what is not. The former is formulated as the actions at hand at each timestep, while the latter is interpreted as the stochastic environment transition probabilities. They have learned a new latent variable to represent future information and aid inference and included mutual information constraints in the optimization problem. Experiments comparing DoC with mainly RCSL/DT and future VAE methods proved the novelty of the paper in terms of the recovery of the optimal policies.", "summary_of_the_review": "Regardless of the derivation confusion, it causes me, it is a novel and original paper to me from both points of view of optimization goal formulation and empirical results. I would consider this an inspiring direction to which I would like to contribute.", "strength_and_weaknesses": "Strength:\n1. The paper has clearly pointed out the problem sets they are trying to improve on. And the way they came to their choices on mutual information sounds natural. \n2. Most of the notations and the derivations are clear to follow\n3. The experiment setup and resulting figures are clear\n\nWeakness:\n1. Some part of the derivation mentioned in the next section is confusing to me\n2. Some of the phrasings are very abstract and can be more rigorous, such as \"By only capturing the controllable factors in the latent variable, DoC can maximize over each action step without also attempting to maximize environment transitions as shown in Figure 1\".", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "GxOMFJa0lE", "reviewer": "Reviewer_kJLT", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "This paper aims to address the issue of inconsistency in return-conditioned supervised learning (RCSL), such as Decision Transformers. Specifically, when RCSL in highly stochastic environments is conditioned on the highest dataset return, the resulting expected return could be lower as the environment randomness is not under the agent's control. To this end, this paper proposes to capture only the agent's controllable factors and minimize the learning dependence on future rewards and transitions — which are environment characteristics. Under reasonable assumptions, the paper proves that their latent variable approach with Mutual Information constraints leads to policies consistent with their conditioning input reward.", "summary_of_the_review": "I like the paper's writing, problem, and proposed approach. However, it is missing some crucial justification as to why it is expected to be better than the best baseline (VAE), which prior work has proposed to mitigate inconsistency. I would happily reconsider my rating if the above issues are addressed, resulting in a clear demonstration of improvement over baselines and ablations.\n\n----  \n[Post-rebuttal]\nThe author response addressed most of my concerns, except \"DoC does not convincingly outperform VAE in control tasks.\" I have increased my score to reflect this.\n", "strength_and_weaknesses": "## Strengths\n- The paper identifies an important problem and solution for RCSL with inconsistent conditioning, especially in stochastic environments.\n- The paper is clearly written with all the necessary details and flow to understand everything.\n- A good set of simple environments are chosen, including a didactic Bernoulli Bandit environment and MuJoCo benchmarks.\n- The paper has theoretical results on the consistency of their proposed method.\n\n\n## Weaknesses\n- **Missing crucial comparisons**\n    + There are two key differences of DoC from prior approaches (i.e., VAE) using latent embedding of future to mitigate inconsistency: (a) Mutual Information (MI) constraint and (b) Inference from a learned prior and value function. The importance of these two components must be ablated on all the environments. Specifically, the following comparisons can be added:\n        * VAE + Inference with learned prior (a separate copy with stopgrad) and value function\n        * DoC w/o MI constraint (= DoC - a)\n        * DoC + conditioning on highest return (= DoC - b)\n- **Comparison against VAE**\n    + Since future-VAE also regularizes the z to the learnable prior conditioned on the past, this ensures that the latent z is incentivized not to use future information. Therefore, DoC's key benefit must come from utilizing the controllable part of the future in z while ignoring the environment transitions. Why is it expected that encoding the controllable information in z will improve the performance? What are example environments with this property?\n    + The difference between the ideology of VAE and DoC being seemingly small is reinforced by the results on the mujoco environments, where the empirical performance of VAE and DoC are pretty similar. Even on reacher, where KL beta is not tuned well, VAE first reachers almost the optimal performance before falling. Therefore, it is not clear that DoC is necessarily better than VAE.\n        * Is more stochasticity in mujoco environments expected to increase the difference in performance between DoC and VAE?\n        * Is any other experiment (quantitative or qualitative) possible to show that the importance of encoding controllable future information in z is helpful?", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "CbRusNM6E1J", "reviewer": "Reviewer_eVkG", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "The paper is in the field of \"Future- or return-conditioned supervised learning\". Here the main idea is to learn, from a batch of data, a policy that is conditioned on a latent variable z, which is often the return of a given trajectory.  However, in stochastic environments, the return is often not driven by the performance of the actions executed,  but by the inherent stochasticity of the environment. \n\nHere, the authors introduce a novel objective function, by penalizing the mutual information between the latent variable and the (environment-driven) stochastic rewards. (The method is (loosely) inspired by the stoic concept of dichotomy of control, which they used to name the method and the paper).  \n\nThe authors  test this method on a set of stochastic benchmarks  (often deterministic, but adjusted to be stochastic) and show better stability and performance", "summary_of_the_review": "The methodology is novel and its performance is empirically shown.   I have some doubts regarding the universality of the approach, which may be cleared up during the rebuttal and subsequent discussions. \n\nTherefore, for now I vote for acceptance with some reservations regarding empirical significance.", "strength_and_weaknesses": "Pros\n- The paper is written very well and good to follow\n- The paper is focused on a relevant question in the field\n- Concepts are well introduced and explained\n- Experiments make sense and are executed nicely: testing in stochastic scenarios of different scale -- visualization, experiment repetition etc.\n\n\nCons:\nMy main critique is around the point of universality:  I feel a bit unsure about how universal the approach is for different forms of stochasticity in the environment and how well it would perform if no stochasticity would be present:\n\n 1. I feel a bit unsure about Eq. (7) (especially in continuous cases): to calculate the mutual information between the reward r_t and z (and s_{t+1}) you say you \"we set ρ to be the marginal distribution of rewards in the dataset\" and that is has to be a  \"fixed sampling distribution of rewards\"  But this sounds like you will have to make simplifying assumptions here about the stochasticity of the s_t and r_t: \n    - are you limited to Gaussian and empirical (histogram) distributions? What happens if the environment is multi-modal or long-tail in its state or reward transitions? \n    -  do you learn the parameters of the conditional distribution \\omega?\n    -  what happens if the environment would be stochastic in some regions and deterministic in others?\n    -  Looking at the experiment, I noticed the algorithm is only tested in situations where randomness is either Bernoulli (6.1, 6.2) or Gaussian (6.3).  \n\n\n2. The authors evaluate their method solely on stochastic tasks (or deterministic tasks, like MuJoCo but made stochastic). I do not know how well the method performs under standard (deterministic) settings. E.g. perhaps it will behave a bit more conservative because it assumes some form of stochastic behavior\n\n\n\n\n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: notable-top-5%