label-width = 2em , item-indent = 2em , before-skip = -after-skip = -after-
item-skip =

# Policy Gradients Incorporating the Future

David Venuto1,2, Elaine Lau2, Doina Precup1,2,3, Ofir Nachum4
1Mila, 2McGill University, 3DeepMind, 4Google Brain
{david.venuto<EMAIL_ADDRESS>
<EMAIL_ADDRESS><EMAIL_ADDRESS>
###### Abstract

Reasoning about the future – understanding how decisions in the present time
affect outcomes in the future – is one of the central challenges for
reinforcement learning (RL), especially in highly-stochastic or partially
observable environments. While predicting the future directly is hard, in this
work we introduce a method that allows an agent to ”look into the future”
without explicitly predicting it. Namely, we propose to allow an agent, during
its training on past experience, to observe what _actually_ happened in the
future at that time, while enforcing an information bottleneck to avoid the
agent overly relying on this privileged information. This gives our agent the
opportunity to utilize rich and useful information about the future trajectory
dynamics in addition to the present. Our method, Policy Gradients
Incorporating the Future (PGIF), is easy to implement and versatile, being
applicable to virtually any policy gradient algorithm. We apply our proposed
method to a number of off-the-shelf RL algorithms and show that PGIF is able
to achieve higher reward faster in a variety of online and offline RL domains,
as well as sparse-reward and partially observable environments.

##  1 Introduction

Fundamentally, reinforcement learning (RL) is composed of gathering useful
information (_exploration_) and assigning credit to that information (_credit
assignment_). Both of these problems present their own unique learning
challenges. In this work, we focus on credit assignment, which refers to the
challenge of matching observed outcomes in the future to decisions made in the
past. Humans appear to do this in a sample efficient manner [17], and so it is
natural to expect our own RL agents to do so as well.

One of the most popular approaches to credit assignment, known as _model-free_
RL, is to learn a value function to approximate the future return given a
starting state and action. The value function is learned using experience of
the agent acting in the environment via temporal difference (TD) methods [39],
which regress the value function to a target based on a combination of
groundtruth returns achieved in the environment and the approximate value
function itself. The need to _bootstrap_ learning of the value function on its
own estimates is known to lead to difficulties in practice, where one must
achieve a careful balance between bias and variance [15, 49, 35, 23]. If a
slight imbalance arises, the consequences can be disastrous for learning [45,
46, 41]. For example, in offline RL this issue is so pronounced that
algorithms must apply strong regularizations on both learned policy and value
function to achieve stable performance [50, 20, 53, 24].

The model-free approach plays dual to the _model-based_ approach, where an
agent learns a dynamics and reward model of the environment, and then learns
an agent to optimize behavior in this model. Thus, credit assignment boils
down to utilizing an appropriate planning algorithm that can perform multiple
rollouts in the model, effectively allowing the agent to “look into the
future” [32, 29] to determine cause-and-effect [40, 31, 2]. While model-based
RL may appear more straightforward, learning an accurate model is a challenge
in practice, presenting its own sample-efficiency problems [48] as well as
memory and computational issues [54]. Model-based approaches are thus most
beneficial when the environment exhibits some level of regularity [1].

Beyond these issues, credit assignment in both model-free and model-based RL
is further exacerbated by _partially observable_ environments, in which the
full environment state is not known to the learning agent. Thus it is
infeasible to predict future events accurately. When applied to such non-
Markovian domains, model-free algorithms relying on bootstrapping and value
function approximation tend to be biased [37]. On the other hand for model-
based approaches, learning an accurate dynamics model in such domains is a
difficult, potentially ill-defined problem [38, 6].

In this work, we aim to circumvent these challenges. We propose a simple
modification to model-free RL that allows the learned policy and value
function to “look into the future” but without the need to learn an accurate
model. Namely, we propose to modify the policy and value function to not only
condition on the presently observed state and action but also on the
subsequent trajectory (sampled by the agent as it was interacting with the
environment) following this state and action. This way, our method mitigates
potential approximation or feasibility issues in accurately modeling the
future. To ensure that the learned policy and value function remains relevant
during inference (i.e., data collection) when the future trajectory is
unavailable, we place an _information bottleneck_ [33, 43] on the additional
inputs, encouraging the learned functions to minimize their reliance on this
privileged information. One may thus view our method as an instance of
_teacher forcing_ or _$Z$ -forcing_ [11, 21] where our student is the learned
policy and value function and the teacher is some function of the information
in the future trajectory.

Practically, our method, Policy Gradients Incorporating the Future (PGIF), is
easy to implement. We use either a backwards RNN or a transformer to inject
downstream information from the observed trajectories by way of latent
variables, with a KL divergence regularization on these latents. We apply PGIF
on top of a variety of off-the-shelf RL algorithms, including RNN-based PPO
[36], SAC [14], and BRAC [50], and evaluate these algorithms on online and
offline RL as well as sparse-reward and partially observable environments. In
all of these domains, we demonstrate the ability of PGIF to achieve higher
returns faster compared to these existing RL algorithms on their own, thus
showing that our proposed method is both versatile and beneficial in practice.

##  2 Background and Notation

We begin by providing a brief overview of the notation and preliminary
concepts that we will use in our later derivations.

Markov Decision Processes (MDPs) MDPs are defined by a tuple
$\langle\mathcal{S},\mathcal{A},\mathbb{P},R,\rho_{0},\gamma\rangle$ where
$\mathcal{S}$ is a set of _states_ , $\mathcal{A}$ is a set of _actions_ ,
$\mathbb{P}$ is a transition kernel giving a probability
$\mathbb{P}(s^{\prime}|s,a)$ over next states given the current state and
action,
$R:\mathcal{S}\times\mathcal{A}\rightarrow[R_{\text{min}},R_{\text{max}}]$ is
a reward function, $\rho_{0}$ is an initial state distribution, and
$\gamma\in[0,1)$ is a discount factor. An agent in this MDP is a stationary
policy $\pi$ giving a probability $\pi(a|s)$ over actions at any state
$s\in\mathcal{S}$. A policy $\pi$ interacts with the MDP by starting at
$s_{0}\sim\rho_{0}$ and then at time $t\geq 0$ sampling an action
$a_{t}\sim\pi(s_{t})$ at which point the MDP provides an immediate reward
$R(s_{t},a_{t})$ and transitions to a next state
$s_{t+1}\sim\mathbb{P}(s_{t},a_{t})$. The interaction ends when the agent
encounters some terminal state $s_{T}$.

The value function $V^{\pi}:\mathcal{S}\to\mathbb{R}$ of a policy is defined
as $V^{\pi}(s)=\mathbb{E}_{\pi}[\sum_{t=0}^{T-1}\gamma^{t}r_{t}|s_{0}=s]$,
where $\mathbb{E}_{\pi}$ denotes the expectation of following $\pi$ in the MDP
and $T$ is a random variable denoting when a terminal state is reached.
Similarly, the state-action value function
$Q^{\pi}:\mathcal{S}\times\mathcal{A}\to\mathbb{R}$ is defined as
$Q^{\pi}(s,a)=\mathbb{E}_{\pi}[\sum_{t=0}^{T-1}\gamma^{t}r_{t}|s_{0}=s,a_{0}=a]$.
The advantage $A^{\pi}$ is then given by
$A^{\pi}(s,a)=Q^{\pi}(s,a)-V^{\pi}(s)$. We denote $\rho_{\pi}$ as the
distribution over trajectories $\tau=(s_{0},a_{0},r_{0},\dots,s_{T})$ sampled
by $\pi$ when interacting with the MDP.

During learning, $\pi$ is typically parameterized (e.g., by a neural network),
and in this case, we use $\pi_{\theta}$ to denote this parameterized policy
with learning parameters given by $\theta$. The policy gradient theorem [42]
states that, in order to optimize the RL objective
$\mathbb{E}_{s_{0}\sim\rho_{0}}[V^{\pi}(s_{0})]$, a parameterized policy
should be updated with respect to the loss function,

$J_{\text{PG}}(\pi_{\theta})=\mathbb{E}_{\tau\sim\rho_{\pi_{\theta}}}\left[\sum_{t=0}^{T-1}\gamma^{t}\cdot\hat{Q}_{t}\log\pi_{\theta}(a_{t}|s_{t})\right],$
(1)

where $\hat{Q}_{t}$ is an unbiased estimate of $Q^{\pi}(s_{t},a_{t})$. In the
simplest case, $\hat{Q}_{t}$ is the empirically observed future discounted
return following $s_{t},a_{t}$. In other cases, an approximate $Q$-value or
advantage function is incorporated to trade-off between the bias and variance
in the policy gradients. When the $Q$ or $V$ value function is parameterized,
we will use $\psi$ to denote its parameters. For example, the policy gradient
loss with a parameterized $Q_{\psi}$ is given by,

$J_{\text{PG}}(\pi_{\theta},Q_{\psi})=\mathbb{E}_{\tau\sim\rho_{\pi_{\theta}}}\left[\sum_{t=0}^{T-1}\gamma^{t}\cdot
Q_{\psi}(s_{t},a_{t})\log\pi_{\theta}(a_{t}|s_{t})\right].$ (2)

The value function $Q_{\psi}$ is typically learned via some regression-based
temporal differencing method. For example,

$J_{\text{TD}}(Q_{\psi})=\mathbb{E}_{\tau\sim\rho_{\pi_{\theta}}}\left[\sum_{t=0}^{T-1}(\hat{Q}_{t}-Q_{\psi}(s_{t},a_{t}))^{2}\right].$
(3)

#### Stochastic Latent Variable Models

In our derivations, we will utilize parameterized policies and value functions
conditioned on auxiliary inputs given by stochastic latent variables. That is,
we consider a latent space $\mathcal{Z}$, typically a real-valued vector
space. We define a parameterized policy that is conditioned on this latent
variable as $\pi_{\theta}(a|s,z)$ for
$a\in\mathcal{A},s\in\mathcal{S},z\in\mathcal{Z}$; i.e., $\pi_{\theta}$ takes
in states and latent variables and produces a distribution over actions. In
this way, one can consider the latent variable $z$ as modulating the behavior
of $\pi_{\theta}$ in the MDP. During interactions with the MDP or during
training, the latent variables themselves are generated by some stochastic
process, thus determining the behavior of $\pi_{\theta}$. For example, in the
simplest case $z$ may be sampled from a latent prior
$p_{\upsilon^{\text{(Z)}}}(z|s)$, parameterized by $\upsilon^{\text{(Z)}}$.
Thus, during interactions with the MDP actions $a_{t}$ are sampled as
$a_{t}\sim\pi_{\theta}(s_{t},z_{t}),z_{t}\sim
p_{\upsilon^{\text{(Z)}}}(s_{t})$.

We treat parameterized latent variable value functions analogously.
Specifically, we consider a latent space $\mathcal{U}$ and a parameterized
$Q$-value function as $Q_{\psi}(s,a,u)$. A prior distribution over these
latent variables may be denoted by $p_{\upsilon^{\text{(U)}}}(u|s)$.

##  3 Policy Gradients Incorporating the Future

Our method aims to allow a policy during training to leverage future
information for learning control. We propose to utilize stochastic latent
variable models to this end. Namely, we propose to train $\pi_{\theta}$ and
$Q_{\psi}$ with latent variables
$(\mathbf{z},\mathbf{u})=\\{(z_{t},u_{t})\\}_{t=0}^{T-1}$ sampled from a
learned function $q_{\phi}(\tau)$ which has access to the full trajectory. For
example, the PGIF form of the policy gradient objective in (2) may be
expressed as

$J_{\text{PGIF,
PG}}(\pi_{\theta},Q_{\psi},q_{\phi})=\mathbb{E}_{\tau\sim\rho_{\pi},(\mathbf{z},\mathbf{u})\sim
q_{\phi}(\tau)}\left[\sum_{t=0}^{T-1}\gamma^{t}\cdot
Q_{\psi}(s_{t},a_{t},u_{t})\log\pi_{\theta}(a_{t}|s_{t},z_{t})\right].$ (4)

The PGIF form of the temporal difference objective in (3) may be expressed
analogously as,

$J_{\text{PGIF,TD}}(Q_{\psi},q_{\phi})=\mathbb{E}_{\tau\sim\rho_{\pi_{\theta}},\mathbf{u}\sim
q_{\phi}(\tau)}\left[\sum_{t=0}^{T-1}(\hat{Q}_{t}-Q_{\psi}(s_{t},a_{t},u_{t}))^{2}\right].$
(5)

It is clear that any RL objective which trains policies and/or value functions
on trajectories can be adapted to a PGIF form in a straightforward manner. For
example, in our experiments we will apply PGIF to an LSTM-based PPO [36], SAC
[14], and BRAC [50].

While the PGIF-style objectives above adequately achieve our aim of allowing a
policy to leverage future trajectory information during training, they also
present a challenge during inference. When performing online interactions with
the environment, one cannot evaluate $q_{\phi}(\tau)$, since the full
trajectory $\tau$ is not yet observed.

Therefore, while we want to give $\pi_{\theta}$ and $Q_{\psi}$ the ability to
look at the full $\tau$ during training, we do not want their predictions to
overly rely on this privileged information. To this end, we introduce a
regularization on $q_{\phi}$ in terms of a KL divergence from a prior
distribution
$p_{\upsilon^{\text{}}}(\tau):=\\{p_{\upsilon^{\text{}}}(z_{t},u_{t}|s_{t})\\}_{t=0}^{T-1}$
which conditions $(z_{t},u_{t})$ only on $s_{t}$. Thus, in the case of policy
gradient, the full loss is,

$J_{\text{PGIF-KL}}(\pi_{\theta},Q_{\psi},q_{\phi})=J_{\text{PGIF,
PG}}(\pi_{\theta},Q_{\psi},q_{\phi})+\beta\mathbb{E}_{\tau\sim\rho_{\pi}}\left[D_{\text{KL}}(q_{\phi}(\tau)\|p_{\upsilon^{\text{}}}(\tau))\right],$
(6)

where $\beta$ is the weight of the divergence term. The introduction of this
prior thus solves two problems: (1) it encourages the learned policies and
value functions to not overly rely on information beyond the immediate state;
(2) it provides a mechanism for inference, namely using latent samples from
$p_{\upsilon^{\text{}}}(s)$ when interacting with the environment.

#### Parameterization of $q_{\phi}$

In our implementation, we parameterize $q_{\phi}(\tau)$ as an RNN operating in
reverse order on $\tau$. Specifically, we use an LSTM network to process the
states in $\tau$ backwards, to yield LSTM hidden states
$\mathbf{b}=\\{b_{t}\\}_{t=0}^{T-1}$. The function $q_{\phi}$ is then given by
Gaussian distributions with mean and variance at time $t$ derived from the
backwards state $b_{t}$. In practice, to avoid potential interfering gradients
from the objectives of $\pi_{\theta}$ and $Q_{\psi}$, we use separate RNNs
with independent parameters $q_{\phi^{(\text{Z})}},q_{\phi^{(\text{U})}}$ for
$z_{t},u_{t}$, respectively. See Figure 1 for a graphical diagram of the
training and inference procedure for PGIF. In our empirical studies below, we
will also show that a transformer [47] can be used in place of an RNN,
providing more computational efficiency and potentially allowing for better
propagation of information over time.

Figure 1: The architecture of the model. Our inference model
$q_{\phi^{(\text{Z})}}$ uses a backwards hidden state $b_{t}$ to approximate
dependencies of $z_{t}$ on the future of the trajectory. The blue line
separates the data collection and policy gradient training steps in our
algorithm and the red lines represent information flowing into the backwards
RNN. Grey variables are used during training and white variables are used
during data collection. Top: we show the training model where the policy
gradient loss is calculated with backwards RNN hidden state information.
Bottom: we show the data collection phase of the algorithm utilizing latent
variables sampled from the latent prior.

###  3.1 Variational Information Bottleneck Interpretation

The KL regularization we employ above may be interpreted as a variational
information bottleneck, constraining the mutual information between the latent
variable distribution and the trajectory $\tau$. Here we provide a brief
derivation establishing this equivalence.

For simplicity, we consider a specific timestep $t\in\mathbb{N}$ and a
starting state $s_{t}=s$. Let $\mathcal{T}_{\geq t}$ denote the random
variable for all information contained after and including timestep $t$ in
trajectory $\tau$ induced by $\pi$. Let $\mathcal{U}_{t}$ be the random
variable for latents $u_{t}$ induced by $q_{\phi}(\tau_{\geq t})$ conditioned
on all steps in the trajectory after and including $t$. Consider a constrained
objective minimizing $J_{\text{PGIF,TD}}$ while enforcing an upper bound
$I_{\text{max}}$ on the mutual information between the distribution of
trajectory steps and the distribution of latent variables $I(\mathcal{T}_{\geq
t},\mathcal{U}_{t}|s_{t}=s)$. This objective is given by,

$\displaystyle\min_{\psi,\phi}$ $\displaystyle
J_{\text{PGIF,TD}}(Q_{\psi},q_{\phi}|s_{t}=s):=\mathbb{E}_{\tau\sim\rho_{\pi_{\theta}}(\cdot|s_{t}=s),u_{t}\sim
q_{\phi}(\tau_{\geq t})}[(\hat{Q}_{t}-Q_{\psi}(s_{t},a_{t},u_{t}))^{2}]$ (7)
$\displaystyle\text{s.t. }I(\mathcal{T}_{\geq t},\mathcal{U}_{t}|s_{t}=s)\leq
I_{\text{max}}.$

Recall the definition of mutual information:

$\displaystyle I(\mathcal{T}_{\geq t},\mathcal{U}_{t}|s_{t}=s)$
$\displaystyle=\int p(\tau_{\geq t},{u}_{t}|s_{t}=s)\log\frac{p(\tau_{\geq
t},{u}_{t}|s_{t}=s)}{p(\tau_{\geq
t}|s_{t}=s)p({u}_{t}|s_{t}=s)}d{u}_{t}d\tau_{\geq t}$ (8)
$\displaystyle=\int\rho_{\pi_{\theta}}(\tau_{\geq
t}|s_{t}=s)q_{\phi}({u}_{t}|\tau_{\geq
t})\log\frac{q_{\phi}({u}_{t}|\tau_{\geq
t})}{p({u}_{t}|s_{t}=s)}d{u}_{t}d\tau_{\geq t},$

where $p({u}_{t}|s_{t}=s)$ is the marginal distribution over the latent
variable $p({u}_{t}|s_{t}=s)=\int q_{\phi}({u}_{t}|\tau_{\geq
t})\rho_{\pi}(\tau_{\geq t}|s_{t}=s)d\tau_{\geq t}$. This marginal is
intractable to compute directly, and so to approximate this marginal we
introduce a variational distribution $h({u}_{t}|s_{t}=s)$. By definition we
know that $D_{\text{KL}}[p({u_{t}}|s_{t}=s)\lVert h({u_{t}}|s_{t}=s)]\geq 0$.
We can then see that $\int p({u}_{t}|s_{t}=s)\log
p({u}_{t}|s_{t}=s)d{u}_{t}\geq\int p({u}_{t}|s_{t}=s)\log
h({u}_{t}|s_{t}=s)d{u}_{t}$. We therefore derive the upper bound for use in
(7) as,

$\displaystyle I(\mathcal{T}_{\geq t},\mathcal{U}_{t}|s_{t}=s)$
$\displaystyle\leq\int\rho_{\pi}(\tau_{\geq
t}|s_{t}=s)q_{\phi}({u}_{t}|\tau_{\geq
t})\log\frac{q_{\phi}({u}_{t}|\tau_{\geq
t})}{h({u}_{t}|s_{t}=s)}d{u}_{t}d\tau_{\geq t}$ (9)
$\displaystyle\leq\mathbb{E}_{\tau\sim\rho_{\pi_{\theta}}(\cdot|s_{t}=s)}\Big{[}D_{\text{KL}}(q_{\phi}({u}|\tau_{\geq
t})\lVert h({u}_{t}|s_{t}=s))\Big{]}.$

We can subsume the constraint into the objective as,

$\displaystyle\min_{\psi,\phi}$ $\displaystyle
J_{\text{PGIF,TD}}(Q_{\psi},q_{\phi}|s_{t}=s)+\beta\left(\mathbb{E}_{\tau\sim\rho_{\pi_{\theta}}(\cdot|s_{t}=s)}\left[D_{{\text{KL}}}\left(q_{\phi}(\tau_{\geq
t})\|h(u_{t}|s_{t}=s)\right)\right]-I_{\text{max}}\right).$

By taking $h$ to be our learned prior $p_{\upsilon}$, we achieve the single
step ($s_{t}=s$), TD analogue of the PGIF objective in (6), offset by a
constant $\beta\cdot I_{\text{max}}$.

###  3.2 Z-Forcing with Auxiliary Losses

While our proposed training architecture enables the policy and value function
to look at the full trajectory $\tau$, in practice it may be difficult for the
trajectory information to propagate, especially in settings with highly
stochastic learning signals. In fact, it is known that such latent variable
models may ignore the latent variables due to optimization issues, completely
negating any potential benefit [4]. To circumvent these issues, we make use of
the idea of _Z-forcing_ [12], which employs auxiliary losses and models to
force the latent variables to encode information about the future. We denote
this loss as $J_{\text{Ax}}(\zeta)$ where $\zeta$ is the set of parameters in
any auxiliary models, and elaborate on the main forms of this loss which we
consider below.

#### State based forcing (Force)

A simple way to force state information to be encoded is to derive conditional
generative models $p_{\zeta}(b_{t}|z_{t})$ over the backwards states given the
inferred latent variables $z_{t}\sim q_{\phi^{{(\text{Z})}}}(z_{t}|b_{t})$,
and similarly for the latents $u_{t}$. We can write this auxiliary objective
as a maximum log-likelihood loss
$J_{\text{Ax}}(\zeta)=-\mathbb{E}_{q_{\phi^{(\text{Z})}}(z_{t}|b_{t})}[\log
p_{\zeta}(b_{t}|z_{t})]$. This way, we enforce the noisy mapping $b_{t}\to
z_{t}$ defined by $q_{\phi^{\text{(Z)}}}$ to not be _too noisy_ so as to
completely remove any information from $b_{t}$.

#### Value prediction networks (VPN)

A more sophisticated approach to force information to be propagated is to use
an autoencoder-like, model-based auxiliary loss. To this end, we take
inspiration from VPNs [26], and apply an auxiliary loss that uses $b_{t}$ to
predict future rewards, values, and discounts. Note that, in principle,
$b_{t}$ already has access to this information, by virtue of the backwards RNN
or transformer conditioned on the future trajectory. Thus, this auxiliary loss
only serves to enforce that the RNN or transformer dutifully propagates this
information from its inputs. We also note that, in contrast to the state based
forcing described above, this approach only enforces $b_{t}$ to contain the
relevant information, and it is up to the RL loss whether this information
should be propagated to the latents $z_{t},u_{t}$.

In more detail, the VPN loss introduces the following additional learnable
functions: (2) Encoding network :$f^{(\text{enc})}:s\rightarrow x$ Outcome
encoder: $f^{(\text{out})}:x,a\rightarrow\hat{\gamma},\hat{r}$

Value predictor: $f^{(\text{val})}:x\rightarrow\hat{V}(x)$ Transition
predictor: $f^{(\text{trans})}:x,a\rightarrow\hat{x}^{\prime}$ Given a
backwards state $b_{t}$, we first embed this to
$x_{t}^{0}=f^{\text{emb}}(b_{t})$ with (a). We then make predictions on future
rewards $\hat{r}_{t}^{l}$, transition dynamics $x_{t}^{l}$, discounts
$\hat{\gamma}_{t}^{l}$, and values $\hat{V}(x_{t}^{l})$ using (b, c, d) for
$l=0,\dots,k$. Finally, we compute the VPN loss as,

$J_{\text{VPN}}^{t}=\sum_{l=0}^{k-1}(R_{t+l}-\hat{V}(x_{t}^{l}))^{2}+(r_{t+l}-\hat{r}_{t}^{l})^{2}+(\log_{\gamma}\gamma_{t+l}-\log_{\gamma}\hat{\gamma}_{t}^{l})^{2},$
(10)

where $R_{t+l}$ is the empirically observed future discounted reward in the
trajectory. The full auxiliary loss $J_{\text{Ax}}(\zeta)$ is then given by
summing up the loss in (10) over all timesteps in all trajectories.

###  3.3 Full Algorithm

The full learning objective for PGIF is thus composed of three components:
First, a latent-variable augmented RL objective, e.g., policy gradient as
shown in (4). Second, a KL regularizer, e.g., as shown in (6). Finally, an
auxiliary loss, given by either state based forcing (Force) or value
prediction networks (VPN). We present an example pseudocode of a PGIF-style
policy gradient with learned value function in Algorithm 1.

See the Appendix for further details, including how to adaptively tune the
coefficients on the KL and auxiliary loss components as well as more specific
pseudocode algorithms for advantage policy gradient and soft actor-critic.

Algorithm 1 PGIF Algorithm with State-Action Value Function Estimation

0: Initial parameters:
$\theta,\upsilon^{\text{(U)}},\upsilon^{\text{(Z)}},\phi^{\text{(U)}},\phi^{\text{(Z)}},\psi,\zeta_{\text{PG}},\zeta_{\text{TD}}$,
Weights:
$\alpha_{\text{PG}},\alpha_{\text{TD}},\beta_{\text{PG}},\beta_{\text{TD}}$

1: for $\text{policy-step }k=0,1,2,\dots,N$ do

2: Collect set of Trajectories $\mathcal{D}=\\{\tau_{i},\dots\\}:$

3: repeat

4: $z_{t}\sim p_{\upsilon^{(\text{Z})}}(z_{t})$

5: Execute: $\pi_{\theta}(a_{t}|s_{t},z_{t})$ and observe $r_{t},s_{t+1}$ from
environment.

6: until episode termination

7: for $\text{Trajectory: }\tau_{i}\in\mathcal{D}$ do

8: $\mathbf{b^{Z}}$ $=\text{BackwardsLSTM}^{\text{Z}}(\tau_{i})$ (Operates
over the entire trajectory)

9: $\mathbf{b^{U}}=\text{BackwardsLSTM}^{\text{U}}(\tau_{i})$

10: $\tau_{i}=\tau_{i}\cup\\{\mathbf{b^{Z}},\mathbf{b^{U}}\\}$

11:
$\forall\\{\mathbf{s},\mathbf{a},\mathbf{r},\mathbf{b^{Z}},\mathbf{b^{U}}\\}\in\mathcal{D}$:

12: Derive $J_{\text{Ax-PG}}(\zeta_{\text{PG}}),J_{\text{Ax-
TD}}(\zeta_{\text{TD}})$ according to any auxiliary loss.

13:
$D_{\text{TD}}=D_{{\text{KL}}}(q_{\phi^{{\text{(U)}}}}(\mathbf{u}|\mathbf{b^{U}})\lVert
p_{\upsilon^{{(\text{U})}}}(\mathbf{u}|\mathbf{s}))$

14:
$J_{\text{TD}}=\mathbb{E}_{\tau\in\mathcal{D}}\left[J_{\text{PGIF,TD}}(Q_{\psi},q_{\phi^{\text{(U)}}})+\alpha_{\text{TD}}J_{\text{Ax-
TD}}(\zeta_{\text{TD}})+\beta_{\text{TD}}D_{\text{TD}}\right]$

15:
$D_{\text{PG}}=D_{{\text{KL}}}(q_{\phi^{\text{(Z)}}}(\mathbf{z}|\mathbf{b^{Z}})\lVert
p_{\upsilon^{{(\text{Z})}}}(\mathbf{z}|\mathbf{s}))$

16: $J_{\text{PG}}=\mathbb{E}_{\tau\in\mathcal{D}}\left[J_{\text{PGIF,
PG}}(\pi_{\theta},Q_{\psi},q_{\phi^{\text{(Z)}}})+\alpha_{\text{PG}}J_{\text{Ax-
PG}}(\zeta_{\text{PG}})+\beta_{\text{PG}}D_{\text{PG}}\right]$

17: Update all parameters w.r.t: $J_{\text{PG}}$ and $J_{\text{TD}}$

##  4 Related Work

We review relevant works in the literature in this section.

#### Incorporating the future

Recent works in model-based RL have considered incorporating the future by way
of dynamically leveraging rollouts of various horizon lengths and then using
them for policy improvement [5] . Z-forcing and stochastic dynamics models
have been applied directly to learning environmental models and for behavioral
cloning while incorporating the future but not for online or offline
continuous control [19]. Our present work is unique for incorporating
Z-forcing and conditioning on the future in the model-free RL setting. A few
other methods explore the future in less direct ways. For example, RL Upside
down [34], uses both reward (or desired return) and state to predict actions,
turning RL into a supervised learning problem. They incorporate future
information by way of using these desired returns along a specified desired
horizon as input.

#### Hindsight

Hindsight credit assignment introduces the notion of incorporating the future
of a trajectory by assigning credit based on the likelihood of an action
leading to an outcome in the future [15]. These methods were extended using a
framework similar to ours, leveraging a backwards RNN to incorporate
information in hindsight [22]. Still, there are a number of differences
compared to our own work. First, in these previous works only the value
function (rather than both the value and policy functions) is provided access
to the future trajectory. Second, there is no KL information bottleneck;
rather information is constrained via an action prediction objective. Third,
these previous works do not employ any Z-forcing, while it is well-known that
learning useful latent variables in the presence of an autoregressive decoder
is difficult without Z-forcing [3]; in fact, in our own preliminary
experiments we found our algorithm performs significantly worse without any
auxiliary losses. Finally, our method is more versatile, being applicable to
off-policy and offline RL settings rather than purely on-policy RL, as in
these previous works. Nevertheless, it is an interesting avenue for future
work to investigate how to combine the best of both approaches, especially
with guarantees on variance reduction of policy gradient estimators [25] and
hierarchical policy learning [51].

Value driven hindsight modelling (HiMo) [13] proposes a hindsight value
function, which is conditioned on future information in the trajectory. In
this work, they primarily use the hindsight value function (separate from the
agent’s value function) to learn a low-dimensional representation, which is
distilled to a non-hindsight representation that is then used by the actual
actor and critic. Thus there are a few key differences from our work: (1) The
hindsight information is only used for value prediction and not policy
prediction; (2) There are no gradients passing from the RL loss to the
representation loss (only the separate value prediction loss is used to train
the representation), thus this method is arguably less end-to-end than PGIF;
(3) The only mechanism for controlling the amount of information in the
representation is through its dimension size and look-ahead, while using a KL
penalty is more flexible.

#### Combining model-based and model-free RL

Our proposed PGIF is a way of enabling a model-free RL algorithm to “look into
the future”. Although we do not learn explicit models of the MDP, our work may
nevertheless be considered as combining elements of both model-based and
model-free RL. Previous works have also claimed to do this, although in
distinct ways. For example, previous work has proposed learning a low
dimensional encoding of the environment and then performing planning in this
abstract representation along with model-free RL [1]. VPNs [26] improve
representations in model-based RL, and use a single network along with
supervised learning to learn transition and reward dynamics in an abstract
representation of the state space. Attention augment agents utilize dynamics
model roll-outs as input during policy optimization [32]. Our work takes a
very different approach and incorporates a representation of future
information using the hidden states of a backwards RNN, while leveraging VPN-
style losses only to enforce information in the RNN to be propagated from
inputs to outputs.

#### Learning stochastic latent variable RNNs

Deriving an approximate posterior over stochastic latent variables [3]
conditioned on a backwards RNN hidden state has been practically useful in
RNNs [12] and has been able to overcome being trapped in local minima [18].
Z-forcing [12] is able to learn useful latent variables that capture higher
level representations even with a strong auto-regressive decoder, by
reconstructing hidden states in a backwards RNN. We use the main methods from
Z-forcing to force our variables to learn a useful representation of the
trajectory for the agent. Our method employs an auxiliary signal similar to
[18] to design a cost with good convergence properties. Our work is a novel
application of these stochastic latent variable RNNs in online and offline
policy gradient methods.

#### Auxiliary objectives in RL

Incorporating auxiliary objectives in RL generally takes the form of intrinsic
rewards. State-space density models have been used to derive intrinsic rewards
that incentivize exploration, challenging the agent to find states that
generate "surprise" [28]. Other works use the prediction error in an inverse
model as a metric for curiosity [30]. The auxiliary objectives we utilize in
this work serve the purpose of improving the incorporation of future dynamics
information into policy optimization.

##  5 Experiments

We now provide a wide array of empirical evaluations of our method, PGIF,
encompassing tasks with delayed rewards, sparse rewards, online access to the
environment, offline access to the environment, and partial observability.

In our online RL experiments, we compare against Soft Actor Critic (SAC) [14]
and Proximal Policy Optimization (PPO) [36] with an LSTM layer in the policy
network. The comparison with PPO is particularly important since this method
also leverages a form of artificial memory of the past (but not forward-
looking like in PGIF). SAC is a state-of-the-art model-free RL method that
employs a policy entropy term in a policy gradient objective and shows optimal
performance and stability when compared with other online deep RL benchmarks.
We show the hyper-parameters for each experiment in the Appendix. In addition,
we explore using our method with a transformer as opposed to an LSTM as the
backwards networks.

###  5.1 Credit Assignment

We first aim to show that our method is highly effective in an environment
where credit assignment is the paramount objective. We examine the Umbrella-
Length task from BSUITE [27], a task involving a long sequential episode where
only the first observation (a forecast of rain or shine) and action (whether
to take an umbrella) matter, while the rest of the sequence contains random
unrelated information. A reward of $+1$ is given at the end of the episode if
the agent chooses correctly to take the umbrella or not depending on the
forecast. This difficult task is used to test the agent’s ability to assign
credit correctly to the first decision. We evaluate PGIF-versions of SAC and
PPO to vanilla SAC [14] as well as a PPO-LSTM [36] baseline. Results are
presented in Table 1. We see that our method is able to achieve best
performance, improving on the baselines by at least 50%. This suggests that
the PGIF agent is able to effeciciently and accurately propogate information
about the final reward to the initial timestep, more so than either the one-
step backups used in SAC or the multi-step return regressions used in PPO-LSTM
can.

Method | Mean Bsuite-score
---|---
PPO-LSTM | $0.33\pm 0.09$
SAC | $0.41\pm 0.03$
PGIF-PPO (VPN) | $0.46\pm 0.09$
PGIF-PPO (Force) | $0.26\pm 0.10$
PGIF-SAC (VPN) | $\textbf{0.60}\pm 0.04$
PGIF-SAC (Force) | $\textbf{0.58}\pm 0.08$
Table 1: Performance on the Umbrella-Length environment. We run our model for
$1000$ episodes steps over $5$ random seeds. The BSuite score is calculated in
terms of the regret normalized
$[\text{random},\text{optimal}]\rightarrow[0,1]$ (higher is better). The
number after $\pm$ is the standard deviation.

###  5.2 Sparse Rewards

To test the performance of our method in a sparse reward setting, we utilize
the AntMaze environment [7] where we have a simple U-shaped maze with a goal
at the end. The agent receives a reward of $+1$ if it reaches the goal (within
an L2 distance of 5) and $0$ elsewhere. The Ant starts at one end of a
U-shaped corridor and must navigate to the goal location at the other end.
This task is particularly challenging since the reward is extremely sparse.

We again compare PGIF-versions of SAC and PPO to vanilla SAC and PPO-LSTM.
Results are presented in Table 2, where we see that the baselines are unable
to make any progress on this challenging task, while PGIF is able to solve the
task to a point where it successfully navigates to the goal location in the
maze 50% of the time. We hypothesize our algorithm has benefits in these
environments since as soon as it obtains a reward signal it can adapt quickly
and make use of the signal by incorporating it in both policy and value
optimization, therefore accelerating learning.

Method | Mean Episodic Reward
---|---
PPO-LSTM | $0.00\pm 0.00$
SAC | $0.00\pm 0.00$
PGIF-PPO (VPN) | $0.20\pm 0.13$
PGIF-PPO (Force) | $0.30\pm 0.15$
PGIF-SAC (VPN) | $\textbf{0.40}\pm 0.16$
PGIF-SAC (Force) | $\textbf{0.50}\pm 0.16$
Table 2: Mean episodic reward on the AntMaze environment over 10 random seeds
trained with $3$ million environment steps. A sparse reward of $+1$ is
obtained at the end of the episode if the agent successfully reaches the goal
state within an L2 distance of $5$. The number after $\pm$ is the standard
error.

###  5.3 Partial Observability

We now aim to show that our method is not only effective in fully-observed
Markovian settings, but also in environments with partial observability. This
set of experiments uses the MuJoCo robotics simulator [44] suite of continuous
control tasks. These are a set of popular environments used in both online and
offline deep RL works [10, 9] and provides an easily comparable benchmark for
evaluating algorithm sample efficiency and reward performance. As in previous
work [52], we introduce an easy modification to these tasks to make the
environment partially observable thereby increasing the difficulty: We zero-
out a random dimension of the state space at each data collection step. This
helps us replicate a more realistic scenario for a robotic agent where not all
of the state space is accessible.

We compare a PGIF-style SAC implementation to vanilla SAC and PPO-LSTM on
these domains. We show the results of these experiments in Fig. 2. We find
that PGIF can provide improved performance on these difficult tasks,
suggesting that PGIF is able to leverage future information in the trajectory
to appropriately avoid uncertainties about the environment, more so than when
only conditioning on the immediate state (vanilla SAC) or even when
conditioning on the entire past trajectory (PPO-LSTM). Interestingly, we find
that the simple state based forcing (Force) performs more consistently better
than the more sophisticated VPN based forcing. See the Appendix for additional
results, including online evaluations without partial observability.

Figure 2: The online episodic mean reward evaluated over 5 episodes every 500
steps for MuJoCo continuous control RL tasks with partial observability. We
show the average over $5$ random seeds. $500,000$ environment step
interactions are used. The shaded area shows the standard error.

###  5.4 Offline RL Evaluations

To asses if our method is effective in an offline RL setting, we evaluate our
proposed algorithm in several continuous control offline RL tasks [8] against
Behavior Regularized Actor Critic (BRAC) [50] and Batch-Constrained Q-learning
(BCQ) [10]. BRAC operates as a straightforward modification of SAC, penalizing
the value function using a measure of divergence (KL) between the behaviour
policy and the learned agent policy. For our PGIF algorithm, we use BRAC as
the starting point.

For these offline MuJoCo tasks, we examine D4RL datasets classified as medium
(where the training of the agent is ended after achieving a "medium" level
performance) and medium expert (where medium and expert data is mixed) [8].
Datasets that contain these sub-optimal trajectories present a realistic
problem for offline RL algorithms. We also include an offline version of the
AntMaze, which is particularly challenging due to sparse rewards. We show the
results of these experiments in Table 3. We find that our method outperforms
the baselines in all but one of the tasks in terms of final episodic reward.
We hypothesize in the medium-expert setting that we perform slightly worse due
to the lack of action diversity which makes learning a dynamics representation
difficult. Interestingly, in contrast to the online results, we find that VPN
based forcing performs better than state based forcing, although even state
based forcing usually performs better than the baseline methods.

Environment | BRAC | PGIF (VPN) | PGIF (Force) | BCQ
---|---|---|---|---
ant-medium | $2731\pm 329$ | $\textbf{3250}\pm 125$ | $2980\pm 164$ | $1851\pm 94$
ant-medium-expert | $2483\pm 329$ | $\textbf{3048}\pm 362$ | $2431\pm 417$ | $2010\pm 133$
hopper-medium | $1757\pm 183$ | $\textbf{2327}\pm 399$ | $1930\pm 44$ | $1722\pm 166$
walker2d-medium | $3687\pm 25$ | $\textbf{3989}\pm 259$ | $3821\pm 341$ | $2653\pm 301$
halfcheetah-medium | $5462\pm 198$ | $6037\pm 324$ | $\textbf{6231}\pm 303$ | $4722\pm 206$
halfcheetah-medium-expert | $\textbf{5580}\pm 105$ | $5418\pm 76$ | $5491\pm 143$ | $4463\pm 88$
antmaze-umaze | $0.5\pm 0.16$ | $\textbf{0.95}\pm 0.0$ | $0.7\pm 0.15$ | $0.8\pm 0.13$
Table 3: Performance on the offline RL tasks showing the average episodic
return. The final average return is shown after training the algorithm for
$500,000$ episodes and then evaluating the policy over 5 episodes. Results
show an average of $5$ random seeds. The value after $\pm$ shows the standard
error.

###  5.5 Transformers Experiments

Using transformers for processing the future trajectory into latent variables
offers a few key benefits over traditional RNN architectures, namely better
computational efficiency with long sequences and improved ability to model
long timescale interactions. These provide a few interesting benefits for use
with our architecture, namely due to the fact that some of our environments
generate sequences up to $1000$ timesteps, which can be onerous for an RNN to
process. In the following experiments, we simply replace our RNN with a
transformer. This is inspired by work that models trajectories using a
transformer [16] for offline RL. Additionally, we show the results of
experiments where we use a $K$-fixed length context of future states as the
input to our transformer. A short length context is much faster to process
than the entire sequence of upstream states, but there is a decrease in
performance, suggesting that the entire trajectory contains useful
information. We show these results in Figure 3. The run-time of training was
decreased by approximately 40%.

Figure 3: The percent improvement evaluated over 5 episodes by replacing the
LSTM backwards network with a transformer. The $30$ timestep transformer has a
context size of a 30 timesteps look-ahead into the future of the trajectory.
We show the average over $5$ random seeds taking the final policy evaluation
mean episodic reward as the value. $500,000$ environment step interactions are
used. The error bars shows the standard error.

##  6 Discussion

In this work, we consider the problem of incorporating information from the
entire trajectory in model-free online and offline RL algorithms, enabling an
agent to use information about the future to accelerate and improve its
learning. Our empirical results attest to the versatility of our method. The
benefits of our method are apparent in both online and offline settings, which
is a rare phenomenon given that many previous offline RL works suggest that
what works well in online RL often transfers poorly to offline settings, and
vice versa. Beyond just online and offline RL, our results encompass partial
observability, sparse rewards, delayed rewards, and sub-optimal datasets,
demonstrating the ability for PGIF to achieve higher reward faster in all
settings.

Our work provides many interesting directions for future work. Namely, one
interesting avenue would be to combine our method with exploration strategies
in addition to credit assignment. Perhaps there are benefits in decreasing
access to future information during initial exploratory agent steps so the
agent is better able to explore. One could also incorporate a metric for
surprise to _gate_ access to this future information in states of high
novelty.

We also wish to highlight potential risks in our work. Specifically, the use
of future information in PGIF may exacerbate biases present in the experience
or offline dataset. For example, it is known that expert datasets lack action
diversity. Further conditioning on the future in this dataset could force
these biases to be incorporated more easily into the learning algorithm. Some
other biases that may arise are a trajectory that contains some marginal
reward and is therefore incorporated into our policy/value function with PGIF.
This could hamper exploration of the agent and prevent discovery of states
that yield optimal reward. Furthermore, we find that our method is much slower
to train than the baselines we compare to, due to the fact that the
architecture requires training an LSTM, with at times, very long trajectories
as input. We must also use the entire trajectories to train this LSTM rather
than transition batches from our replay buffer, further increasing the
training complexity. Investigating the transformer architecture rather than an
LSTM would be an interesting future direction.

## References

  * Fra [2019] Combined reinforcement learning via abstract representations.  _Proceedings of the AAAI Conference on Artificial Intelligence_ , 2019\.
  * [2] Zaheer Abbas, Samuel Sokota, Erin Talvitie, and Martha White.  Selective dyna-style planning under limited model capacity.  In _Proceedings of the 37th International Conference on Machine Learning_.
  * Bayer and Osendorfer [2015] Justin Bayer and Christian Osendorfer.  Learning stochastic recurrent networks, 2015.
  * Bengio et al. [2015] Samy Bengio, Oriol Vinyals, Navdeep Jaitly, and Noam Shazeer.  Scheduled sampling for sequence prediction with recurrent neural networks.  In _Proceedings of the 28th International Conference on Neural Information Processing Systems_ , 2015.
  * Buckman et al. [2018] Jacob Buckman, Danijar Hafner, George Tucker, Eugene Brevdo, and Honglak Lee.  Sample-efficient reinforcement learning with stochastic ensemble value expansion.  In _Advances in Neural Information Processing Systems_ , 2018.
  * Bush and Pineau [2009] Keith Bush and Joelle Pineau.  Manifold embeddings for model-based reinforcement learning under partial observability.  In _Advances in Neural Information Processing Systems_ , 2009.
  * Florensa et al. [2017] Carlos Florensa, David Held, Markus Wulfmeier, Michael Zhang, and Pieter Abbeel.  Reverse curriculum generation for reinforcement learning.  _Proceedings of Machine Learning Research_ , 2017.
  * Fu et al. [2020] Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine.  D4rl: Datasets for deep data-driven reinforcement learning, 2020.
  * Fujimoto et al. [2018] Scott Fujimoto, Herke van Hoof, and David Meger.  Addressing function approximation error in actor-critic methods.  In _Proceedings of the 35th International Conference on Machine Learning_ , 2018.
  * Fujimoto et al. [2019] Scott Fujimoto, David Meger, and Doina Precup.  Off-policy deep reinforcement learning without exploration.  In _Proceedings of the 36th International Conference on Machine Learning_ , 2019.
  * Goyal et al. [2017] Anirudh Goyal, Alessandro Sordoni, Marc-Alexandre Côté, Nan Rosemary Ke, and Yoshua Bengio.  Z-forcing: Training stochastic recurrent networks.  In _Proceedings of the 31st International Conference on Neural Information Processing Systems_ , 2017.
  * Goyal et al. [2021] Anirudh Goyal, Alex Lamb, Jordan Hoffmann, Shagun Sodhani, Sergey Levine, Yoshua Bengio, and Bernhard Schölkopf.  Recurrent independent mechanisms.  In _International Conference on Learning Representations_ , 2021.
  * Guez et al. [2020] Arthur Guez, Fabio Viola, Theophane Weber, Lars Buesing, Steven Kapturowski, Doina Precup, David Silver, and Nicolas Heess.  Value-driven hindsight modelling.  In _Advances in Neural Information Processing Systems 33: Annual Conference_ , 2020.
  * Haarnoja et al. [2018] Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine.  Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor.  In _Proceedings of the 35th International Conference on Machine Learning_ , 2018.
  * Harutyunyan et al. [2019] Anna Harutyunyan, Will Dabney, Thomas Mesnard, Mohammad Gheshlaghi Azar, Bilal Piot, Nicolas Heess, Hado P van Hasselt, Gregory Wayne, Satinder Singh, Doina Precup, and Remi Munos.  Hindsight credit assignment.  In _Advances in Neural Information Processing Systems_ , 2019.
  * Janner et al. [2021] Michael Janner, Qiyang Li, and Sergey Levine.  Reinforcement learning as one big sequence modeling problem.  _arXiv preprint arXiv:2106.02039_ , 2021.
  * Johnson-Laird [2010] Philip N. Johnson-Laird.  Mental models and human reasoning.  _Proceedings of the National Academy of Sciences_ , 2010.
  * Karl et al. [2017] M. Karl, Maximilian Sölch, J. Bayer, and P. V. D. Smagt.  Deep variational bayes filters: Unsupervised learning of state space models from raw data.  _ArXiv_ , abs/1605.06432, 2017.
  * Ke et al. [2019] Nan Rosemary Ke, Amanpreet Singh, Ahmed Touati, Anirudh Goyal, Yoshua Bengio, Devi Parikh, and Dhruv Batra.  Modeling the long term future in model-based reinforcement learning.  In _International Conference on Learning Representations_ , 2019.
  * Kumar et al. [2019] Aviral Kumar, Justin Fu, Matthew Soh, George Tucker, and Sergey Levine.  Stabilizing off-policy q-learning via bootstrapping error reduction.  In _Advances in Neural Information Processing Systems_ , 2019.
  * Lamb et al. [2016] Alex M Lamb, Anirudh Goyal, Ying Zhang, Saizheng Zhang, Aaron C Courville, and Yoshua Bengio.  Professor forcing: A new algorithm for training recurrent networks.  In _Advances in Neural Information Processing Systems_ , 2016.
  * Mesnard et al. [2021] Thomas Mesnard, Theophane Weber, Fabio Viola, Shantanu Thakoor, Alaa Saade, Anna Harutyunyan, Will Dabney, Thomas S Stepleton, Nicolas Heess, Arthur Guez, Eric Moulines, Marcus Hutter, Lars Buesing, and Remi Munos.  Counterfactual credit assignment in model-free reinforcement learning.  In _Proceedings of the 38th International Conference on Machine Learning_ , 2021.
  * [23] Volodymyr Mnih, Adria Puigdomenech Badia, Mehdi Mirza, Alex Graves, Timothy Lillicrap, Tim Harley, David Silver, and Koray Kavukcuoglu.  Asynchronous methods for deep reinforcement learning.  In _Proceedings of The 33rd International Conference on Machine Learning_.
  * Nachum et al. [2019] Ofir Nachum, Bo Dai, Ilya Kostrikov, Yinlam Chow, Lihong Li, and Dale Schuurmans.  Algaedice: Policy gradient from arbitrary experience, 2019.
  * Nota et al. [2021] Chris Nota, Philip Thomas, and Bruno C. Da Silva.  Posterior value functions: Hindsight baselines for policy gradient methods.  In _Proceedings of the 38th International Conference on Machine Learning_ , 2021.
  * Oh et al. [2017] Junhyuk Oh, Satinder Singh, and Honglak Lee.  Value prediction network.  In _Advances in Neural Information Processing Systems_ , 2017.
  * Osband et al. [2020] Ian Osband, Yotam Doron, Matteo Hessel, John Aslanides, Eren Sezener, Andre Saraiva, Katrina McKinney, Tor Lattimore, Csaba Szepesvári, Satinder Singh, Benjamin Van Roy, Richard Sutton, David Silver, and Hado van Hasselt.  Behaviour suite for reinforcement learning.  In _International Conference on Learning Representations_ , 2020.
  * Ostrovski et al. [2017] Georg Ostrovski, Marc G. Bellemare, Aäron van den Oord, and Rémi Munos.  Count-based exploration with neural density models.  In _Proceedings of the 34th International Conference on Machine Learning_ , 2017.
  * Pascanu et al. [2017] Razvan Pascanu, Yujia Li, Oriol Vinyals, Nicolas Heess, Lars Buesing, Sebastien Racanière, David Reichert, Théophane Weber, Daan Wierstra, and Peter Battaglia.  Learning model-based planning from scratch, 2017.
  * Pathak et al. [2017] Deepak Pathak, Pulkit Agrawal, Alexei A. Efros, and Trevor Darrell.  Curiosity-driven exploration by self-supervised prediction.  In _Proceedings of the 34th International Conference on Machine Learning_ , 2017.
  * Peng et al. [2018] Baolin Peng, Xiujun Li, Jianfeng Gao, Jingjing Liu, and Kam-Fai Wong.  Deep Dyna-Q: Integrating planning for task-completion dialogue policy learning.  In _Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics)_ , 2018.
  * Racanière et al. [2017] Sébastien Racanière, Theophane Weber, David Reichert, Lars Buesing, Arthur Guez, Danilo Jimenez Rezende, Adrià Puigdomènech Badia, Oriol Vinyals, Nicolas Heess, Yujia Li, Razvan Pascanu, Peter Battaglia, Demis Hassabis, David Silver, and Daan Wierstra.  Imagination-augmented agents for deep reinforcement learning.  In _Advances in Neural Information Processing Systems_ , 2017.
  * Saxe et al. [2018] Andrew Michael Saxe, Yamini Bansal, Joel Dapello, Madhu Advani, Artemy Kolchinsky, Brendan Daniel Tracey, and David Daniel Cox.  On the information bottleneck theory of deep learning.  In _International Conference on Learning Representations_ , 2018.
  * Schmidhuber [2020] Juergen Schmidhuber.  Reinforcement learning upside down: Don’t predict rewards – just map them to actions, 2020.
  * Schulman et al. [2016] John Schulman, Philipp Moritz, Sergey Levine, Michael Jordan, and Pieter Abbeel.  High-dimensional continuous control using generalized advantage estimation.  In _Proceedings of the Fourth International Conference of Learning Representations_ , 2016.
  * Schulman et al. [2017] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov.  Proximal policy optimization algorithms.  _arXiv_ , 1707.06347, 2017.
  * Singh [1994] Satinder P. Singh.  Learning without state-estimation in partially observable markovian decision processes.  In _In Proceedings of the Eleventh International Conference on Machine Learning_ , 1994.
  * Suematsu and Hayashi [1999] Nobuo Suematsu and Akira Hayashi.  A reinforcement learning algorithm in partially observable environments using short-term memory.  In _Advances in Neural Information Processing Systems_ , 1999.
  * Sutton [1988] Richard S. Sutton.  Learning to predict by the methods of temporal differences.  _Machine Learning_ , 1988.
  * Sutton [1991] Richard S. Sutton.  Dyna, an integrated architecture for learning, planning, and reacting.  _SIGART Bull._ , 1991.
  * Sutton and Barto [2018] Richard S. Sutton and Andrew G. Barto.  _Reinforcement Learning: An Introduction_.  A Bradford Book, 2018.  ISBN 0262039249.
  * Sutton et al. [1999] Richard S. Sutton, David McAllester, Satinder Singh, and Yishay Mansour.  Policy gradient methods for reinforcement learning with function approximation.  In _Proceedings of the 12th International Conference on Neural Information Processing Systems_ , 1999.
  * Tishby et al. [1999] Naftali Tishby, Fernando C. Pereira, and William Bialek.  The information bottleneck method.  1999\.
  * Todorov et al. [2012] E. Todorov, T. Erez, and Y. Tassa.  MuJoCo: A physics engine for model-based control.  In _2012 IEEE/RSJ International Conference on Intelligent Robots and Systems_ , 2012.
  * Tsitsiklis and Van Roy [1996] John N. Tsitsiklis and Benjamin Van Roy.  Analysis of temporal-difference learning with function approximation.  In _Proceedings of the 9th International Conference on Neural Information Processing Systems_ , 1996.
  * van Hasselt et al. [2018] Hado van Hasselt, Yotam Doron, Florian Strub, Matteo Hessel, Nicolas Sonnerat, and Joseph Modayil.  Deep reinforcement learning and the deadly triad, 2018.
  * Vaswani et al. [2017] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  In _Advances in neural information processing systems_ , pages 5998–6008, 2017.
  * Wang et al. [2020] Tingwu Wang, Xuchan Bao, Ignasi Clavera, Jerrick Hoang, Yeming Wen, Eric Langlois, Shunshi Zhang, Guodong Zhang, Pieter Abbeel, and Jimmy Ba.  Benchmarking model-based reinforcement learning, 2020.
  * Weaver and Tao [2001] Lex Weaver and Nigel Tao.  The optimal reward baseline for gradient-based reinforcement learning.  In _Proceedings of the Seventeenth Conference on Uncertainty in Artificial Intelligence_ , 2001.
  * Wu et al. [2020] Yifan Wu, George Tucker, and Ofir Nachum.  Behavior regularized offline reinforcement learning, 2020.  URL https://openreview.net/forum?id=BJg9hTNKPH.
  * Wulfmeier et al. [2021] Markus Wulfmeier, Dushyant Rao, Roland Hafner, Thomas Lampe, Abbas Abdolmaleki, Tim Hertweck, Michael Neunert, Dhruva Tirumala, Noah Siegel, Nicolas Heess, and Martin Riedmiller.  Data-efficient hindsight off-policy option learning, 2021.
  * Yang and Nachum [2021] Mengjiao Yang and Ofir Nachum.  Representation matters: Offline pretraining for sequential decision making, 2021.
  * Zhang et al. [2021] Chi Zhang, Sanmukh Rao Kuppannagari, and Viktor Prasanna.  {BRAC}+: Going deeper with behavior regularized offline reinforcement learning, 2021.  URL https://openreview.net/forum?id=bMCfFepJXM.
  * Łukasz Kaiser et al. [2020] Łukasz Kaiser, Mohammad Babaeizadeh, Piotr Miłos, Błażej Osiński, Roy H Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, Afroz Mohiuddin, Ryan Sepassi, George Tucker, and Henryk Michalewski.  Model based reinforcement learning for atari.  In _International Conference on Learning Representations_ , 2020.

