Dichotomy of Control: Separating What You Can Control from What You Cannot

Introduction

Offline reinforcement learning (RL) aims to extract an optimal policy solely from an existing dataset of previous interactions (Fujimoto et al., 2019; Wu et al., 2019; Kumar et 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 et 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 et al., 2022) and RL via Supervised Learning (RvS) (Emmons et al., 2021), have gained popularity due to their algorithmic simplicity and ease 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 a goal state or generic future event (Codevilla et al., 2018; Ghosh et al., 2019; Lynch et al., 2020). RCSL trains a policy to imitate actions associated with a 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.

Despite the empirical advantages that come with supervised training (Emmons et al., 2021; Kumar et al., 2021), RCSL can be highly suboptimal in stochastic environments (Brandfonbrener et al., 2022), where the future an RCSL policy conditions on (e.g., return) can be primarily determined by randomness in the environment rather than the data collecting policy itself. Figure 1 (left) illustrates an example, where conditioning an RCSL policy on the highest return observed in the dataset ($r=100$) leads to a policy ($a_{1}$) that relies on a stochastic transition of very low probability ($T=0.01$) to achieve the desired return of $r=100$; by comparison the choice of $a_{2}$ is much better in terms of average return, as it surely achieves $r=10$. The crux of the issue is that the RCSL policy is inconsistent with its conditioning input. Conditioning the policy on a desired return (i.e., $100$) to act in the environment leads to a 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 if 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 future-conditioned supervised learning methods. These techniques have been shown to be effective in imitation learning (Singh et al., 2020; Pertsch et al., 2020), offline Q-learning (Ajay et al., 2020), and online policy gradient (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 to a learned prior conditioned only on past information – 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 regularizer makes no distinction between future information that is controllable versus that which is not, such an approach will still exhibit inconsistency, in the sense that the latent variable representation may contain information about the future that is due only to environment stochasticity.

It is clear that the major issue with both RCSL and naïve variational methods is that they make no distinction between stochasticity of the policy (controllable) and stochasticity of the environment (uncontrollable). An optimal policy should maximize over the controllable (actions) and take expectations over uncontrollable (e.g., transitions) as shown in Figure 1 (right). This implies that, under a 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 information 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-conditioned supervised learning framework termed dichotomy of control (DoC), which, in Stoic terms (Shapiro, 2014), has “the serenity to accept the things it cannot change, courage to change the things it can, and wisdom to know the difference.” DoC separates mechanisms within a policy’s control (actions) from those beyond a policy’s control (environment stochasticity). To achieve this separation, we condition the policy on a latent variable representation of the future while minimizing the mutual information between the latent variable and future stochastic rewards and transitions in the environment. 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 (right). 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 naïve variational methods on highly stochastic environments.

Related Work

Since offline RL algorithms (Fujimoto et al., 2019; Wu et al., 2019; Kumar et al., 2020) can be sensitive to hyper-parameters and difficult to 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 is deterministic and near-expert demonstrations are available (Brandfonbrener et al., 2022). RCSL learns to predict behaviors (actions) 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 a conditional policy in practice. Our work also conditions policies on richer trajectory-level information in the form of a latent variable representation of the future, with additional theoretical justifications of such conditioning 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 Štrupl et al. (2022) presented counter-examples where online RvS can diverge in stochastic environments. Brandfonbrener et al. (2022) identified near-determinism as a necessary condition for RCSL to achieve optimality guarantees similar to other offline RL algorithms but did not propose a solution for RCSL in stochastic settings. Paster et al. (2022) identified this same issue with stochastic transitions and proposed to cluster offline trajectories and condition the policy on the average cluster returns. However, the approach in Paster et al. (2022) has technical limitations (see Appendix C), does not account for reward stochasticity, and still conditions the policy on (expected) returns, which can lead to undesirable policy-averaging, i.e., a single policy covering two very different behaviors (clusters) that happen to have the same return. Villaflor et al. (2022) also identifies overly optimistic behavior of DT and proposes to use discrete $\beta$-VAE to induce diverse future predictions a policy can condition on. This approach only differs the issue with stochastic environments to stochastic latent variables, i.e., the latent variables will still contain stochastic environment information that the policy cannot reliably reproduce.

Various works have explored learning a 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., 2021), offline RL (Ajay et al., 2020; Zhou et al., 2020), or online RL (Fox et al., 2017; Krishnan et al., 2017; Goyal et 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 hierarchical 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 constraints which can allow RCSL to effectively combat environment stochasticity, as we will see in our work.

Preliminaries

We consider the problem of learning a decision-making agent to interact with a sequential, finite-horizon environment. At time $t=0$, the agent observes an initial state $s_{0}$ determined by the environment. After observing $s_{t}$ at a timestep $0\leq t\leq H$, the agent chooses an action $a_{t}$. After the action is applied the environment yields an immediate scalar reward $r_{t}$ and, if $t<H$, a next state $s_{t+1}$. We use $\tau:=(s_{t},a_{t},r_{t})_{t=0}^{H}$ to denote a generic episode generated from interactions with the environment, and use $\tau_{i:j}:=(s_{t},a_{t},r_{t})_{t=i}^{j}$ to denote a generic sub-episode, with the understanding that $\tau_{0:-1}$ refers to an empty sub-episode. The return associated with an episode $\tau$ is defined as $R(\tau):=\sum_{t=0}^{H}r_{t}$.

We will use $\mathcal{M}$ to denote the environment. We assume that $\mathcal{M}$ is determined by a stochastic reward function $\mathcal{R}$, stochastic transition function $\mathcal{T}$, and unique initial state $s_{0}$, so that $r_{t}\sim\mathcal{R}(\tau_{0:t-1},s_{t},a_{t})$ and $s_{t+1}\sim\mathcal{T}(\tau_{0:t-1},s_{t},a_{t})$ during interactions with the environment. Note that these definitions specify a history-dependent environment, as opposed to a less general Markovian environment.

In future- or return-conditioned supervised learning, one uses a fixed training data distribution $\mathcal{D}$ of episodes $\tau$ (collected by unknown and potentially multiple agents) to learn a policy $\pi$, where $\pi$ is trained to predict $a_{t}$ conditioned on the history $\tau_{0:t-1}$, the observation $s_{t}$, and an additional conditioning variable $z$ that may depend on both the past and future of the episode. For example, in return-conditioned supervised learning, policy training minimizes the following objective over $\pi$:

To apply an RCSL-trained policy $\pi$ during inference — i.e., interacting online with the environment — one must first choose a specific $z$.For simplicitly, we assume $z$ is chosen at timestep $t=0$ and held constant throughout an entire episode. As noted in Brandfonbrener et al. (2022), this protocol also encompasses instances like DT (Chen et al., 2021) in which $z$ at timestep $t$ is the (desired) return summed starting at $t$. For example, one might set $z$ to be the maximal return observed in the dataset, in the hopes of inducing a behavior policy which achieves this high return. Using $\pi_{z}$ as a shorthand to denote the policy $\pi$ conditioned on a specific $z$, we define the expected return $V_{\mathcal{M}}(\pi_{z})$ of $\pi_{z}$ in $\mathcal{M}$ as,

Ideally the expected return induced by $\pi_{z}$ is close to $z$, i.e., $z\approx V_{\mathcal{M}}(\pi_{z})$, so that acting according to $\pi$ conditioned on a high return induces behavior which actually achieves a high return. However, RCSL training according to Equation 1 will generally yield policies that are highly inconsistent in stochastic environments, meaning that the achieved returns may be significantly different than $z$ (i.e., $V_{\mathcal{M}}(\pi_{z})\neq z$). This has been highlighted in various previous works (Brandfonbrener et al., 2022; Paster et al., 2022; Štrupl et al., 2022; Eysenbach et al., 2022; Villaflor et al., 2022), and we provided our own example in Figure 1.

A number of future-conditioned supervised learning approaches propose to learn a stochastic latent variable embedding of the future, $q(z|\tau)$, while regularizing $q$ with a KL-divergence from a learnable prior conditioned only on the past $p(z|s_{0})$ (Ajay et al., 2020; Venuto et al., 2021; Lynch et al., 2020), thereby minimizing:

One could consider adopting such a future-conditioned objective in RCSL. However, since the KL regularizer makes no distinction between observations the agent can control (actions) from those it cannot (environment stochasticity), the choice of coefficient $\beta$ applied to the regularizer introduces a ‘lose-lose’ trade-off. Namely, as noted in Ajay et al. (2020), if the regularization coefficient is too large ($\beta\geq 1$), the policy will not learn diverse behavior (since the KL limits how much information of the future actions is contained in $z$); while if the coefficient is too small ($\beta<1$), the policy’s learned behavior will be inconsistent (in the sense that $z$ will contain information of environment stochasticity that the policy cannot reliably reproduce). The discrete $\beta$-VAE incoporated by Villaflor et al. (2022) with $\beta<1$ 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 a practical framework for optimizing DoC’s constrained objective in practice and an inference scheme to enable better-than-dataset behavior via a learned value function and prior.

As elaborated in the prevous section, whether $z(\tau)$ is the return $R(\tau)$ or more generally a stochastic latent variable with distribution $q(z|\tau)$, existing RCSL methods fail to satisfy consistency because they insufficiently enforce the type of future information $z$ can contain. Our key observation is that $z$ should not include any information due to environment stochasticity, i.e., any information about a future $r_{t},s_{t+1}$ that is not already known given the previous history up to that point $\tau_{0:t-1},s_{t},a_{t}$. Accordingly, we modify the RCSL objective from Equation 1 to incorporate a conditional mutual information constraint between $z$ and each pair $r_{t},s_{t+1}$ in the future:

The first part of the DoC objective conditions the policy on a latent variable representation of the future, similar to the first part of the future-conditioned VAE objective in Equation 3. However, unlike Equation 3, the DoC objective enforces a much more precise constraint on $q$, given by the MI constraints in Equation 4.

2 Dichotomy of Control in Practice

To satisfy the mutual information constraints in Equation 4 we transform the MI to a contrastive learning objective. Specifically, for the constraint on $r$ and $z$ (and similarly on $s_{t+1}$ and $z$) one can derive,

where $\rho$ is some fixed sampling distribution of rewards. In practice, we set $\rho$ to be the marginal distribution of rewards in the dataset. Hence we express the first term of Equation 6 via an optimization over $\omega$, i.e.,

As is standard in RCSL approaches, the policy learned by DoC requires an appropriate conditioning input $z$ to be chosen during inference. To choose a desirable $z$ associated with high return, we propose to (1) enumerate or sample a large number of potential values of $z$, (2) estimate the expected return for each of these values of $z$, (3) choose the $z$ with the highest associated expected return to feed into the policy. To enable such an inference-time procedure, we need to add two more components to the method formulation: First, a prior distribution $p(z|s_{0})$ from which we will sample a large number of values of $z$; second, a value function $V(z)$ with which we will rank the potential values of $z$. These components are learned by minimizing the following objective:

Note that we apply a stop-gradient to $q(z|\tau)$ when learning $p$ so as to avoid regularizing $q$ with the prior. This is unlike the VAE approach, which by contrast advocates for regularizing $q$ via the prior. See Algorithm 1 for inference pseudocode (and Appendix D for training pseudocode).

Consistency Guarantees for Dichotomy of Control

A future-conditioned policy $\pi$ and value function $V$ are consistent for a specific conditioning input $z$ if the expected return of $z$ predicted by $V$ is equal to the true expected return of $\pi_{z}$ in the environment: $V(z)=V_{\mathcal{M}}(\pi_{z})$.

To guarantee consistency of $\pi,V$, we will make the following two assumptions:

Given these two assumptions, we can then establish the following consistency guarantee for DoC.

Suppose DoC yields $\pi,V,q$ with $q$ satisfying the MI constraints:

While the results above are focused on environments and policies that are non-Markovian, one can extend Theorem 4 to Markovian environments and policies. This result is somewhat surprising, as the assignments of $z$ to episodes $\tau$ induced by $q$ are necessarily history-dependent, and projecting the actions appearing in these clusters to a non-history-dependent policy would seemingly lose important information. However, a Markovian assumption on the rewards and transitions of the environment is sufficient to ensure that no ‘important’ information will be lost, at least in terms of the satisfying requirements for consistency in Definition 1. Alternative notions of consistency are not as generally applicable; see Appendix C.

With these two assumptions, we can then establish the analogue to Theorem 4, which relaxes the dependency on history for both the policy $\pi$ and the MI constraints:

Suppose DoC yields $\pi,V,q$ with $q$ satisfying the MI constraints:

Experiments

We conducted an empirical evaluation to ascertain the effectiveness of DoC. For this evaluation, we considered three settings: (1) a Bernoulli bandit problem with stochastic rewards, based on a canonical ‘worst-case scenario’ for RCSL (Brandfonbrener et al., 2022); (2) the FrozenLake domain from (Brockman et al., 2016), where the future VAE approach proves ineffective; and finally (3) a modified set of OpenAI Gym (Brockman et al., 2016) environments where we introduced environment 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. For DT, we use the same implementation and hyperparameters as Chen et 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.

Consider a two-armed bandit as shown in Figure 2 (left). The two arms, $a_{1},a_{2}$, have stochastic rewards drawn from Bernoulli distributions of $\texttt{Bern}(1-p)$ and $\texttt{Bern}(p)$, respectively. In the offline dataset, the $a_{1}$ arm with reward $\texttt{Bern}(1-p)$ is pulled with probability $\pi_{D}(a_{1})=p$. When $p$ is small, this corresponds to the better arm only being pulled occasionally. Under this setup, $\pi_{\text{RCSL}}(a_{1}|r=1)=\pi_{\text{RCSL}}(a_{2}|r=1)=0.5$, which is highly suboptimal compared to always pulling the optimal arm $a_{1}$ with reward $\texttt{Bern}(1-p)$ for $p<0.5$.

We train tabular DoC and baselines on 1000 samples where the superior arm with $r\sim\texttt{Bern}(1-p)$ is pulled with probability $p$ for $p\in\{0.1,...,0.5\}$. Figure 2 (right) shows that RCSL and percentage BC (filtered by $r=1$) always result 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 $p$ values, but is sensitive to the KL regularization coefficient when $p$ is close to $0.5$.

2 Evaluating Stochastic Transitions in FrozenLake

Next, we consider the FrozenLake environment with stochastic transitions where the agent taking an action has probability $p$ of moving in the intended direction, and probability $0.5\cdot(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 environment ($p=\frac{1}{3}$) 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 $\epsilon$ to lower the performance of the offline data so that BC is highly suboptimal.

Figure 3 presents the visualization (left) and results (right) for this task. When the offline data is closer to being expert ($\epsilon=0.3$), DT, future VAE, and DoC perform similarly with better performance in more deterministic environments. As the offline dataset becomes more suboptimal ($\epsilon=0.5$), DoC starts to dominate across all levels of transition stochasticity. When the offline data is highly suboptimal ($\epsilon=0.7$), DT and future VAE has little advantage over BC, whereas DoC continues to learn policies with reasonable performance.

3 Evaluating Stochastic Gym MuJoCo

We now consider a set of Gym MuJoCo environments including Reacher, Hopper, 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 collection and evaluation for all but AntMaze environments. Specifically, the Gaussian noise we introduce to the actions has 0 mean and standard deviation of the form $(1-e^{-0.01\cdot t})\cdot\sin(t)\cdot\sigma$ where $t$ is the step number and $\sigma\in$. For AntMaze where the dataset has already been collected in the deterministic environment by D4RL, we add gaussian noise with 0.1 standard deviation to the reward uniformly with probability 0.1 (both to the dataset and during evaluation).

Figure 4 shows the average performance (across 5 seeds) of DT, future VAE, 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 $\beta$ hyperparameter, and simply using $\beta=1$ can result in the falure case as shown in Reacher-v2.

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 a supervised policy to only control the controllable (actions). Theoretically, DoC learns consistent policies, guaranteeing that they achieve the future or 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 of control serves as a stepping stone in solving sequential decision making with large-scale supervised learning.

Thanks to George Tucker for reviewing draft versions of this manuscript. Thanks to Marc 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

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 $z$ in the training data are equal to the rewards and dynamics of the environment.

Suppose $DoC$ yields $q$ satisfying the MI constraints:

We show the derivations relevant to reward, with those for next-state being analogous. We start with the definition of mutual information:

and so we immediately have the desired result.

We will further employ the following lemma, which takes us most of the way to proving Theorem 4:

Suppose DoC yields $\pi,q$ with $q$ satisfying the MI constraints:

Moreover, by Lemma 8 we may write the second and third probabilities as

A.1 Theorem Proof

Using Assumption 3, we can express $V(z)$ as,

Appendix B Proof of Theorem 7

We begin by proving a result under stricter conditions, namely, when the MI constraints retain the conditioning on history.

Suppose DoC yields $\pi,V,q$ with $q$ satisfying the MI constraints:

Thus, for any $t,\hat{s},\hat{a}$ we have

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 $\pi_{z}$ (Assumption 6).

Before continuing to the main proof, we present the following analogue to Lemma 8:

Suppose $DoC$ yields $q$ satisfying the MI constraints:

The proof is analogous to the proof of Lemma 8.

We can now tackle the proof of Theorem 7. To do so, we start by interpreting the episodes $\tau$ in the training data $\mathcal{D}$ as coming from a modified Markovian environment $\mathcal{M}^{\dagger}$. Specifically, we define $\mathcal{M}^{\dagger}$ as an environment with the same state space as $\mathcal{M}$ but with an action space consisting of tuples $(a,r,s^{\prime})$, where $a$ is an action from the action space of $\mathcal{M}$, $r$ is a scalar, and $s^{\prime}$ is a state from the state space of $\mathcal{M}$. We define the reward and transition functions of $\mathcal{M}^{\dagger}$ to be deterministic, so that the reward and next state associated with $(a,r,s^{\prime})$ is $r$ and $s^{\prime}$, respectively. This way, we may interpret any episode $\tau=(s_{t},a_{t},r_{t})_{t=0}^{H}$ in $\mathcal{M}$ as an episode

in the modified environment $\mathcal{M}^{\dagger}$. Denoting $\mathcal{D}^{\dagger}$ 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 deterministic. Thus, the policy $\pi^{\dagger}$ defined as

we have, by definition of $\pi_{z}^{\dagger}$,

In a similar vein, by definition of $\pi_{z}^{\dagger}$ and Lemma 11 we have,

Appendix C Invalidity of Alternative Consistency Frameworks

Paster et al. (2022) propose a similar but distinct notion of consistency compared to ours (i.e., Definition 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:

Suppose $\mathcal{M}$ is Markovian and $\mathcal{D},q$ are given such that

A simple counter-example may be constructed by considering the Markovian environment displayed in Figure 5. The environment has three states. The first state gives a choice of two actions ($a_{0}\in\{0,1\}$), and each action deterministically transitions to the same second state. The second state again provides a choice of two actions ($a_{1}\in\{0,1\}$), and each of these again deterministically transitions to the same terminal state. Thus, episodes in this environment are uniquely determined by choice of $a_{0},a_{1}$. There are four unique episodes:

We now construct $q$ as a deterministic function, clustering these four trajectories into two distinct $z$:

Suppose $\mathcal{D}$ includes $\tau_{0},\tau_{1},\tau_{2},\tau_{3}$ with equal probability. Since the environment is deterministic, the conditions of Theorem 2.1 in Paster et al. (2022) are trivially satisfied. Learning a policy $\pi$ with respect to $z_{0}$ yields

However, it is clear that interacting with $\pi(\cdot|\cdot,z_{0})$ in the environment will lead to $\tau_{2},\tau_{3}$ with non-zero probability, while $\tau_{2},\tau_{3}$ are never associated with $z_{0}$ in the data $\mathcal{D}$. Contradiction.

Appendix D Pseudocode for DoC training

Appendix E Experiment Details

We use the same hyperparameters as the publically available Decision Transformer (Chen et al., 2021) implementation. For VAE, we additionally learn a 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.

E.2 Details of the offline datasets

We train a DQN (Mnih et al., 2013) policy for 100k steps in the original 4x4 FrozenLake Gym environment with stochasticity level $p=\frac{1}{3}$. We then modify $p$ to simulate environments of different stochasticity levels, while collecting 100 trajectories of maximum length 100 at each level using the trained DQN agent with probability $\epsilon$ of selecting a random action as opposed to the action output by the DQN agent to emulate offline data with different quality.

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, we modify the original 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-e^{-0.01\cdot t})\cdot\sin(t)\cdot\sigma$ where $t$ is the step number and $\sigma\in$. 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.

For the AntMaze task, we 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.1 to 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.

Appendix F Additional Results