##  Appendix A Implementation Considerations

We will introduce some implementation details for our model, specifically how
the parameters for each distribution are learned.

Data Collection Model: In the first phase of the actor critic process we are
in a _generation_ phase of the data. Here, we have a prior over the latent
variables $p_{\upsilon^{{\text{Z}}}}(z_{t})$. This prior distribution can be
defined as a Gaussian distribution with parameters as
$p_{\upsilon^{\text{Z}}}(z_{t})=\mathcal{N}(z_{t};\mu^{\text{\text{Z}}}_{t},\log\sigma^{\text{Z}}_{t})$.
To learn the parameters of this Gaussian prior we utilize an MLP ($f$) defined
by $[\mu_{t}^{\text{Z}},\log\sigma^{\text{\text{Z}}}_{t}]=f^{\text{Z}}(s_{t})$
which takes the current state as input.

Now that we have the learned parameters, we can sample $z_{t}\sim
p_{\upsilon^{\text{Z}}}(z_{t})$ using the reparameterization trick given the
learned mean and standard deviation. We denote this as
$z_{t}\sim\text{reparameterize}(\mu_{t}^{\text{\text{Z}}},\log\sigma^{\text{Z}}_{t})$.
At each time-step during the generation phase, we encounter a state from our
environment. We find the distribution parameters and then sample $z_{t}$ as
denoted previously. The latent variable policy is then executed:
$a_{t}\sim\pi_{\theta}(a_{t}|s_{t},z_{t})$. By applying action $a_{t}$ on the
environment we then obtain $s_{t+1},r_{t}$, and repeat until termination. This
leaves us with a sequence of states for use in inference. We denote the full
trajectories generated in this phase with, $\tau$, and
$\tau_{(s)}=\\{s_{0},\dots,s_{T}\\}$ being the trajectories only containing
states. Let us define the set of state only trajectories as,
$\\{\tau_{(s)}^{0},\dots,\tau_{(s)}^{n}\\}\in\mathcal{D}_{(s)}$.

Training Model: The 2nd part of our actor critic algorithm where we learn the
policy using the data in the inference phase. This is where we approximate the
true posterior distribution over our latent variables. As in Z-forcing, we
will use a _backwards_ RNN, which takes the reversed sequence of states,
$\tau_{(s)}^{\leftarrow}$ as input. We define this network as,

$b_{t}=f^{\leftarrow}(s_{t+1},b_{t+1}).$ (11)

Each state $b_{t}$ therefore contains information about future states and can
be used to shape the approximate posterior distribution over latent variables
to contain this information. We define the inference network with a normal
distribution,

$q_{\phi^{\text{Z}}}(z_{t}|b_{t})=\mathcal{N}(z_{t};\mu_{t}^{\text{\text{q}}},\log\sigma_{t}^{\text{q}}),$
(12)

where, as before, the parameters are learned with an MLP defined by
$[\mu_{t}^{{\text{q}}},\log\sigma_{t}^{\text{q}}]=f^{\text{\text{q}}}(b_{t})$.

We derive a conditional generative model $p_{\zeta}(b|z)$ over the backwards
states given the inferred latent variables $z_{t}\sim
q_{\phi^{\text{\text{Z}}}}(z_{t}|b_{t})$. Similar to before, we define the
parameters of this Gaussian as
$p_{\zeta}(b_{t}|z_{t})=\mathcal{N}(z_{t};\mu_{t}^{\text{Ax}},\log\sigma_{t}^{\text{Ax}})$,
with a MLP that produces
$[\mu_{t}^{\text{Ax}},\log\sigma_{t}^{\text{Ax}}]=f^{\text{Ax}}(z_{t})$.

##  Appendix B Full Algorithm

We show the full algorithms for PGIF-PPO (Algorithm 2) with State based
forcing and PGIF-SAC (Algorithm 4) with state based forcing in this section.
Using VPN forcing simply changes the auxiliary loss calculations. The shared
Backwards RNN Pass algorithm is in Algorithm 3.

Algorithm 2 Algorithm with Advantage Policy Gradient and State Based Forcing

0: Initial policy parameters: $\theta$, Prior parameters:
$\upsilon^{{\text{Z}}},\phi^{{\text{Z}}}$, KL weight: $\beta$, Auxiliary loss
Parameters: $\theta^{\text{Ax}}$, Auxiliary loss weight: $\alpha$ Backwards
net input parameters: $\theta^{\text{in}}$, Backwards RNN parameters:
$\theta^{\text{bkw}}$, Backwards net output parameters: $\theta^{{\text{bkw-
out}}}$, Initial hidden states $\mathbf{h_{i}}$.

1: for $\text{policy-step }k=0,1,2,\dots,N$ do

2: Collect set of Trajectories $\mathcal{D}=\\{\tau_{i},\dots\\}:$

3: repeat

4:
$[\mu_{t}^{\text{Z}},\log\sigma^{\text{Z}}_{t}]=f_{\upsilon^{\text{Z}}}(s_{t})$

5:
$z_{t}\sim\text{reparameterization}(\mu_{t}^{\text{Z}},\log\sigma^{\text{Z}}_{t})$

6: Execute: $\pi_{\theta}(a_{t}|s_{t},z_{t})=f_{\theta}(s_{t},z_{t})$

7: Observe $s_{t+1},r_{t}$ from environment.

8: $\tau_{i}=\tau_{i}\cup\\{s_{t},a_{t},s_{t+1},r_{t}\\}$

9: until episode termination

10: Compute advantage estimates $\hat{A}_{t}^{k}$ using GAE.

11: for $\text{Trajectory: }\tau_{i}\in\mathcal{D}$ do

12: $\mathbf{b_{i}}$, $\mathbf{h_{i}}$
$=\text{BackwardsPass}(\tau_{i},\mathbf{h}_{{i}},\theta^{(in)},\theta^{\text{bkw}},\theta^{{\text{bkw-
out}}})$

13: $\tau_{i}=\tau_{i}\cup\\{\mathbf{b_{i}}\\}$

14: Reset: $J_{\text{PG}}=0$

15: for timestep $t\in\\{0,\dots,T\\}$ do

16: $\forall\\{s_{t},b_{t},a_{t}\\}\in\mathcal{D}$:
$[\mu_{t}^{\text{Z}},\log\sigma_{t}^{\text{Z-PG}}]=f_{\phi^{\text{Z}}}(b_{t},s_{t})$

17:
$z_{t}^{\text{PG}}\sim\text{reparameterization}(\mu_{t}^{\text{Z-PG}},\log\sigma^{{\text{Z-PG}}}_{t})$

18:
$\pi_{\theta}(a_{t}|s_{t},z_{t}^{\text{PG}})=f_{\theta}(s_{t},z_{t}^{\text{PG}})$

19:
$[\mu_{t}^{\text{Ax}},\log\sigma_{t}^{\text{Ax}}]=f_{\theta^{\text{Ax}}}(z_{t}^{\text{PG}})$

20:
$J_{\text{KL}}=\text{KLDivergence}(\mu_{t}^{\text{Z-PG}},\log\sigma^{{\text{Z-PG}}}_{t},\mu_{t}^{\text{Z}},\log\sigma^{\text{Z}}_{t})$

21:
$J_{\text{Ax}}=\text{LogProbGaussian}(b_{t},\mu_{t}^{\text{Ax}},\log\sigma_{t}^{\text{Ax}})$

22:
$J_{\text{PG}}\stackrel{{\scriptstyle+}}{{=}}\mathbb{E}_{\tau_{i}\in\mathcal{D}}\Big{[}\log\pi_{\theta}(a_{t}|s_{t},z_{t}^{\text{PG}})\hat{A}^{k}_{t}\Big{]}-\alpha
J_{\text{Ax}}-\beta J_{\text{KL}}$

23: Update all parameters w.r.t: $J_{\text{PG}}$

24: Update value function estimates with any method.

Algorithm 3 BackwardsPass

0: Trajectory: $\tau$, Hidden states $h$, Input MLP Parameters:
$\theta^{\text{in}}$, Backwards RNN Parameters: $\theta^{\text{bkw}}$, Output
net parameters: $\theta^{{\text{bkw-out}}}$.

1: $\tau^{-1}=\text{Reverse-Order}(\tau)$

2: $x_{\text{bkw}}=f_{\theta^{\text{in}}}(\tau^{-1})$

3: $h_{\text{bkw}}=f_{\theta^{\text{bkw}}}(x_{\text{bkw}},h)$ (Operates over
the entire sequence)

4: $x_{\text{bkw-out}}=f_{\theta^{\text{bkw-out}}}(h_{\text{bkw}})$,

5: $x_{\text{bkw-out}}=\text{Reverse-Order}(x_{{\text{bkw-out}}})$

6: $h_{\text{bkw}}=\text{Reverse-Order}(h_{{\text{bkw}}}$)

7: $\textbf{return }x_{\text{bkw-out}},h_{\text{bkw}}$

Algorithm 4 SAC Algorithm with State Based Forcing

0: Initial policy parameters: $\theta$, Z forcing parameters:
$\upsilon^{{\text{Z}}},\upsilon^{\text{U}},\phi^{{\text{Z}}},\phi^{\text{U}}$,
KL weight: $\beta$, Auxiliary loss Parameters: $\theta^{\text{Ax-
PG}},\theta^{\text{Ax-TD}}$, Auxiliary loss weight: $\alpha$ Backwards net
input parameters: $\theta^{\text{in}}$, Backwards RNN parameters:
$\theta^{\text{bkw-Z}},\theta^{\text{bkw-U}}$, Backwards net output
parameters: $\theta^{\text{bkw-out-Z}},\theta^{\text{bkw-out-U}}$, Initial
hidden states $\mathbf{h_{i}}^{\text{Z}},\mathbf{h_{i}}^{\text{U}}$.

1: for $\text{policy-step }k=0,1,2,\dots,N$ do

2: Collect set of Trajectories $\mathcal{D}=\\{\tau_{i},\dots\\}:$

3: repeat

4:
$[\mu_{t}^{\text{Z}},\log\sigma^{\text{Z}}_{t}]=f_{\upsilon^{\text{Z}}}(s_{t})$

5:
$z_{t}\sim\text{reparameterization}(\mu_{t}^{\text{Z}},\log\sigma^{\text{Z}}_{t})$

6: Execute: $\pi_{\theta}(a_{t}|s_{t},z_{t})=f_{\theta}(s_{t},z_{t})$

7: Observe $s_{t+1},r_{t}$ from environment.

8: $\tau_{i}=\tau_{i}\cup\\{s_{t},a_{t},s_{t+1},r_{t}\\}$

9: until episode termination

10: for $\text{Trajectory: }\tau_{i}\in\mathcal{D}$ do

11: $\mathbf{b_{i}}^{\text{Z}}$, $\mathbf{h_{i}}^{\text{Z}}$
$=\text{BackwardsPass}^{\text{Z}}(\tau_{i},\mathbf{h}_{{i}}^{\text{Z}},\theta^{\text{in-Z}},\theta^{\text{bkw-Z}},\theta^{\text{bkw-
out-Z}})$

12: $\mathbf{b_{i}}^{\text{U}}$, $\mathbf{h_{i}}^{\text{U}}$
$=\text{BackwardsPass}^{\text{U}}(\tau_{i},\mathbf{h}_{{i}}^{\text{U}},\theta^{\text{in-U}},\theta^{\text{bkw-U}},\theta^{\text{bkw-
out-U}})$

13:
$\tau_{i}=\tau_{i}\cup\\{\mathbf{b_{i}}^{\text{U}},\mathbf{b_{i}}^{\text{Z}}\\}$

14: Reset: $J_{\text{PG}},J_{\text{Q}}=0$

15: for timestep $t\in\\{0,\dots,T\\}$ do

16:
$\forall\\{s_{t},b_{t},a_{t},s_{t+1},b^{\text{Z}}_{t},b^{\text{U}}_{t}\\}\in\mathcal{D}$

17: Compute $a_{t+1}$

18:
$[\mu_{t+1}^{\text{Z}},\log\sigma_{t+1}^{\text{Z}}]=f_{\upsilon^{\text{Z}}}(s_{t+1})$

19:
$z_{t+1}\sim\text{reparameterization}(\mu_{t+1}^{\text{Z}},\log\sigma^{{\text{Z}}}_{t+1})$

20:
$a_{t+1}\sim\pi_{\theta}(a_{t+1}|s_{t+1},z_{t+1})=f_{\theta}(s_{t+1},z_{t+1})$

21: Compute target Q (U-Tar to denote distribution parameters)

22:
$[\mu_{t+1}^{\text{U-Tar}},\log\sigma_{t+1}^{\text{U-Tar}}]=f_{\upsilon^{\text{U}}}(s_{t+1})$

23:
$u_{t+1}^{\text{Tar}}\sim\text{reparameterization}(\mu_{t+1}^{\text{U-Tar}},\log\sigma^{\text{U-Tar}}_{t+1})$

24:
$y(r_{t},s_{t+1},d)=r_{t}+\gamma(1-d)\Big{(}Q_{\psi\text{-target}}(s_{t+1},a_{t+1},u_{t+1}^{\text{Tar}})-\alpha\log\pi_{\theta}(a_{t+1}|s_{t+1},z_{t+1}))\Big{)}$

25: Update Q functions

26:
$[\mu_{t}^{\text{U}},\log\sigma_{t}^{\text{U}}]=f_{\phi^{\text{U}}}(b_{t}^{\text{Q}},s_{t})$

27:
$u_{t}\sim\text{reparameterization}(\mu_{t+1}^{\text{U}},\log\sigma^{{\text{U}}}_{t+1})$

28: $J_{\text{KL-
TD}}=\text{KLDivergence}(\mu_{t}^{\text{U}},\log\sigma^{\text{U}}_{t},\mu_{t}^{\text{U-Tar}},\log\sigma^{\text{\text{U-Tar}}}_{t})$

29: $[\mu_{t}^{\text{Ax-TD}},\log\sigma_{t}^{\text{Ax-
TD}}]=f_{\theta^{\text{Ax-TD}}}(u_{t}^{\text{}})$

30: $J_{\text{Ax-
TD}}=\text{LogProbGaussian}(b_{t}^{\text{Q}},\mu_{t}^{\text{Ax-
TD}},\log\sigma^{\text{Ax-TD}}_{t})$

31:
$J_{Q}\stackrel{{\scriptstyle+}}{{=}}(Q_{\psi}(s_{t},a_{t},u_{t})-y(r_{t},s_{t+1},d))+\alpha
J_{\text{Ax-TD}}+\beta J_{\text{KL-TD}}$

32: Update Target Networks

33: $(\psi-\text{target})\leftarrow\rho(\psi-\text{target})+(1-\rho)\psi$

34: Update Policy

35:
$[\mu_{t}^{\text{Z-PG}},\log\sigma_{t}^{\text{Z-PG}}]=f_{\phi^{\text{Z}}}(b_{t}^{\text{Z}},s_{t})$

36:
$z_{t}^{\text{PG}}\sim\text{reparameterization}(\mu_{t}^{\text{Z-PG}},\log\sigma^{{\text{Z-PG}}}_{t})$

37:
$a_{t}\sim\pi_{\theta}(a_{t}|s_{t},z_{t}^{\text{PG}})=f_{\theta}(s_{t},z_{t}^{\text{PG}})$

38: $[\mu_{t}^{\text{Ax-PG}},\log\sigma_{t}^{\text{Ax-
PG}}]=f_{\theta^{\text{Ax-PG}}}(z_{t}^{\text{PG}})$

39: $J_{\text{KL-
PG}}=\text{KLDivergence}(\mu_{t}^{\text{Z-PG}},\log\sigma^{{\text{Z-PG}}}_{t},\mu_{t}^{\text{Z}},\log\sigma^{\text{Z}}_{t})$

40: $J_{\text{Ax-
PG}}=\text{LogProbGaussian}(b_{t}^{\text{Z}},\mu_{t}^{\text{Ax-
PG}},\log\sigma_{t}^{\text{Ax-PG}})$

41:
$J_{\text{PG}}\stackrel{{\scriptstyle+}}{{=}}\mathbb{E}_{\tau_{i}\in\mathcal{D}}\Big{[}Q_{\psi}(s_{t},a_{t},u_{t}^{\text{}})-\alpha\log\pi_{\theta}(a_{t}|s_{t},z_{t}^{\text{PG}})\Big{]}+\alpha
J_{\text{Ax-PG}}+\beta J_{\text{KL-PG}}$

42: Update all parameters w.r.t: $J_{\text{PG}}$ or $J_{\text{Q}}$

43: Update value function estimates with any method.

##  Appendix C Experimental Parameters

In this section, we give the hyperparameters used for each of our experiments
in Tables 4, 5, and 6.

Parameter | Value
---|---
Optimizer | Adam
Learning rate | $5e^{-4}$
Batch size | $250$
Actor and Critic network dimensions | $(300,200)$
RNN Dim | 15
RNN-Embedding Dim | 20
Z or Q Dim | 5
Z-force Neural net dim | 20
Initial random exploration steps | 10000
Replay Buffer Size | 1,000,000 steps
Discount | 0.99
Evaluation Episodes | 5
Table 4: Parameters used for PGIF and SAC in the online experiments with MuJoCo. Parameter | Value
---|---
Optimizer | Adam
Learning rate | $1e^{-3}$
Value penalty | True
Batch size | $250$
Actor and Critic network dimensions | $(300,200)$
Q value ensemble | $2$
RNN Dim | $15$
RNN-Embedding Dim | $20$
Z or Q Dim | $10$
Z-force Neural net dim | $20$
Initial random exploration steps | $10000$
Replay Buffer Size |  $1,000,000$ steps
Discount | $0.99$
Evaluation Episodes | $5$
BRAC Value Penalty $\alpha$ | $0.1$
Table 5: Parameters used for transformers experiments Parameter | Value
---|---
Hidden Dim | 128
Activation | ReLu
Heads | 1
Heads | 1
Layers | 3
Attention Dropout | 0.1
Table 6: Parameters used for the transformers experiments.

Training the behavior policy for BRAC offline RL experiments uses behavior
cloning with 300,000 time steps.

For the purpose of the backwards RNN, since we have variable episode length we
pad the episode sequence to the maximum length with zeros. The maximum number
of timesteps is generally $1000$ in MuJoCo. The training of the RNN with this
long sequence and the pre-processing steps involved are expensive
computationally. For offline experiments, our replay buffer is episodic, so we
train with an entire episode at each update. For online we train with a batch
of randomly sampled steps from the replay buffer (and obtain the respective
episodes that those steps were in for the training of the backwards RNN).

For the weight parameters $\alpha$ and $\beta$, we recommend that the weight
on the KL be increased at a small fixed constant rate during training so that
the final returned policy uses minimal information about the future. We use an
initial weight of $1e^{-4}$ for the KL and auxiliary losses and increase it by
$1e^{-6}$ at each training iteration for experiments to a max weight of $1$.
Thus, at the end of training the policy and value functions should be
minimally reliant on this privileged information.

##  Appendix D Online Experiment Learning Curves

In this section, we show the training curves for the online experiments with
full observability in Figure 4. We compare PGIF against SAC only, since we see
that PPO performs much worse from initial experiments. Note that in
practicality, it is best to reduce the maximum number of steps in the MuJoCo
environments to 500 for better computational efficiency in the backwards RNN
steps (without much change in final reward). We note that our method appears
to have a small reduction in standard error compared to vanilla SAC.

Figure 4: The online episodic mean reward evaluated over 5 episodes every 500
steps for MuJoCo continuous control RL tasks with full observability. We show
the average over $5$ random seeds. $500,000$ environment step interactions are
used. The shaded area shows the standard error.

We aim to investigate the effect PGIF learning with a lower batch size. We
notice that PGIF has better performance at lower batch sizes. It may be an
interesting direction to further investigate this anomaly. In Figure 5, we
show the performance of PGIF with a batch size of 25.

Figure 5: The online episodic mean reward evaluated over 5 episodes every 500
steps for MuJoCo tasks with full observability. We show the average over $5$
random seeds. $500,000$ environment step interactions are used. The model is
updated with a batch size of 25. The shaded area shows the standard error.

##  Appendix E Computing Infrastructure

The cluster used to run these experiments has $688$ NVIDIA V100-SXM2 GPUs.