# Do Transformer World Models Give
Better Policy Gradients?

Michel Ma, Tianwei Ni, Clement Gehring, Pierluca D’Oro11footnotemark: 1,
Pierre-Luc Bacon
Mila, Université de Montréal
{michel.ma, tianwei.ni<EMAIL_ADDRESS>
{pierluca.doro<EMAIL_ADDRESS>Equal contribution.

###### Abstract

A natural approach for reinforcement learning is to predict future rewards by
unrolling a neural network world model, and to backpropagate through the
resulting computational graph to learn a policy. However, this method often
becomes impractical for long horizons since typical world models induce hard-
to-optimize loss landscapes. Transformers are known to efficiently propagate
gradients over long horizons: could they be the solution to this problem?
Surprisingly, we show that commonly-used transformer world models produce
_circuitous gradient paths_ , which can be detrimental to long-range policy
gradients. To tackle this challenge, we propose a class of world models called
Actions World Models (AWMs), designed to provide more direct routes for
gradient propagation. We integrate such AWMs into a policy gradient framework
that underscores the relationship between network architectures and the policy
gradient updates they inherently represent. We demonstrate that AWMs can
generate optimization landscapes that are easier to navigate even when
compared to those from the simulator itself. This property allows transformer
AWMs to produce better policies than competitive baselines in realistic long-
horizon tasks.

###  1 Introduction

Given a class of parameterized policies, policy optimization methods aim to
find parameters ${\bm{\theta}}$ that maximize a performance criterion $J$ –
typically the expected sum of rewards over the episodes of experiences
generated by an agent in an environment. When the environment dynamics are
known and differentiable, the gradient of the objective $J$ can be directly
evaluated through automatic differentiation and used to optimize $J$ by
gradient ascent. If the dynamics are unknown, a differentiable world model can
be trained on interaction data, whose predicted states will be backpropagated
through. This model-based approach to reinforcement learning (RL) has been
explored early on in the field with Werbos (1974); Schmidhuber (1990); Miller
et al. (1995) and led to contemporary work in deep RL such as Hafner et al.
(2022; 2023). While backpropagation through time Werbos (1990) using a world
model is conceptually appealing and readily compatible with deep learning
practice, most attempts to scale it up have been unsuccessful beyond a few
dozen steps Hafner et al. (2022); Ghugare et al. (2023). Indeed, even in the
presence of perfectly accurate differentiable simulators, policy gradients
obtained by backpropagation through time can quickly become unstable over long
horizons Suh et al. (2022); Metz et al. (2021).

(a) Markovian world model.

(b) History World Model.

(c) Actions World Model.

Figure 1: Diagram illustrating gradient flows through different world model
types from states to actions. Circuitous (longer than necessary) gradient
paths go through connections highlighted in red. An Actions World Model has no
circuitous gradient paths, allowing gradients to directly flow from states to
actions through a single application of a world model.

This problem of stability in differentiating through world models over long
horizons echoes similar challenges observed in sequence modeling. The latter
field has experienced numerous advances in deep network architectures
specifically aimed at alleviating this challenge. As a notable example,
attention-based models such as transformers were designed to directly
propagate gradients from any output to any input arbitrarily far in the past,
within their context length Schmidhuber (1992); Vaswani et al. (2017). It is,
therefore, natural to ponder whether using a transformer as a world model and
differentiating through it could overcome the difficulties encountered in
model-based policy optimization. Or, in other words, to ask the question: _do
transformer world models give better policy gradients?_ In this paper, we
provide a new answer to this question: surprisingly, transformers do not
necessarily confer the same benefit for policy optimization when used as world
models conditioned on the full history of states and actions. Instead,
conditioning solely on a sequence of actions tends to yield better-behaved
policy gradients. This, in turn, leads to improved temporal credit assignment
and policy performance.

We first show that backpropagating naively through a transformer predicting
the next state conditioned on the full history Micheli et al. (2022); Chen et
al. (2022); Robine et al. (2023) does not necessarily lead to better policy
gradients in long-horizon tasks. We demonstrate, conceptually and
theoretically, that this phenomenon can be attributed to the existence of
_circuitous gradient paths_ : long state-to-state gradient paths generated
during the autoregressive unrolling of the model (see Figure 1).

To avoid the creation of those longer-than-necessary gradient paths, we
propose an alternative class of world models, which we call _Actions World
Models_ (AWMs), of the form
${\hat{s}_{t+1}=g(s_{1},a_{1},\dots,a_{t-1},a_{t})}$ conditioned solely on the
history of past actions and an initial state. An AWM is forced to create its
own internal dynamics, potentially very different from the one of the
environment, while being able to predict future states. This inductive bias
constrains the temporal accumulation of gradients to only happen within the
neural network of the AWM, fully harnessing its gradient properties.

We show theoretically that an AWM directly inherits the gradient propagation
properties of the underlying network architecture. Thus, AWMs yield policy
gradients subject to bounds similar to the ones previously derived in deep
sequence models Pascanu et al. (2013b); Kerg et al. (2020). Additionally, our
theoretical framework provides a unified perspective on many backpropagation-
based policy optimization methods by exposing the relationship between the
network architecture and the policy gradient updates that they induce. For
instance, it allows us to cast backpropagation through unrolled Markovian
dynamics Heess et al. (2015) as a specific AWM instantiated with a vanilla
recurrent neural network. Our framework paves the way for new algorithmic
developments by fostering novel synergies between the study of network
architectures in deep learning and the policy optimization updates they
implicitly express.

Finally, through a series of experiments, we showcase the remarkable empirical
properties of backpropagation-based policy optimization with AWMs. We
demonstrate in illustrative domains with chaotic or non-differentiable
dynamics that AWMs can give better policy gradients than the simulator itself,
by forcing the policy optimization landscape to be easy to optimize. This
leads to successful policy optimization via backpropagation even when the same
procedure would fail with the true underlying model. On a testbed of realistic
long-horizon tasks Howe et al. (2022), we benchmark policy optimization with
transformer AWMs against other competitive model-free and model-based methods,
showing superior performance.

###  2 Background

Problem definition We are interested in deterministic111 See Laidlaw et al.
(2023) for a discussion on the importance of the study of deterministic MDPs.
discrete-time finite-horizon Markov Decision Processes (MDPs) (Fairbank,
2014), defined as $\mathcal{M}=(\mathcal{S},\mathcal{A},f,r,H,s_{1})$, where
$\mathcal{S}\subseteq\mathbb{R}^{n}$ is the state space,
$\mathcal{A}\subseteq\mathbb{R}^{m}$ is the action space,
$f:\mathcal{S}\times\mathcal{A}\to\mathcal{S}$ is the differentiable
transition dynamics, $r:\mathcal{S}\to\mathbb{R}$ is the known differentiable
reward function, $H$ is the horizon and $s_{1}\in\mathcal{S}$ is the initial
state. The behavior of an agent in the environment is described by a policy
$\pi_{{\bm{\theta}}}:\mathcal{S}\to\mathcal{A}$, belonging to a space of
parameterized differentiable deterministic222This assumption can be relaxed to
a stochastic policy using the reparameterization trick (Kingma & Welling,
2014). stationary Markov policies
$\Pi=\\{\pi_{{\bm{\theta}}}:{\bm{\theta}}\in\mathbb{R}^{d}\\}$. For clarity,
we focus on stationary rewards and dynamics, but our results could readily be
adapted to the non-stationary case.

Backpropagation-based Policy Optimization The goal of the agent is to maximize
the cumulative episodic rewards. In the rest of the paper, we consider the
following problem:

$\begin{gathered}\begin{aligned}
\text{maximize}\quad&J^{f}({\bm{\theta}};H):=\sum_{t=1}^{H}r(s_{t}),&\text{subject
to}\quad&s_{t+1}=f(s_{t},a_{t}),\,a_{k}=\pi_{\bm{\theta}}(\text{sg}[s_{k}])\enspace.\end{aligned}\end{gathered}$
(1)

where $\mathrm{sg}[\cdot]$ is the stop-gradient operator. Model-based policy
gradient methods solve this problem by unrolling (simulating) the dynamics
forward and computing the gradient of the corresponding unconstrained problem.
The presence of the stop gradient operator in our formulation aligns with the
contemporary practice (Hafner et al., 2022; 2023; Ghugare et al., 2023). Its
inclusion in our formulation does not impact performance (see Appendix F.1 for
an ablation) and simplifies our theoretical analysis.

Input: Initial buffer $\mathcal{B}$, initial policy parameters
${\bm{\theta}}$, initial model parameters ${\bm{\psi}}$, learning rates
$\\{\alpha_{{\bm{\theta}}},\alpha_{{\bm{\psi}}}\\}$, world model class
$\mathcal{W}$.

1: while not exceeding training steps do

2: Collect an episode with $\pi_{{\bm{\theta}}}$ and add it to $\mathcal{B}$

3: for each world model learning step do

4:
${\bm{\psi}}\leftarrow{\bm{\psi}}-\alpha_{{\bm{\psi}}}\nabla_{{\bm{\psi}}}\ell^{\mathcal{W}}(\tau;{\bm{\psi}}),\quad\tau\sim\mathcal{B}$

5: end for

6: for each policy learning step do

7: Compute $J^{\mathcal{W}}({\bm{\theta}};{\bm{\psi}})$ by unrolling the world
model

8: Compute $\nabla_{\theta}J^{\mathcal{W}}({\bm{\theta}};{\bm{\psi}})$ by
backpropagation

9:
${\bm{\theta}}\leftarrow{\bm{\theta}}+\alpha_{{\bm{\theta}}}\nabla_{{\bm{\theta}}}J^{\mathcal{W}}({\bm{\theta}};{\bm{\psi}})$

10: end for

11: end while

Algorithm 1 Backpropagation-based Policy Optimization (BPO)

When the true dynamics $f$ are unknown, a transition function
$\hat{f}_{\bm{\psi}}(s_{t},a_{t})$ belonging to a space of parameterized
differentiable functions
$\mathcal{F}=\\{f_{\bm{\psi}}:{\bm{\psi}}\in\mathbb{R}^{d_{\psi}}\\}$ must be
learned to perform policy optimization. Given transitions
$s_{t},a_{t},s_{t+1}$ sampled from the environment, an approximate Markovian
world model $\hat{f}_{\bm{\psi}}$ can be learned with the loss
${\ell^{\hat{f}}(\tau;{\bm{\psi}})=\sum_{t=1}^{H-1}\|s_{t+1}^{\tau}-\hat{f}_{\bm{\psi}}(s_{t}^{\tau},a_{t}^{\tau})\|^{2}}$,
where ${\tau=(s_{1}^{\tau},a_{1}^{\tau},\dots,s_{H}^{\tau})}$ is a trajectory.
For the remainder of this paper, consider $\|\bullet\|$ to denote the $L^{2}$
norm. Given a learned model, policy and initial state, the corresponding
return $J^{\hat{f}}({\bm{\theta}};{\bm{\psi}})$ can be computed through
rollouts and the gradient
$\nabla_{{\bm{\theta}}}J^{\hat{f}}({\bm{\theta}};{\bm{\psi}})$ can be obtained
by backpropagation. We call this basic algorithm, outlined in Algorithm 1,
_Backpropagation-based Policy Optimization_ (BPO).

###  3 Policy Gradient Computation with History World Models

When running Algorithm 1, the performance of the resulting policy depends on
the quality of the gradient approximation. Effective long-term credit
assignment capabilities are required for computing the gradient
$\nabla_{a_{t}}r(s_{t+k})$, as provided by the world model, for potentially
very large values of $k$ in challenging control tasks. This means
understanding how a reward at a temporally distant time changes in response to
variations in a specific action taken at an earlier moment. Therefore, the
long-term credit assignment properties of a model-based policy optimization
algorithm are intimately tied to the structure of the unrolled world model.

One fundamental fact about backpropagation is that its success in learning
long-term dependencies is intimately linked to the length of the paths
involved in the neural network’s gradient computation (Pascanu et al., 2013a;
Zhang et al., 2016). Indeed, the success of transformers is often attributed
to this phenomenon: “one key factor affecting the ability to learn such
dependencies is the length of the paths forward and backward signals have to
traverse in the network” (Vaswani et al., 2017). When trained for sequence
modeling tasks, transformers stand out compared to other architectures because
the length of these paths does not linearly increase with the sequence length.

Recent model-based RL approaches have introduced the use of transformer
History World Models (HWMs), which predict the next state as
$\hat{s}_{t}=h(s_{1:{t-1}},a_{1:{t-1}})$ based on the full history of states
and actions (Micheli et al., 2022; Robine et al., 2023). Their use leads to
improved performance due to better prediction abilities but, surprisingly,
these approaches were not able to take advantage of the gradient propagation
properties of the transformer architecture, as exemplified by previous work
which only unrolls them for a few steps (Chen et al., 2022).

In practice, an HWM is unrolled in this manner to compute the cumulative
reward for a given policy:

$\begin{gathered}\begin{aligned}
\text{maximize}\quad&J^{h}({\bm{\theta}};H):=\sum_{t=1}^{H}r(\hat{s}_{t}),&\text{subject
to}\quad&\hat{s}_{t+1}=h(\hat{s}_{1:t},a_{1:k}),\,a_{k}=\pi_{\bm{\theta}}(\text{sg}[\hat{s}_{k}])\enspace.\end{aligned}\end{gathered}$
(2)

HWMs are typically learned by employing a prediction objective
$\ell^{h}(\tau;{\bm{\psi}})=\sum_{t=1}^{H-1}\|s_{t+1}^{\tau}-h_{\bm{\psi}}(s_{1:t}^{\tau},a_{1:t}^{\tau})\|^{2}$,
similar to the one employed for a Markovian model, where $s_{1:t}^{\tau}$ and
$a_{1:t}^{\tau}$ denote subsequences of states and actions from an episode
$\tau$. While history-conditioned dynamics are in principle not necessary in
an MDP333Policy gradients can be computed through equation 2 similarly in
partially observable MDPs (POMDP), where states are replaced with
observations. The resulting analysis would, therefore, still hold in a POMDP.,
one may hope that the gradient properties of transformers might manifest
themselves positively in the policy gradient. It stands to reason that the
policy gradients through a transformer may be able to more effectively capture
long-term dependencies than a Markovian model would. Nonetheless, this
apparently natural usage of a transformer has unintended consequences for
policy optimization.

Visually, we show in Figure 1 the gradient paths induced by unrolling a
_state-based_ world model, whether Markovian (Figure 1a) or history-dependent
(Figure 1b), from state predictions to actions. The impact of an action on a
future state is measured by navigating the paths connecting the two generated
by a world model. Some of these are _circuitous gradient paths_ : these paths
do not go from state to action directly but instead include auto-regressive
state predictions from the model itself. Such paths are perilous for long-
horizon policy optimization; errors, noise, and gradient behaviors that may be
good internally can rapidly degrade outside the model, reducing the possible
advantages of backpropagating through transformers.

In particular, the following theorem shows that even if the gradient of a
transformer HWM is bounded, the resulting policy gradient can still grow
exponentially. The proof of this and the rest of our results can be found in
Appendix A.

######  Theorem 1.

Let the gradient norm of $h$ with respect to its inputs be bounded by $L_{a}$
and $L_{s}$: $\|\frac{\partial h(\hat{s}_{1:t},a_{1:t})}{\partial a_{k}}\|\leq
L_{a}$ and $\|\frac{\partial
h(\hat{s}_{1:t},a_{1:t})}{\partial\hat{s}_{i}}\|\leq L_{s}$ for all
$s_{1:t},a_{1:t},k,i$. Let $r$ be the $L_{r}$-Lipschitz reward function from a
Markov Decision Process $\mathcal{M}$, $\Pi_{\bm{\theta}}$ a parametric space
of differentiable deterministic $L_{\pi}$-policies. Given
$\pi_{\bm{\theta}}\in\Pi_{\bm{\theta}}$, the norm of the policy gradient
$\nabla_{\bm{\theta}}J^{h}({\bm{\theta}};H)$ of $\pi_{\bm{\theta}}$ under a
History World Model $h$ grows asymptotically as a function of the horizon $H$
as:

$\displaystyle\|\nabla_{\theta}J^{h}({\bm{\theta}};H)\|$
$\displaystyle=O(HL_{r}+H^{2}L_{\pi}+H^{2}L_{a}+H^{2}L_{s}^{H})=O(L_{s}^{H})\enspace.$

Figure 2: Transformer AWMs outperforms all BPO baselines in chaotic
environments. Final performances of BPO with different world models on the
double-pendulum environment (10 seeds $\pm$ std).

The exponential dependency, related to exploding gradients in RNNs, is
apparent in Figure 1b, where the longest path from any given state to action
is no shorter than the longest path through a Markovian model. The bound is
shown to be tight in Appendix A.1. In the one-dimensional case, when all
gradients are positive scalars, the above bound can be turned into a lower
bound.

To illustrate the negative effects of circuitous gradient paths in practice,
we conduct experiments in a double-pendulum environment where the agent must
move the double-pendulum to a target position. We describe and visualize the
task more precisely in Section 5. The results in Figure 2 show that BPO does
not lead to successful policies with both Markovian models and HWMs in this
notoriously chaotic environment. This result demonstrates that state-
conditioned transformers do not provide any benefits to policy gradients due
to circuitous gradient paths, unlike in the traditional supervised learning
setting where they were initially introduced. In the rest of the paper, we
will see how using transformers as a different type of world model allows us
not only to overcome the limitations of Markovian models and HWMs, but also to
outperform policy optimization performed by differentiating through the
dynamics of the simulator, as well as competitive model-free approaches.

###  4 Policy Gradient Computation with Actions World Models

Policy gradients obtained from state-based world models, whether Markovian
models or HWMs, are necessarily computed by backpropagating through sequences
of states generated by autoregressively unrolling the models. Each state
prediction uses as input previously generated states; consequently, the
longest path between any action and reward traverses through circuitous
gradient paths in the state space, _regardless of the underlying world model
architecture_. Through careful examination of Figure 1 and Theorem 1, it
quickly becomes apparent that the state dependency in $h$ or $f$ is at the
root cause of such paths.

Inspired by this observation, we propose to remove all circuitous gradient
paths in the simplest and most direct way: conditioning the world model on
actions only. We thus consider a model which predicts a future state given an
initial state and a sequence of actions. The resulting model, which we call an
_Actions World Model_ (AWM), is of the form
$\hat{s}_{t+1}=g(s_{1},a_{1},a_{2},\ldots,a_{t})$. This structure for the
world model constrains gradients to pass directly from rewards to actions
through the network parameterizing the AWM. Note that previous work showed
that actions are theoretically sufficient to predict future states in
deterministic MDPs and often sufficient in practice for most stochastic
environments (Rezende et al., 2020).

Given an AWM, a policy is learned by computing and employing (as in Algorithm
1) the gradient of the objective:

$\begin{gathered}\begin{aligned}
\text{maximize}\quad&J^{g}({\bm{\theta}};H):=\sum_{t=1}^{H}r(\hat{s}_{t}),&\text{subject
to}\quad&\hat{s}_{t+1}=g(s_{1},a_{1:t}),\,a_{k}=\pi_{\bm{\theta}}(\mathrm{sg}[\hat{s}_{k}])\enspace.\end{aligned}\end{gathered}$
(3)

An AWM can be trained in a similar manner to a Markovian model or an HWM, by
using the loss
${\ell^{g}(\tau;{\bm{\psi}})=\sum_{t=1}^{H-1}\|s_{t+1}^{\tau}-g_{\bm{\psi}}(s_{1},a_{1:t}^{\tau})\|^{2}}$,
computed likewise using states from an episode $\tau$.

We will now show that, despite the seemingly counter-intuitive removal of
state inputs, AWMs naturally connect the theory of policy gradients and
sequence modeling with neural networks and provide a framework in which
transformers can give high-quality policy gradients.

####  4.1 Markovian Models are RNN Actions World Models

The intuition that policy gradients computed through an unrolled Markovian
model prescribed by equation 1 resembles the process of backpropagating
through a recurrent neural network has been pointed out in multiple prior
works Heess et al. (2015); Metz et al. (2021); Suh et al. (2022); Zhang et al.
(2023). Yet, a formal account of this relationship has never been shown. Via
the framework provided by AWMs, the next proposition shows how recurrent
neural networks and policy gradients are connected, offering a first glimpse
at the connection between neural network architectures, AWMs and policy
gradient computation in MDPs.

######  Proposition 1.

Let $f\text{-}\texttt{RNN}$ be a recurrent network with its recurrent cell
being the dynamics $f$ of the MDP $\mathcal{M}$, and
$g_{f\text{-}\texttt{RNN}}$ denote an AWM instantiated with
$f\text{-}\texttt{RNN}$. Then,

$\displaystyle\nabla_{\bm{\theta}}J^{g_{f\text{-}\texttt{RNN}}}({\bm{\theta}};H)=\nabla_{\bm{\theta}}J^{f}({\bm{\theta}};H).$

The above proposition tells us that the policy gradient computed through a
Markovian model is, in fact, equivalent to the one computed through an AWM
when instantiating $g$ as a recurrent neural network with a specific recurrent
cell. Crucially, this not only provides grounding for gradient estimation with
AWMs but also solidifies a fundamental fact that will be analyzed in-depth in
this section: policy gradient computation by differentiating through unrolled
Markovian models can be understood to be fundamentally ill-behaved due to its
correspondence to an RNN structure.

####  4.2 Theoretical Properties of Actions World Models

Through the concept of Actions World Models, we have established a direct
connection between deep sequence models and policy gradient computation. We
will now exploit this connection to characterize the asymptotic behavior of
the policy gradient depending on the underlying neural network architecture
employed as an AWM. To do so, we leverage the following theorem.

######  Theorem 2.

Let $r$ be the $L_{r}$-Lipschitz reward function from a Markov Decision
Process $\mathcal{M}$, $\Pi_{\bm{\theta}}$ a parametric space of
differentiable deterministic $L_{\pi}$-policies. Given
$\pi_{\bm{\theta}}\in\Pi_{\bm{\theta}}$, the norm of the policy gradient
$\nabla_{\bm{\theta}}J^{g}({\bm{\theta}};H)$ of $\pi_{\bm{\theta}}$ under an
action world model $g$ as a function of the horizon $H$ can upper bounded as:

$\displaystyle\left\|\nabla_{\bm{\theta}}J^{g}({\bm{\theta}};H)\right\|\leq
L_{r}L_{\pi}\sum_{t=1}^{H}\sum_{k=1}^{t-1}\left\|\frac{\partial
g(s_{1},a_{1:t-1})}{\partial a_{k}}\right\|.$

This theorem holds for any differentiable AWM $g$, and establishes a worst-
case relationship between the Jacobian of the AWM w.r.t. its inputs (i.e.,
actions) and the policy gradient. Notably, the result implies that the policy
gradient does not explode if the Jacobian of the AWM does not explode. In
contrast, this may not be true even for well-behaved HWM.

We will now leverage the generality of the previous result to characterize the
asymptotic behavior of the policy gradient computed using different neural
network architectures for an AWM. First, let us consider an AWM instantiated
with a simple recurrent neural network:

$\displaystyle\text{(}g_{\texttt{RNN}}\text{)}$ $\displaystyle
x_{t+1}=\sigma(W_{x}x_{t})+W_{a}a_{t}+b;$
$\displaystyle\hat{s}_{t+1}=W_{o}x_{t+1},$ (4)

where $\sigma$ is an activation function with gradient norm bounded by
$\left\|diag(\sigma^{\prime}(x))\right\|\leq\frac{1}{\beta}$ for some constant
$\beta$. Then, the following result holds.

######  Corollary 2.1.

Let $g_{\texttt{RNN}}$ be an Actions World Model instantiated with a recurrent
neural network as in Equation 4 and $\eta=\|W_{x}^{T}\|\frac{1}{\beta}$. The
asymptotic behavior of the norm of the policy gradient
$\nabla_{\bm{\theta}}J^{g_{\texttt{RNN}}}({\bm{\theta}};H)$ as a function of
the horizon H can be described as:

$\displaystyle\left\|\nabla_{{\bm{\theta}}}J^{g_{\texttt{RNN}}}({\bm{\theta}};H)\right\|=O\left(\eta^{H}\right).$

Corollary 2.1 shows that, in the worst case, the policy gradient computed with
an AWM instantiated with an RNN backbone can explode exponentially fast with
respect to the problem horizon if the spectral radius of $W_{x}$ exceeds
$\beta$. Just like Theorem 1, the above bound can be written as a lower bound
when all gradients are positive scalars, therefore it is tight. This result
will allow us to compare the gradient provided by an RNN AWM with the one
provided by a transformer AWM, but it is also connected with gradient
computation with Markovian models: as seen in Proposition 1, using a Markovian
model to compute the policy gradient, either given or learned, implies an RNN-
structure for the AWM. Theorem 2.1 is a consequence of foundational theory of
RNNs (Bengio et al., 1994; Pascanu et al., 2013b), thus formalizing the
intuition that the recurrent nature of policy gradients through Markovian
models make both exploiting differentiable simulators $f$ and learned models
$\hat{f}_{\bm{\psi}}$ difficult (Metz et al., 2021; Heess et al., 2015).

Let us now apply Theorem 2 to a layer of self-attention AWM (Vaswani et al.,
2017). Given $Q\in\mathbb{R}^{1\times d_{z}}$, $K\in\mathbb{R}^{n\times
d_{z}}$, $V\in\mathbb{R}^{n\times d_{o}}$, we define
$\text{Attention}(Q,K,V):=\text{softmax}(QK^{T})V$. Using a set of weight
matrices, a self-attention AWM predicts the next state as follows444Without
loss of generality, the dependency on the initial state is not included in the
self-attention layer.:

$\displaystyle\text{(}g_{\texttt{ATT}}\text{)}$
$\displaystyle\hat{s}_{t+1}=\text{Attention}(a_{t}W_{q}^{T},a_{1:t}W_{k}^{T},a_{1:t}W_{v}^{T})\enspace.$
(5)

We can now characterize the behavior of the policy gradient computed with a
self-attention AWM.

######  Corollary 2.2.

Let $g_{\texttt{ATT}}$ be an attention-based Actions World Model instantiated
with self-attention as in equation 5. The asymptotic behavior of the norm of
the policy gradient
$\nabla_{{\bm{\theta}}}J^{g_{\texttt{ATT}}}({\bm{\theta}};H)$ as a function of
the horizon $H$ can be described as:

$\displaystyle\left\|\nabla_{{\bm{\theta}}}J^{g_{\texttt{ATT}}}({\bm{\theta}};H)\right\|=O\left(H^{3}\right).$

This result shows that the norm of the policy gradient computed through a
self-attention AWM has a worst-case polynomial dependency on the horizon
instead of an exponential one shown in Theorem 1 and Corollary 2.1. Since the
output of a transformer model only depends on the sequence length through the
attention mechanism, this characterization captures the gradient dynamics of
the transformer model itself. It shows that, by eliminating all circuitous
gradient paths, AWMs provide a framework for transformers to translate their
favorable gradient properties towards better policy gradients, unlike state-
conditioned models.

(a) One-bounce environment overview.

(b) Example trajectory with non-differentiable point at $t\approx 0.25$.

(c) Final return with respect to the initial action for different models.

Figure 3: AWMs ignore non-differentiable points in the state space. (a) After
the block is pushed with some initial action, it bounces off the wall,
instantaneously reversing its velocity. (b) Visualization of the point of non-
differentiability in the state space. (c) Learning a Markovian model or a HWM
causes catastrophic compounding errors, but an AWM can still accurately model
the final reward when varying the initial action. Learned dynamics are trained
offline on a dataset collected using random actions.

###  5 Experiments

Backpropagation through Actions World Models produces policy gradients with no
circuitous gradient paths. Consequently, the longest gradient path from an
action to any reward strictly depends on the AWM internals, while the path
through state-based world models scales with the horizon, regardless of the
model architecture. This section empirically explores how this phenomenon
benefits the policy optimization process (see Appendix C for hyperparameters
of all experiments).

####  5.1 Empirical Properties of Action World Models

By removing the need to explicitly model the state dynamics auto-regressively,
AWMs can internally build a more appropriate world model for policy
optimization. Using environments inspired by real-world physical systems, we
highlight two particular properties of AWMs: skipping points of complex state-
space dynamics and smoothing out ill-behaved policy optimization landscapes.

AWMs overcome non-differentiable dynamics Unlike state-based world models,
AWMs directly map actions to future states without explicitly predicting
intermediate states. This mapping allows AWMs to skip points of potentially
complex dynamics in the state space, which are not always necessary for policy
optimization. For example, consider points of non-differentiability (e.g.,
contact points) in real-world systems (e.g., physical systems). These points
make directly using the simulator dynamics for BPO impossible and can cause
catastrophic compounding errors for state-conditioned world models. However,
these points often exist even if the dependency of rewards on actions is
smooth and well-behaved everywhere. In such cases, AWMs can be used to side-
step the non-differentiability of the state dynamics. To illustrate this
point, we examine a _one-bounce_ environment, shown in Figure 3. In this task,
an agent must push a block towards a wall with some initial action, such that
the block bounces off the wall and ends up at some predetermined goal state
after $H$ steps. A single terminal reward is given at $t=H$ measuring the
distance of the block to the goal state. Even though the state trajectory is
non-differentiable due to the wall bounce, the final state of the block is
actually well-behaved with respect to initial actions (Figure 3b and Figure
3c). Consequently, we show in Figure 3c that AWMs can accurately predict the
final reward while an unrolled Markovian or HWM cannot. Notably, the resulting
AWM thus allows to easily use backpropagation-based policy optimization even
in the presence of an underlying non-differentiable environment dynamics.

AWMs overcome chaotic dynamics In many real-world systems (e.g., robots), the
chaotic dynamics of an environment can create a remarkably complex mapping
from actions to rewards over time or a highly non-smooth return landscape
(Rahn et al., 2023). In such systems, even when a differentiable simulator is
available, it gives ill-behaved gradients and generates an optimization
landscape that is hard to navigate (Suh et al., 2022). We now show that
transformer AWMs are able to take full advantage of the inductive biases of
the underlying neural networks, naturally generating policy optimization
landscapes that are easy to navigate. This happens regardless of the
complexity of the real dynamics while preserving the global problem structure
and allowing good policies to be found more easily. To demonstrate this, we
use a prototypical chaotic system, a double-pendulum task depicted in Figure
4a. In this environment, the agent’s initial action determines the initial
angular position of the inner pendulum. Afterwards, the system is rolled out
for $H$ steps. The objective is to get the pendulums to be in some predefined
goal state after $H$ steps. A terminal reward is given measuring the distance
between the final observed state and the desired goal state. We illustrate in
Figure 4b and Figure 4c the chaotic nature of the task, showing respectively
that the norm of the true gradient of the return with respect to the initial
action grows exponentially with the episode length $H$, and that the true
return landscape is extremely difficult to navigate for long horizons. Figure
4b shows that transformers provide stable gradients when used as AWMs, while
the gradient provided by HWMs explodes similarly to the true gradients. AWMs
generate a smooth and accurate approximation of the return landscape (Figure
4c), which gradient ascent can easily navigate, leading to better policies
when used for backpropagation-based policy optimization compared to all the
other world models (Figure 2).

(a) Double-pendulum.

(b) Gradient norms.

(c) Final returns with respect to the initial action for different models when
$H=100$.

Figure 4: Transformer AWMs smooths out chaotic dynamics. (a) A double-pendulum
environment where an initial position must be chosen in order to achieve some
pre-determined goal state after $H$ steps. Different transition models are
learned on a data set of random trajectories. (b) The mean gradient norm of
the final state with respect to the initial action for each model is computed
over 50 different random actions for different horizons. (c) Final return
according to different models with respect to different initial actions for
$H=100$.

####  5.2 Benchmarking Credit Assignment with AWMs

We now show that transformers employed as AWMs can give better policies in
environments which require long-horizon planning. We focus on eight tasks from
the Myriad testbed Howe et al. (2022). Myriad provides continuous MDPs
inspired by real-world problems Lenhart & Workman (2007) with configurable
horizons. In Myriad, the agent requires better temporal credit assignment
capabilities to solve problems with longer episode lengths, in contrast to
other RL benchmarks (e.g., focused on simple robotics locomotion), which
typically feature a short effective horizon Laidlaw et al. (2023); Ni et al.
(2023) regardless of the episode length. Additional information about Myriad
and its experiments are provided in Appendix F. All aggregate results use the
IQM Agarwal et al. (2021).

(a) Final performances of BPO for different horizons.

(b) Learning curves compared to other baselines.

Figure 5: Policy optimization with transformer AWMs gives better policies for
long horizons. (a) Final performance of BPO with different world models on
Myriad (10 seeds $\pm$ 95% C.I.). (b) Learning curves of BPO through a
transformer AWM, a SAC agent, and an Online-DT agent on 20 and 100 length
horizons (10 seeds $\pm$ 95% C.I.).

Figure 5a shows the final performance of various world models used in BPO with
increasing episode lengths. The transformer AWM is the only model to perform
reliably well for increasingly long horizons, while other methods struggle
when the horizon exceeds 100 steps. Indeed, when used as AWMs, transformer
world models give better policy gradients than standard Markovian models and
RNN world models. However, when transformer world models are conditioned on
states, their long-term policy gradients are no better than the gradients
induced by Markovian models as measured by performance in Figure 5a. These
results corroborate our previous findings: transformer HWMs, RNN AWMs, and
Markovian models all generate long policy gradient paths, which makes long-
term credit assignment difficult.

We also compare our method with other relevant baselines in Figure 5b, namely
model-free Soft Actor-Critic (SAC) Haarnoja et al. (2018) and the recently
popularized Online Decision Transformer (Online-DT) Zheng et al. (2022);
Schmidhuber (2019); Srivastava et al. (2019). In addition to BPO with
transformer AWMs scaling better with the problem horizon, BPO is also
significantly more sample efficient than the two baselines.

###  6 Related Work

Backpropagation-based Policy Optimization The concept of learning a policy by
differentiating through the environment dynamics has been a foundational
aspect of early RL methods (Werbos, 1974; Schmidhuber, 1990). Estimators for
such a policy gradient have been referred to with different names, including
pathwise derivatives (Clavera et al., 2020; Hafner et al., 2022; 2023) and
value gradients (Fairbank, 2014; Heess et al., 2015; Amos et al., 2021). AWMs
connect this literature to the one on sequence modeling with neural networks,
with a focus on gradient propagation (Hochreiter, 1991; Bengio et al., 1994;
Hochreiter, 1998; Pascanu et al., 2013b) and on architectures that allow for
improved credit assignment (Hochreiter & Schmidhuber, 1997; Vaswani et al.,
2017; Kerg et al., 2020). Finally, our work can be seen as an instance of
gradient-aware model-based RL (D’Oro et al., 2020; D’Oro & Jaskowski, 2020;
Abachi et al., 2020), in which we modify the architecture of a model of the
dynamics to obtain better policy gradients (Ma et al., 2021).

Multi-step state prediction with actions Prior works used multi-step models
related to AWMs, analyzing their role as partial models (Rezende et al.,
2020), in the context of tree search (Schrittwieser et al., 2019), or training
them with a multi-step latent prediction loss (Gregor et al., 2019; Schwarzer
et al., 2020), without however ever analyzing their gradient properties.
Notably, models used in these prior works are rarely shown to work beyond a
few dozen steps of unrolling or imagination, while AWMs demonstrably scale
favorably for horizons well beyond established limitations.

Sequence models in RL In partially observable MDPs, sequence models have been
extensively used as history encoders to maximize returns with or without world
models Hausknecht & Stone (2015); Ni et al. (2021); Hafner et al. (2023).
Other works have recently emerged treating MDPs as sequence modeling Chen et
al. (2021); Zheng et al. (2022); Janner et al. (2021), in imitation learning
or offline RL problem settings, based on return-conditioned models. Sequence
models have also been used to reshape rewards in sparse reward settings Hung
et al. (2019); Arjona-Medina et al. (2019); Liu et al. (2019). In contrast to
all of these, our framework uses an action-only conditioned model to directly
improve long-term policy gradients and naturally assigns credit by
backpropagation as opposed to reward-shaping heuristics. Appendix G contains a
more detailed discussion of how all these methods differ from ours.

###  7 Conclusions and Limitations

We presented a model-based policy optimization framework based on
backpropagation through Actions World Models: models of the dynamics of the
environment which only take an initial state and a sequence of actions as
inputs. We highlighted that, in contrast to the more common world models that
compute the gradient by autoregressively unrolling their state predictions,
AWMs do not suffer from the existence of circuitous gradient paths where the
gradient can flow outside of the underlying network architecture. We showed
that their structure allows AWMs to fully exploit transformers’ favorable
properties, taking advantage, both in theory and in practice, of their ability
to perform long-term temporal credit assignment by leveraging short gradient
paths.

Theoretically, we demonstrated how transformer AWMs do not suffer from ill-
behaved gradients coming from circuitous gradient paths which harm the quality
of the gradients provided by more common Markovian models and HWMs.
Empirically, we showed that BPO using AWMs can be successful even when BPO
using the underlying simulator fails, due to how AWMs leverage transformers’
inductive biases to obtain an easier-to-navigate policy optimization
landscape. We demonstrated that policies trained with transformer AWMs
outperform existing model-based and model-free approaches, including recently
proposed algorithms based on transformers, in realistic domains requiring
long-term planning.

While our theoretical analysis is developed within a deterministic setting,
and our experiments are conducted in low-dimensional yet realistic domains,
our framework establishes a new connection between gradient propagation
techniques for sequence modeling and backpropagation-based policy gradients.
Future work could focus on scaling our approach, making RL even more relevant
to real-world applications where long-horizon capabilities are crucial.

##### Acknowledgements and Disclosure of Funding

The authors thank Nathan Rahn and Jürgen Schmidhuber for insightful
discussions and useful suggestions on the paper, the Mila community for
creating a stimulating research environment, Digital Research Alliance of
Canada and Nvidia for computational resources. This work was partially
supported by CIFAR, Google, Facebook AI Chair, IVADO, and Gruppo Ermenegildo
Zegna.

### References

  * Abachi et al. (2020) Romina Abachi, Mohammad Ghavamzadeh, and Amir massoud Farahmand.  Policy-aware model learning for policy gradient methods.  _ArXiv_ , abs/2003.00030, 2020.
  * Agarwal et al. (2021) Rishabh Agarwal, Max Schwarzer, Pablo Samuel Castro, Aaron C. Courville, and Marc G. Bellemare.  Deep reinforcement learning at the edge of the statistical precipice.  In Marc’Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan (eds.), _Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual_ , pp. 29304–29320, 2021.  URL https://proceedings.neurips.cc/paper/2021/hash/f514cec81cb148559cf475e7426eed5e-Abstract.html.
  * Amos et al. (2021) Brandon Amos, Samuel Stanton, Denis Yarats, and Andrew Gordon Wilson.  On the model-based stochastic value gradient for continuous reinforcement learning learning.  In Ali Jadbabaie, John Lygeros, George J. Pappas, Pablo A.;Parrilo, Benjamin Recht, Claire J. Tomlin, and Melanie N. Zeilinger (eds.), _Proceedings of the 3rd Conference on Learning for Dynamics and Control_ , volume 144 of _Proceedings of Machine Learning Research_ , pp. 6–20. PMLR, 07 – 08 June 2021.  URL https://proceedings.mlr.press/v144/amos21a.html.
  * Arjona-Medina et al. (2019) Jose A Arjona-Medina, Michael Gillhofer, Michael Widrich, Thomas Unterthiner, Johannes Brandstetter, and Sepp Hochreiter.  Rudder: Return decomposition for delayed rewards.  _Advances in Neural Information Processing Systems_ , 32, 2019.
  * Bengio et al. (1994) Yoshua Bengio, Patrice Y. Simard, and Paolo Frasconi.  Learning long-term dependencies with gradient descent is difficult.  _IEEE Trans. Neural Networks_ , 5(2):157–166, 1994.  doi: 10.1109/72.279181.  URL https://doi.org/10.1109/72.279181.
  * Chen et al. (2022) Chang Chen, Yi-Fu Wu, Jaesik Yoon, and Sungjin Ahn.  Transdreamer: Reinforcement learning with transformer world models.  _arXiv preprint arXiv:2202.09481_ , 2022.
  * Chen et al. (2021) Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Michael Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch.  Decision transformer: Reinforcement learning via sequence modeling.  In Marc’Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan (eds.), _Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual_ , pp. 15084–15097, 2021.  URL https://proceedings.neurips.cc/paper/2021/hash/7f489f642a0ddb10272b5c31057f0663-Abstract.html.
  * Clavera et al. (2020) Ignasi Clavera, Yao Fu, and Pieter Abbeel.  Model-augmented actor-critic: Backpropagating through paths.  In _8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020_. OpenReview.net, 2020.  URL https://openreview.net/forum?id=Skln2A4YDB.
  * D’Oro & Jaskowski (2020) Pierluca D’Oro and Wojciech Jaskowski.  How to learn a useful critic? model-based action-gradient-estimator policy optimization.  In Hugo Larochelle, Marc’Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien Lin (eds.), _Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual_ , 2020.  URL https://proceedings.neurips.cc/paper/2020/hash/03255088ed63354a54e0e5ed957e9008-Abstract.html.
  * D’Oro et al. (2020) Pierluca D’Oro, Alberto Maria Metelli, Andrea Tirinzoni, Matteo Papini, and Marcello Restelli.  Gradient-aware model-based policy search.  In _The Thirty-Fourth AAAI Conference on Artificial Intelligence, AAAI 2020, The Thirty-Second Innovative Applications of Artificial Intelligence Conference, IAAI 2020, The Tenth AAAI Symposium on Educational Advances in Artificial Intelligence, EAAI 2020, New York, NY, USA, February 7-12, 2020_ , pp. 3801–3808. AAAI Press, 2020.  doi: 10.1609/aaai.v34i04.5791.  URL https://doi.org/10.1609/aaai.v34i04.5791.
  * Fairbank (2014) Michael Fairbank.  _Value-gradient learning_.  PhD thesis, City University London, 2014.
  * Ghugare et al. (2023) Raj Ghugare, Homanga Bharadhwaj, Benjamin Eysenbach, Sergey Levine, and Ruslan Salakhutdinov.  Simplifying model-based rl: Learning representations, latent-space models, and policies with one objective, 2023.
  * Gregor et al. (2019) Karol Gregor, Danilo Jimenez Rezende, Frederic Besse, Yan Wu, Hamza Merzic, and Aäron van den Oord.  Shaping belief states with generative environment models for RL.  _CoRR_ , abs/1906.09237, 2019.  URL http://arxiv.org/abs/1906.09237.
  * Gu et al. (2022) Albert Gu, Karan Goel, and Christopher Ré.  Efficiently modeling long sequences with structured state spaces, 2022.
  * 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.  _CoRR_ , abs/1801.01290, 2018.  URL http://arxiv.org/abs/1801.01290.
  * Hafner et al. (2022) Danijar Hafner, Timothy Lillicrap, Mohammad Norouzi, and Jimmy Ba.  Mastering atari with discrete world models, 2022.
  * Hafner et al. (2023) Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap.  Mastering diverse domains through world models, 2023.
  * Hausknecht & Stone (2015) Matthew Hausknecht and Peter Stone.  Deep recurrent q-learning for partially observable mdps.  In _2015 aaai fall symposium series_ , 2015.
  * Heess et al. (2015) Nicolas Heess, Gregory Wayne, David Silver, Timothy P. Lillicrap, Tom Erez, and Yuval Tassa.  Learning continuous control policies by stochastic value gradients.  In Corinna Cortes, Neil D. Lawrence, Daniel D. Lee, Masashi Sugiyama, and Roman Garnett (eds.), _Advances in Neural Information Processing Systems 28: Annual Conference on Neural Information Processing Systems 2015, December 7-12, 2015, Montreal, Quebec, Canada_ , pp. 2944–2952, 2015.  URL https://proceedings.neurips.cc/paper/2015/hash/148510031349642de5ca0c544f31b2ef-Abstract.html.
  * Hochreiter (1991) Sepp Hochreiter.  Untersuchungen zu dynamischen neuronalen netzen.  1991\.  URL https://api.semanticscholar.org/CorpusID:60091947.
  * Hochreiter (1998) Sepp Hochreiter.  The vanishing gradient problem during learning recurrent neural nets and problem solutions.  _International Journal of Uncertainty, Fuzziness and Knowledge-Based Systems_ , 6(02):107–116, 1998.
  * Hochreiter & Schmidhuber (1997) Sepp Hochreiter and Jürgen Schmidhuber.  Long short-term memory.  _Neural Computation_ , 9(8):1735–1780, November 1997.  doi: 10.1162/neco.1997.9.8.1735.  URL https://doi.org/10.1162/neco.1997.9.8.1735.
  * Howe et al. (2022) Nikolaus Howe, Simon Dufort-Labbé, Nitarshan Rajkumar, and Pierre-Luc Bacon.  Myriad: a real-world testbed to bridge trajectory optimization and deep learning.  In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), _Advances in Neural Information Processing Systems_ , volume 35, pp. 29801–29815. Curran Associates, Inc., 2022.  URL https://proceedings.neurips.cc/paper_files/paper/2022/file/c0b91f9a3587bf35287f41dba5d20233-Paper-Datasets_and_Benchmarks.pdf.
  * Hung et al. (2019) Chia-Chun Hung, Timothy Lillicrap, Josh Abramson, Yan Wu, Mehdi Mirza, Federico Carnevale, Arun Ahuja, and Greg Wayne.  Optimizing agent behavior over long time scales by transporting value.  _Nature communications_ , 10(1):5223, 2019.
  * Janner et al. (2021) Michael Janner, Qiyang Li, and Sergey Levine.  Offline reinforcement learning as one big sequence modeling problem.  In _Advances in Neural Information Processing Systems_ , 2021.
  * Kerg et al. (2020) Giancarlo Kerg, Bhargav Kanuparthi, Anirudh Goyal, Kyle Goyette, Yoshua Bengio, and Guillaume Lajoie.  Untangling tradeoffs between recurrence and self-attention in artificial neural networks.  In Hugo Larochelle, Marc’Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien Lin (eds.), _Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual_ , 2020.  URL https://proceedings.neurips.cc/paper/2020/hash/e2065cb56f5533494522c46a72f1dfb0-Abstract.html.
  * Kingma & Welling (2014) Diederik P. Kingma and Max Welling.  Auto-encoding variational bayes.  In Yoshua Bengio and Yann LeCun (eds.), _2nd International Conference on Learning Representations, ICLR 2014, Banff, AB, Canada, April 14-16, 2014, Conference Track Proceedings_ , 2014.  URL http://arxiv.org/abs/1312.6114.
  * Laidlaw et al. (2023) Cassidy Laidlaw, Stuart Russell, and Anca Dragan.  Bridging RL theory and practice with the effective horizon.  In _Thirty-seventh Conference on Neural Information Processing Systems_ , 2023.  URL https://openreview.net/forum?id=Lr2swAfwff.
  * Lenhart & Workman (2007) Suzanne Lenhart and John Workman.  Optimal control applied to biological models.  01 2007.
  * Liu et al. (2019) Yang Liu, Yunan Luo, Yuanyi Zhong, Xi Chen, Qiang Liu, and Jian Peng.  Sequence modeling of temporal credit assignment for episodic reinforcement learning.  _arXiv preprint arXiv:1905.13420_ , 2019.
  * Ma et al. (2021) Michel Ma, Pierluca D’Oro, Yoshua Bengio, and Pierre-Luc Bacon.  Long-term credit assignment via model-based temporal shortcuts.  In _Deep RL Workshop NeurIPS 2021_ , 2021.  URL https://openreview.net/forum?id=doy35IAGewq.
  * Metz et al. (2021) Luke Metz, C Daniel Freeman, Samuel S Schoenholz, and Tal Kachman.  Gradients are not all you need.  _arXiv preprint arXiv:2111.05803_ , 2021.
  * Micheli et al. (2022) Vincent Micheli, Eloi Alonso, and François Fleuret.  Transformers are sample efficient world models.  _arXiv preprint arXiv:2209.00588_ , 2022.
  * Miller et al. (1995) W.T. Miller, R.S. Sutton, and P.J. Werbos.  _Neural Networks for Control_.  A Bradford book. MIT Press, 1995.  ISBN 9780262631617.  URL https://books.google.ca/books?id=prjMtIr_yT8C.
  * Ni et al. (2021) Tianwei Ni, Benjamin Eysenbach, and Ruslan Salakhutdinov.  Recurrent model-free rl can be a strong baseline for many pomdps.  _arXiv preprint arXiv:2110.05038_ , 2021.
  * Ni et al. (2023) Tianwei Ni, Michel Ma, Benjamin Eysenbach, and Pierre-Luc Bacon.  When do transformers shine in rl? decoupling memory from credit assignment.  _arXiv preprint arXiv:2307.03864_ , 2023.
  * Pascanu et al. (2013a) Razvan Pascanu, Caglar Gulcehre, Kyunghyun Cho, and Yoshua Bengio.  How to construct deep recurrent neural networks.  _arXiv preprint arXiv:1312.6026_ , 2013a.
  * Pascanu et al. (2013b) Razvan Pascanu, Tomás Mikolov, and Yoshua Bengio.  On the difficulty of training recurrent neural networks.  In _Proceedings of the 30th International Conference on Machine Learning, ICML 2013, Atlanta, GA, USA, 16-21 June 2013_ , volume 28 of _JMLR Workshop and Conference Proceedings_ , pp. 1310–1318. JMLR.org, 2013b.  URL http://proceedings.mlr.press/v28/pascanu13.html.
  * Radford et al. (2019) Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al.  Language models are unsupervised multitask learners.  _OpenAI blog_ , 1(8):9, 2019.
  * Rahn et al. (2023) Nathan Rahn, Pierluca D’Oro, Harley Wiltzer, Pierre-Luc Bacon, and Marc G Bellemare.  Policy optimization in a noisy neighborhood: On return landscapes in continuous control.  In _Thirty-seventh Conference on Neural Information Processing Systems_ , 2023.
  * Rezende et al. (2020) Danilo J. Rezende, Ivo Danihelka, George Papamakarios, Nan Rosemary Ke, Ray Jiang, Theophane Weber, Karol Gregor, Hamza Merzic, Fabio Viola, Jane Wang, Jovana Mitrovic, Frederic Besse, Ioannis Antonoglou, and Lars Buesing.  Causally correct partial models for reinforcement learning.  _CoRR_ , abs/2002.02836, 2020.  URL https://arxiv.org/abs/2002.02836.
  * Robine et al. (2023) Jan Robine, Marc Höftmann, Tobias Uelwer, and Stefan Harmeling.  Transformer-based world models are happy with 100k interactions.  _arXiv preprint arXiv:2303.07109_ , 2023.
  * Schmidhuber (2019) Juergen Schmidhuber.  Reinforcement learning upside down: Don’t predict rewards–just map them to actions.  _arXiv preprint arXiv:1912.02875_ , 2019.
  * Schmidhuber (1990) Jürgen Schmidhuber.  Making the world differentiable: On using self-supervised fully recurrent neural networks for dynamic reinforcement learning and planning in non-stationary environments, 1990.
  * Schmidhuber (1992) Jürgen Schmidhuber.  Learning to control fast-weight memories: An alternative to dynamic recurrent networks.  _Neural Computation_ , 4(1):131–139, 1992.
  * Schrittwieser et al. (2019) Julian Schrittwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, Timothy P. Lillicrap, and David Silver.  Mastering atari, go, chess and shogi by planning with a learned model.  _CoRR_ , abs/1911.08265, 2019.  URL http://arxiv.org/abs/1911.08265.
  * Schwarzer et al. (2020) Max Schwarzer, Ankesh Anand, Rishab Goel, R. Devon Hjelm, Aaron C. Courville, and Philip Bachman.  Data-efficient reinforcement learning with momentum predictive representations.  _CoRR_ , abs/2007.05929, 2020.  URL https://arxiv.org/abs/2007.05929.
  * Srivastava et al. (2019) Rupesh Kumar Srivastava, Pranav Shyam, Filipe Mutz, Wojciech Jaśkowski, and Jürgen Schmidhuber.  Training agents using upside-down reinforcement learning.  _arXiv preprint arXiv:1912.02877_ , 2019.
  * Suh et al. (2022) Hyung Ju Suh, Max Simchowitz, Kaiqing Zhang, and Russ Tedrake.  Do differentiable simulators give better policy gradients?  In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvári, Gang Niu, and Sivan Sabato (eds.), _International Conference on Machine Learning, ICML 2022, 17-23 July 2022, Baltimore, Maryland, USA_ , volume 162 of _Proceedings of Machine Learning Research_ , pp. 20668–20696. PMLR, 2022.  URL https://proceedings.mlr.press/v162/suh22b.html.
  * Vaswani et al. (2017) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  In Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Garnett (eds.), _Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA_ , pp. 5998–6008, 2017.  URL https://proceedings.neurips.cc/paper/2017/hash/3f5ee243547dee91fbd053c1c4a845aa-Abstract.html.
  * Werbos (1990) Paul J Werbos.  Backpropagation through time: what it does and how to do it.  _Proceedings of the IEEE_ , 78(10):1550–1560, 1990.
  * Werbos (1974) Paul John Werbos.  _Beyond regression: New tools for prediction and analysis in the behavioral sciences_.  PhD thesis, Harvard University, Cambridge, MA, 1974.
  * Williams & Zipser (1989) Ronald J. Williams and David Zipser.  A learning algorithm for continually running fully recurrent neural networks.  _Neural Computation_ , 1(2):270–280, 1989.  doi: 10.1162/neco.1989.1.2.270.
  * Wolf et al. (2019) Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, and Jamie Brew.  Huggingface’s transformers: State-of-the-art natural language processing.  _arXiv preprint arXiv:1910.03771_ , 2019.
  * Zhang et al. (2016) Saizheng Zhang, Yuhuai Wu, Tong Che, Zhouhan Lin, Roland Memisevic, Russ R Salakhutdinov, and Yoshua Bengio.  Architectural complexity measures of recurrent neural networks.  _Advances in neural information processing systems_ , 29, 2016.
  * Zhang et al. (2023) Shenao Zhang, Boyi Liu, Zhaoran Wang, and Tuo Zhao.  Model-based reparameterization policy gradient methods: Theory and practical algorithms, 2023.
  * Zheng et al. (2022) Qinqing Zheng, Amy Zhang, and Aditya Grover.  Online decision transformer.  In _international conference on machine learning_ , pp. 27042–27059. PMLR, 2022.

## Appendix

###  Appendix A Proofs

See 1

###### Proof.

The policy gradient computed through equation 2 can be expanded using the
chain rule into the following expression:

$\displaystyle\nabla_{\bm{\theta}}J^{h}({\bm{\theta}};H)$
$\displaystyle=\sum_{t=1}^{H}\frac{\partial
r(\hat{s}_{t})}{\partial\hat{s}_{t}}\sum_{k=1}^{t-1}\frac{dh(\hat{s}_{1:t-1},a_{1:t-1})}{da_{k}}\frac{\partial\pi_{\bm{\theta}}(\hat{s}_{k})}{\partial{\bm{\theta}}}$
$\displaystyle\frac{dh(\hat{s}_{1:t-1},a_{1:t-1})}{da_{k}}$
$\displaystyle=\frac{\partial h(\hat{s}_{1:t-1},a_{1:t-1})}{\partial
a_{k}}+\sum_{i=k+1}^{t-1}\frac{\partial
h(\hat{s}_{1:t-1},a_{1:t-1})}{\partial\hat{s}_{i}}\frac{dh(\hat{s}_{1:i-1},a_{1:i-1})}{da_{k}}$

We begin by showing how the norm of the policy gradient depends on
$\frac{dh(\hat{s}_{1:t},a_{1:t})}{da_{k}}.$

$\displaystyle\|\nabla_{\bm{\theta}}J^{h}({\bm{\theta}};H)\|$
$\displaystyle=\left\|\sum_{t=1}^{H}\frac{\partial
r(\hat{s}_{t})}{\partial\hat{s}_{t}}\sum_{k=1}^{t-1}\frac{dh(\hat{s}_{1:t-1},a_{1:t-1})}{da_{k}}\frac{\partial\pi_{\bm{\theta}}(\hat{s}_{k})}{\partial{\bm{\theta}}}\right\|$
By Lipschitz assumption
$\displaystyle\leq\left\|\sum_{t=1}^{H}L_{r}\sum_{k=1}^{t-1}\frac{dh(\hat{s}_{1:t-1},a_{1:t-1})}{da_{k}}L_{\pi}\right\|$
Triangle inequality
$\displaystyle\leq\sum_{t=1}^{H}L_{r}\sum_{k=1}^{t-1}\left\|\frac{dh(\hat{s}_{1:t-1},a_{1:t-1})}{da_{k}}\right\|L_{\pi}\enspace.$

Now, we show how $\|\frac{dh(\hat{s}_{1:t-1},a_{1:t-1})}{da_{k}}\|$ grows
asymptotically with $t$ and $k$.

$\displaystyle\left\|\frac{dh(\hat{s}_{1:t-1},a_{1:t-1})}{da_{k}}\right\|$
$\displaystyle=\left\|\frac{\partial h(\hat{s}_{1:t-1},a_{1:t-1})}{\partial
a_{k}}+\sum_{i=k+1}^{t-1}\frac{\partial\hat{s}_{t}}{\partial\hat{s}_{i}}\frac{dh(\hat{s}_{1:i-1},a_{1:i-1})}{da_{k}}\right\|$
Triangle inequality $\displaystyle\leq\left\|\frac{\partial
h(\hat{s}_{1:t-1},a_{1:t-1})}{\partial
a_{k}}\right\|+\sum_{i=k+1}^{t-1}\left\|\frac{\partial
h(\hat{s}_{1:t-1},a_{1:t})}{\partial\hat{s}_{i}}\frac{dh(\hat{s}_{1:i-1},a_{1:i-1})}{da_{k}}\right\|$
Cauchy-Schwarz inequality $\displaystyle\leq\left\|\frac{\partial
h(\hat{s}_{1:t-1},a_{1:t-1})}{\partial
a_{k}}\right\|+\sum_{i=k+1}^{t-1}\left\|\frac{\partial
h(\hat{s}_{1:t-1},a_{1:t-1})}{\partial\hat{s}_{i}}\right\|\left\|\frac{dh(\hat{s}_{1:i-1},a_{1:i-1})}{da_{k}}\right\|$
By assumption $\displaystyle\leq
L_{a}+\sum_{i=k+1}^{t-1}L_{s}\left\|\frac{dh(\hat{s}_{1:i-1},a_{1:i-1})}{da_{k}}\right\|\enspace.$

In the recursive expansion of the previous line, notice the linear dependency
w.r.t. $L_{a}$, and exponential dependency w.r.t. $L_{s}$ due to the
compounding products. Therefore,
$\|\frac{dh(\hat{s}_{1:t},a_{1:t})}{da_{k}}\|=O(L_{a}+L_{s}^{t-k})\enspace.$
We can then put everything together to derive the asymptotic behavior.

$\displaystyle\|\nabla_{\bm{\theta}}J^{h}({\bm{\theta}};H)\|$
$\displaystyle\leq\sum_{t=1}^{H}L_{r}\sum_{k=1}^{t-1}\left\|\frac{dh(\hat{s}_{1:t-1},a_{1:t-1})}{da_{k}}\right\|L_{\pi}$
$\displaystyle=O\left(HL_{r}+H^{2}L_{\pi}+H^{2}\left\|\frac{dh(\hat{s}_{1:H-1},a_{1:H-1})}{da_{1}}\right\|\right)$
$\displaystyle=O(HL_{r}+H^{2}L_{\pi}+H^{2}L_{a}+H^{2}L_{s}^{H})$

∎

See 1

###### Proof.

We begin by providing a formal definition of the recurrent network with its
recurrent cell being the dynamics $f$ of the MDP as:

$\displaystyle g_{f\text{-}\texttt{RNN}}(s_{1},a_{1:t})$
$\displaystyle:=f(s_{t},a_{t})$ $\displaystyle s_{t}$
$\displaystyle:=f(s_{t-1},a_{t-1})=f(g_{f\text{-}\texttt{RNN}}(s_{1},a_{1:t-1}),a_{t})$

Begin by developing the LHS of the equation.

$\displaystyle\nabla_{\bm{\theta}}^{g_{f\text{-}\texttt{RNN}}}J({\bm{\theta}};H)$
$\displaystyle=\sum_{t=1}^{H}\frac{\partial r(s_{t})}{\partial
s_{t}}\sum_{k=1}^{t-1}\frac{\partial s_{t}}{\partial
a_{k}}\frac{\partial\pi_{\bm{\theta}}(s_{k})}{\partial{\bm{\theta}}}$

where $s_{t}=f(s_{t-1},a_{t-1})$ by construction (RNN assumption). Therefore,
apply the chain rule $\frac{\partial s_{t}}{\partial a_{k}}$:

$\displaystyle\nabla_{\bm{\theta}}^{g_{f\text{-}\texttt{RNN}}}J({\bm{\theta}};H)$
$\displaystyle=\sum_{t=1}^{H}\frac{\partial r(s_{t})}{\partial
s_{t}}\sum_{k=1}^{t-1}\frac{\partial s_{k+1}}{\partial
a_{k}}\bigg{(}\prod_{i={k+1}}^{t-1}\frac{\partial s_{i+1}}{\partial
s_{i}}\bigg{)}\frac{\partial\pi_{\bm{\theta}}(s_{k})}{\partial{\bm{\theta}}}$

Now, for the RHS of the equation. The policy gradient through equations 1 can
be written as:

$\displaystyle\nabla_{\bm{\theta}}J^{f}({\bm{\theta}};H)$
$\displaystyle=\sum_{t=1}^{H}\frac{\partial r(s_{t})}{\partial
s_{t}}\sum_{k=1}^{t-1}\frac{\partial s_{k+1}}{\partial
a_{k}}\frac{\partial\pi_{\bm{\theta}}(s_{k})}{\partial{\bm{\theta}}}\left(\prod_{i={k+1}}^{t-1}\frac{\partial
s_{i+1}}{\partial s_{i}}\right).$
$\displaystyle\nabla_{\bm{\theta}}^{g_{f\text{-}\texttt{RNN}}}J({\bm{\theta}};H)$
$\displaystyle=\nabla_{\bm{\theta}}J^{f}({\bm{\theta}};H).$

∎

See 2

###### Proof.

With a slight abuse of notation, we use $g_{t}$ to replace
$g(s_{1},a_{1:t-1})$.

$\nabla_{\bm{\theta}}J^{g}({\bm{\theta}};H)=\left\|\sum_{t=1}^{H}\frac{\partial
r(\hat{s}_{t})}{\partial\hat{s}_{t}}\sum_{k=1}^{t-1}\frac{\partial
g_{t}}{\partial
a_{k}}\frac{\partial\pi(\hat{s}_{k})}{\partial{\bm{\theta}}}\right\|\leq
L_{r}L_{\pi}\left\|\sum_{t=1}^{H}\sum_{k=1}^{t-1}\frac{\partial
g_{t}}{\partial a_{k}}\right\|\leq
L_{r}L_{\pi}\sum_{t=1}^{H}\sum_{k=1}^{t-1}\left\|\frac{\partial
g_{t}}{\partial a_{k}}\right\|.$

∎

See 2.1

###### Proof.

We begin by defining the transition function of the recurrent network as
$x_{t}=\sigma(W_{x}x_{t-1})+W_{a}a_{t}+b$, where $\sigma$ is some activation
function with gradient norm bounded by
$\left\|diag(\sigma^{\prime}(x))\right\|\leq\frac{1}{\beta}$. Further,
consider a linear output cell $s_{t+1}=W_{o}x_{t}$. Now we begin by showing
that $\left\|\frac{\partial x_{t}}{\partial
a_{k}}\right\|\leq||W_{a}^{T}||(||W_{x}^{T}||\frac{1}{\beta})^{t-k}$:

$\displaystyle\frac{\partial x_{t}}{\partial a_{k}}$
$\displaystyle=\frac{\partial x_{k}}{\partial
a_{k}}\prod_{i=k}^{t-1}\frac{\partial x_{i+1}}{\partial x_{i}}$
$\displaystyle\left\|\frac{\partial x_{t}}{\partial a_{k}}\right\|$
$\displaystyle=\left\|W_{a}^{T}\prod_{i=k}^{t-1}\frac{\partial
x_{i+1}}{\partial x_{i}}\right\|$ Cauchy-Schwarz Inequality
$\displaystyle\leq||W_{a}^{T}||\prod_{i=k}^{t-1}\left\|\frac{\partial
x_{i+1}}{\partial
x_{i}}\right\|=||W_{a}^{T}||\prod_{i=k}^{t-1}\left\|W_{x}^{T}diag(\sigma^{\prime}(x_{i}))\right\|$
Cauchy-Schwarz Inequality
$\displaystyle\leq||W_{a}^{T}||\prod_{i=k}^{t-1}\left\|W_{x}^{T}\right\|\left\|diag(\sigma^{\prime}(x_{i}))\right\|$
By assumption
$\displaystyle\leq||W_{a}^{T}||\prod_{i=k}^{t-1}\left\|W_{x}^{T}\right\|\frac{1}{\beta}=||W_{a}^{T}||\left(||W_{x}^{T}||\frac{1}{\beta}\right)^{t-k}\enspace.$

We develop the expression $\frac{\partial\hat{s}_{t}}{\partial
a_{k}}=\frac{\partial\hat{s}_{t}}{\partial x_{t-1}}\frac{\partial
x_{t-1}}{\partial a_{k}}$, applying the norm and plugging in the above result
we obtain:

$\displaystyle\left\|\frac{\partial\hat{s}_{t}}{\partial
a_{k}}\right\|\leq\left\|W_{o}\right\|\left\|W_{a}^{T}\right\|\left(||W_{x}^{T}||\frac{1}{\beta}\right)^{t-k-1}$

We then use the bound found in Theorem 1:

$\displaystyle\left\|\nabla_{\bm{\theta}}J^{g}({\bm{\theta}};H)\right\|$
$\displaystyle\leq
L_{r}L_{\pi}\sum_{t=1}^{H}\sum_{k=1}^{t-1}\left\|\frac{\partial
g_{t}}{\partial
a_{k}}\right\|=L_{r}L_{\pi}\sum_{t=1}^{H}\sum_{k=1}^{t-1}\left\|\frac{\partial\hat{s}_{t}}{\partial
a_{k}}\right\|$ $\displaystyle\leq
L_{r}L_{\pi}\sum_{t=1}^{H}\sum_{k=1}^{t-1}\left\|W_{o}\right\|\left\|W_{a}^{T}\right\|\left(||W_{x}^{T}||\frac{1}{\beta}\right)^{t-k-1}=L_{r}L_{\pi}\sum_{t=1}^{H}\sum_{k=1}^{t-1}\left\|W_{o}\right\|\left\|W_{a}^{T}\right\|\eta^{t-k-1}\enspace,$

where $\eta=||W_{x}^{T}||\frac{1}{\beta}$.

$\displaystyle\left\|\nabla_{\bm{\theta}}^{g}J({\bm{\theta}};H)\right\|$
$\displaystyle\leq
L_{r}L_{\pi}\sum_{t=1}^{H}\sum_{k=1}^{t-1}\left\|W_{o}\right\|\left\|W_{a}^{T}\right\|\eta^{t-k-1}$
$\displaystyle=O(H^{2}||W_{o}||\left\|W_{a}^{T}\right\|+H^{2}\eta^{H})$
$\displaystyle=O(\eta^{H})\enspace.$

∎

See 2.2

###### Proof.

We begin with the definition of the self-attention Actions World Model. Let
$Q\in\mathbb{R}^{1\times d_{z}}$, $K\in\mathbb{R}^{n\times d_{z}}$,
$V\in\mathbb{R}^{n\times d_{o}}$, attention can be defined as :

$\text{Attention}(Q,K,V):=\text{softmax}(QK^{T})V$

A self-attention AWM can then be defined in the following way with weight
matrices $W_{q}\in\mathbb{R}^{d_{z}\times d_{a}}$,
$W_{k}\in\mathbb{R}^{d_{z}\times d_{a}}$, $W_{v}\in\mathbb{R}^{d_{s}\times
d_{a}}$.

$\texttt{ATT}_{t}(a_{1},...a_{t-1}):=\text{Attention}(a_{t-1}W_{q}^{T},a_{1:t-1}W_{k}^{T},a_{1:t-1}W_{v}^{T})=\sum_{i=1}^{t-1}c_{i}(a_{i}W_{v}^{T})\enspace=\hat{s}_{t},$

where $c_{i}=\text{softmax}_{i}(a_{t-1}W_{q}^{T}W_{k}a_{1:t-1}^{T})$. The
subscript on the softmax operator indicates the $i$th index. Now, we begin by
showing the expression for $\frac{\partial g_{t}}{\partial a_{k}}$:

$\displaystyle\frac{\partial g_{t}}{\partial a_{k}}$
$\displaystyle=\sum_{i=1}^{t-1}\frac{\partial}{\partial
a_{k}}c_{i}(a_{i}W_{v}^{T})$ $\displaystyle=\sum_{i=1}^{t-1}\frac{\partial
c_{i}}{\partial
a_{k}}(a_{i}W_{v}^{T})+c_{i}\frac{\partial(a_{i}W_{v}^{T})}{\partial a_{k}}$
By Softmax derivative
$\displaystyle=c_{k}W_{v}^{T}+\sum_{i=1}^{t-1}\bigg{(}c_{i}(1\\{i=k\\}-c_{k})(a_{i}W_{v}^{T})\bigg{)}\enspace.$

Then, we take the norm:

$\displaystyle\left\|\frac{\partial g_{t}}{\partial a_{k}}\right\|$
$\displaystyle\leq\|c_{k}W_{v}^{T}\|+\sum_{i=1}^{t-1}\left\|c_{i}(1\\{i=k\\}-c_{k})(a_{i}W_{v}^{T})\right\|\enspace\text{Triangle
Inequality}$ Cauchy-Schwarz Inequality
$\displaystyle\leq\|c_{k}W_{v}^{T}\|+\sum_{i=1}^{t-1}\left\|c_{i}(1\\{i=k\\}-c_{k})\right\|\left\|(a_{i}W_{v}^{T})\right\|$
By definition of Softmax
$\displaystyle\leq\|W_{v}^{T}\|+\sum_{i=1}^{t-1}\left\|(a_{i}W_{v}^{T})\right\|$
Assuming bounded actions
$\displaystyle\leq\|W_{v}^{T}\|+\alpha\sum_{i=1}^{t-1}\left\|W_{v}^{T}\right\|\enspace,\text{where
$a_{i}\leq\alpha\enspace\forall\enspace i$}\enspace.$

Finally, we use the bound derived in Theorem 1 to finalize the proof:

$\displaystyle\left\|\nabla_{\bm{\theta}}J^{g}({\bm{\theta}};H)\right\|$
$\displaystyle\leq
L_{r}L_{\pi}\sum_{t=1}^{H}\sum_{k=1}^{t-1}\bigg{(}\|W_{v}^{T}\|+\alpha\sum_{i=1}^{t-1}\|W_{v}^{T}\|\bigg{)}$
$\displaystyle=O(H^{3}\alpha\|W_{v}^{T}\|)$ $\displaystyle=O(H^{3})\enspace.$

∎

####  A.1 Tightness of Theorem 1

The result of Theorem 1 argues that policy gradients computed through equation
2 may explode as the horizon grows. To make this claim, we comment on the
tightness of the upper bound provided in Theorem 1. Consider the one-
dimensional case, where all the gradients are positive scalars, and
$\|\frac{\partial h(\hat{s}_{1:t},a_{1:t})}{\partial a_{k}}\|=L_{a}$ and
$\|\frac{\partial h(\hat{s}_{1:t},a_{1:t})}{\partial\hat{s}_{i}}\|=L_{s}$ for
all $\hat{s}_{1:t},a_{1:t},k,i$. Then, we show that

$\displaystyle\|\nabla_{\bm{\theta}}J^{h}({\bm{\theta}};H)\|$
$\displaystyle=\Omega(L_{s}^{H})$

###### Proof.

We begin the proof of this lower bound in a similar way, by considering only
one of the terms of the summation. For clarity, Let
$h_{i}=h(\hat{s}_{1:i-1},a_{1:i-1})$.

$\displaystyle\|\nabla_{\bm{\theta}}J^{h}({\bm{\theta}};H)\|$
$\displaystyle=\left\|\sum_{t=1}^{H}\frac{\partial
r(\hat{s}_{t})}{\partial\hat{s}_{t}}\sum_{k=1}^{t-1}\frac{dh_{t}}{da_{k}}\frac{\partial\pi_{\bm{\theta}}(\hat{s}_{k})}{\partial{\bm{\theta}}}\right\|$
$\displaystyle\geq\left\|\frac{\partial
r(\hat{s}_{H})}{\partial\hat{s}_{H}}\frac{dh_{H}}{da_{1}}\frac{\partial\pi_{\bm{\theta}}(\hat{s}_{1})}{\partial{\bm{\theta}}}\right\|$
$\displaystyle=\left\|\frac{\partial
r(\hat{s}_{H})}{\partial\hat{s}_{H}}\right\|\left\|\frac{dh_{H}}{da_{1}}\right\|\left\|\frac{\partial\pi_{\bm{\theta}}(\hat{s}_{1})}{\partial{\bm{\theta}}}\right\|$
$\displaystyle\geq\left\|\frac{\partial
r(\hat{s}_{H})}{\partial\hat{s}_{H}}\right\|\left\|\frac{\partial\pi_{\bm{\theta}}(\hat{s}_{1})}{\partial{\bm{\theta}}}\right\|\left\|\bigg{(}\frac{\partial
h_{H}}{\partial a_{1}}+\sum_{i=2}^{H-1}\frac{dh_{H}}{\partial
s_{i}}\frac{dh_{i}}{da_{1}}\bigg{)}\right\|\enspace.$

Consider now only the longest chain of products in the summation term of the
last line. This corresponds to the longest path of the policy gradient through
the history world model.

$\displaystyle\|\nabla_{\bm{\theta}}J^{h}({\bm{\theta}};H)\|$
$\displaystyle\geq\left\|\frac{\partial
r(\hat{s}_{H})}{\partial\hat{s}_{H}}\right\|\left\|\frac{\partial\pi_{\bm{\theta}}(\hat{s}_{1})}{\partial{\bm{\theta}}}\right\|\left\|\bigg{(}\frac{\partial
h_{H}}{\partial a_{1}}+\frac{dh_{2}}{\partial
a_{1}}\prod_{i=2}^{H-1}\frac{dh_{i}}{\partial s_{i}}\bigg{)}\right\|$
$\displaystyle=\Omega(L_{a}+L_{s}^{H})\enspace.$

∎

####  A.2 Tightness of Corollary 2.1

The result of Corollary 2.1 argues that policy gradients through Markovian
models may explode as the horizon grows. In order to make this claim, we
comment on the tightness of the bound. Just like for Theorem 1, consider the
one-dimensional case, where all the gradients are positive scalars and assume
further the gradient norm of the activation function $\sigma$ is lower bounded
by $\gamma>0$ (which is the case for popular activation functions such as
sigmoid). We will show that, given these conditions, the bound can be written
as a lower bound:

$\displaystyle\|\nabla_{\bm{\theta}}J^{g_{\texttt{RNN}}}({\bm{\theta}};H)\|=\Omega(\eta^{H})\enspace.$

###### Proof.

We begin the proof by lower bounding the expression for the recurrent cell
$\frac{\partial x_{t}}{\partial a_{k}}$:

$\displaystyle\frac{\partial x_{t}}{\partial a_{k}}$
$\displaystyle\geq||W_{a}^{T}||(||W_{x}^{T}||\gamma)^{t-k}$
$\displaystyle\left\|\frac{\partial\hat{s}_{t}}{\partial a_{k}}\right\|$
$\displaystyle\geq\|W_{o}\|||W_{a}^{T}||(||W_{x}^{T}||\gamma)^{t-k}\enspace.$

The bound found in Theorem 2 can also be rewritten as an equality given these
additional assumptions, since the Cauchy-Schwarz and triangle inequalities
become equalities in the one-dimensional positive case:
$\|\nabla_{\bm{\theta}}J^{g}({\bm{\theta}};H)\|=L_{r}L_{\pi}\sum_{t=1}^{H}\sum_{k=1}^{t-1}\left\|\frac{\partial
g_{t}}{\partial a_{k}}\right\|$. Finally, putting these new results together,
we find that

$\displaystyle\|\nabla_{\bm{\theta}}J^{g_{\texttt{RNN}}}({\bm{\theta}};H)\|$
$\displaystyle\geq
L_{r}L_{\pi}\sum_{t=1}\sum_{k=1}^{t-1}\|W_{o}\|\|W_{a}^{T}\|(\|W_{x}^{T}\|\gamma)^{t-k}$
$\displaystyle\geq
L_{r}L_{\pi}\|W_{o}\|\|W_{a}^{T}\|(\|W_{x}^{T}\|\gamma)^{H-1}$
$\displaystyle=\Omega(\eta^{H})\enspace,$

where $\eta=\|W_{x}^{T}\|\gamma$. ∎

###  Appendix B Training an AWM with Teacher Forcing

The policy gradient framework provided by Actions World Models can also be
used to interpret training design choices used in RL through the lens of
sequence modeling. As a concrete example, recall that the transition function
of a Markovian model can be trained with the following loss:

$\displaystyle\ell^{\hat{f}}(\tau;{\bm{\psi}})=\sum_{t=1}^{H-1}\|s_{t+1}^{\tau}-\hat{f}_{\bm{\psi}}(s_{t}^{\tau},a_{t}^{\tau})\|^{2}\enspace.$
(6)

We present the following remark.

######  Remark 1.

Training an AWM instantiated with an RNN to predict states using teacher
forcing is equivalent to training its recurrent cell with the loss function
from Equation 6.

Teacher forcing (Williams & Zipser, 1989) is an algorithm used to train
recurrent neural networks. Instead of autoregressively unrolling them during
training, it enforces the outputs (i.e., in this case, the environment/RNN
states) to be the ones coming from a ground-truth sequence. In other words,
when training an AWM as an RNN with teacher forcing, we are essentially
training its recurrent cell as a one-step model, as traditionally done in
model-based RL applied to a Markovian setting.

###  Appendix C Reinforcement Learning Hyperparameters

####  C.1 Model-based Hyperparameters

All model-based backpropagation-based policy optimization methods, including
the Markovian agent, AWMs and HWMs were sweeped on three entropy constants for
exploration: $[0.1,0.01,0.001]$, and the best performing results are reported.
Indeed, in addition to the objective function $J$, the policies are optimized
with an additional entropy maximization factor as done in Amos et al. (2021)
for exploration. All actions, and states (for History World Models) are
embedded with a linear layer with an output size of 72. The RNN is initialized
with two hidden layers, with a hidden layer size of (64, 64). The self-
attention transformers use a similar architecture to the GPT-2 model (Radford
et al., 2019) implemented by the Hugging Face Transformer library (Wolf et
al., 2019). Our transformers stack 2 layers and 3 heads of self-attention
modules, with hidden layers of size 64. Timesteps are added as an input to
every input of all world models, since we are in a finite-horizon setting. The
Markovian world models predicts transitions as a difference function:
$\hat{s}_{t}=\hat{f}(s_{t-1},a_{t-1})+s_{t-1}$, using two hidden layers of
size (64, 64), and ReLU activation functions. Gradient norms are clipped at a
value of 100 for all policy gradients. All other hyper-parameters relating to
the policy learning algorithm are shown in Table 1.

Hyperparameter | Value
---|---
Number of Environment steps | 200000
Dynamics replay ratio | 2
Policy replay ratio | 16
Dynamics batch size | 64
Policy batch size | 16
Dynamics learning rate | 0.001
Policy learning rate | 0.0001
Replay buffer size | 1e6
Warmup steps | 1500
Table 1: Hyper-parameters for all BPO algorithms that do not pertain to the
world model hyper-parameters.

####  C.2 Model-free Hyperparameters

We use a model-free soft-actor critic Haarnoja et al. (2018) as a benchmark
for the myriad experiments. The critic of this agent is modeled as a 2-layer
MLP with 256 hidden units each. Entropy regularization is done using a
constant entropy constant. Results in Figure 5b show the best performing
results after doing a grid search over the following hyper-parameters:
learning rate $=[0.001,0.0001,0.00001]$ and entropy constant
$=[1.,0.01,0.001]$. The remaining hyperparameters are shown in Table 2.

Hyperparameter | Value
---|---
Number of environment steps | 200000
Critic replay ratio | 2
Policy replay ratio | 16
Batch size | 128
Discount factor | 0.995
Polyak averaging factor | 0.995
Replay buffer size | 1e6
Warmup steps | 1500
Table 2: Hyper-parameters for all model-free results on the Myriad
environments.

###  Appendix D Offline Experiments

The experiments showed in Figure 3 and Figure 4 are in the offline setting. In
both settings, 100000 transitions are collected for all experiments using a
uniform random policy. The world models are then trained on data sampled from
these transitions for 1000 steps. The batch size and learning rate are set to
the same values as shown in Table 1.

###  Appendix E Online Double-pendulum Experiments

The goal position of the environment is calculated using some fixed initial
angular position (initial action) for all episode lengths. In our experiments,
the optimal action is -0.4, which corresponds to an initial angle of
$-0.4\times 180^{\circ}.$ The results in the double-pendulum experiments use
the same hype-parameters as the Myriad experiments, with two exceptions: the
number of environment steps used for optimization is 100000 instead of 200000,
and the leaerning rate is sweeped for values in [0.01, 0.001, 0.0001, 0.00001,
0.000001] due to the unstable gradients. The complete learning curves are
shown in Figure 6.

Figure 6: Learning curves of different world models used for BPO on the
double-pendulum environment for different horizons. The horizons presented are
$H=[5,10,20,50,100]$ in order (10 seeds $\pm$ std).

###  Appendix F Myriad Environments

We give a short overview of each of the Myriad environments used in this work.
For more details, refer to Howe et al. (2022) and Lenhart & Workman (2007).

The underlying dynamics of each of these environments are described by a set
of ordinary differential equations, which are then discretized using Euler’s
method for discrete-time optimal control. In practice, to change a problem’s
horizon, we fix the duration of the experiment to ensure it is still
interesting and simply discretize using a smaller or larger step size. We
normalize the returns, where $0$ is the expected performance of a random
policy and $1$ is the performance of an optimal policy provided by Howe et al.
(2022)555Optimal policies are computed using trajectory optimization on the
underlying differential equations.. Our experiments are conducted on eight
environments: cancer treatment, bioreactor, mould fungicide, bacteria,
harvest, invasive plant, HIV treatment, and timber harvest. Each environment’s
state and action space varies between one and five dimensions.

Cancer treatment follows the normalized density of a cancerous tumour
undergoing chemotherapy. The actions at every time step correspond to the
strength of the chemotherapy drug at a given time. The goal is to minimize the
size of the tumour over a set fixed duration while also minimizing the amount
of drugs administered to the patient.

Bioreactor seeks to minimize the total amount of a chemical contaminant that
naturally degrades in the presence of bacteria. The actions here allow the
agent to feed the bacteria, increasing its population and increasing the rate
of the contaminants’ degradation. However, a cost is associated with feeding
the bacteria.

Mould fungicide models the concentration of a mould population. The goal is to
minimize its population by applying a fungicide, which has an associated cost
to apply.

Bacteria looks to maximize a bacteria population through the application of a
chemical nutrient that stimulates growth. On top of the associated cost of
applying the chemical nutrient, the chemical also produces a byproduct that
might, in turn, hinder bacterial growth.

Harvest models the growing population of some vegetables, and the goal is to
maximize the harvested yield of this population. While harvesting directly
contributes to the reward, it consequently slows down the population’s
exponential growth.

Invasive plant seeks to minimize the presence of an invasive plant species
through interventions that remove a proportion of the invasive population.
These actions have an associated cost.

HIV treatment follows the evolution of uninfected and infected cells in the
presence of a virus. The actions correspond to a drug administered that
affects the virus’ rate of infection. The use of the drug must also be
minimized.

Timber harvest is similar to the harvest environment, except the harvested
population is infinite. Instead, harvested timber can be converted into
capital, which can then be re-invested in the harvesting operation,
stimulating company growth. The goal is to maximize revenue.

To the best of our knowledge, this paper presents the first results using
reinforcement learning for the environments in Myriad. The final performances
of all BPO methods and the SAC agent for each of the eight environments are
shown in Figures 7. Note that not only does our method (BPO with transformer
AWM) have the best aggregate scores for long horizons, but it also achieves
the best performance for the longest horizon of 500 steps for each
environment.

Figure 7: Policy gradients through transformer AWMs give better policies for
long horizons for every environment.  Learning curves of BPO with different
world models on eight of the Myriad environments (10 seeds $\pm$ min/max)

.

####  F.1 Ablation on Stop Gradient

Most of the analysis in this work relies on the assumption that the gradient
is stopped through the policy input, as seen in equation 1, equation 2 and
equation 4. Just like it was noted in prior work Hafner et al. (2022; 2023);
Ghugare et al. (2023), we also perform an ablation on whether the stop-
gradient operator on the policy input has any major impact on the final
performance on the Myriad benchmark. The results in Figure 8 confirm previous
remarks that detaching the policy input from the policy gradient has little to
no effect on final performance.

Figure 8: Stop gradients on the policy input do not affect final performances
on the Myriad benchmark. Final performances of BPO with a transformer AWM and
a transformer HWM, with and without the full policy gradient (10 seeds $\pm$
95% C.I.). _Stop Grad_ denotes the policy gradient used throughout the main
text of this paper, while _Full Grad_ denotes the full unbiased policy
gradient without any stop gradient operators.

####  F.2 Ablation on LSTMs

Our analysis suggests that policy gradients through an AWM inherit certain
gradient properties of the underlying world model architecture. Long-short-
term memory networks (LSTMs) Hochreiter (1998) were also suggested as an
alternative to vanilla RNNs due to their favorable long-term gradients.
Indeed, we show in an ablation study that LSTM AWMs also perform quite well on
the Myriad benchmark in Figure 9. Thus, the Actions World Model is agnostic to
the specific neural network architecture, allowing any modern sequence models
to produce favorable policy gradients, which could be an interesting direction
to explore other more advanced sequence models such as state space models Gu
et al. (2022). We hypothesize that LSTMs are quite effective on environments
from Myriad due to the recurrent inductive bias that is compatible with all
MDPs. We suspect transformer AWMs show greater promise for large-scale
experiments, especially in POMDPs.

Figure 9: Ablation with LSTM Actions World Models on Myriad. Learning curves
of BPO with a transformer AWM, LSTM AWM, and a Markovian world model on eight
of the Myriad environments (10 seeds $\pm$ min/max).

####  F.3 Additional Details on Online Decision Transformer

Recently, a sequence modeling perspective of reinforcement learning has shown
promising results on various continuous control tasks in an offline RL setting
(Chen et al., 2021; Zheng et al., 2022; Janner et al., 2021). Although our
method differs significantly both conceptually and in the problems they solve
(see Appendix G), we show experimentally, for completeness, that the online
decision transformer (Zheng et al., 2022) performs poorly on the low
dimensional Myriad suite. We use the code and hyperparameters provided by
Zheng et al. (2022) in a purely online setup, with a sweep on the number of
trajectories gathered per iteration due to the online nature of our problem
setup. Importantly, decision transformers (DT) must be conditioned on the
return-to-go (RTG) to derive desired policies. Prior works (Chen et al., 2021;
Zheng et al., 2022; Janner et al., 2021) have shown that DTs are robust to
these hyperparameters and can achieve good and sometimes better performance
even when the RTG is set to an out-of-distribution return that is impossible
to attain.

We show in Figure 10 that DT performs poorly on the Myriad tasks in an online
setting. We summarize some of the takeaways from these experiments:

  * •

Decision transformers are still not well suited for a purely online training
regime, exhibiting worse performance than simple one-step model-based methods.
Moreover, the action-sequence model outperforms the decision transformer in
every single environment.

  * •

Decision transformers require some expert knowledge in many domains, and
simply overshooting the return-to-go can yield sub-optimal and sometimes
catastrophic results.

  * •

These methods do not scale well with the horizon unlike BPO with Actions World
Models. Even on relatively short horizons such as 100 time steps, the online
decision transformer’s performance quickly drops when compared to its
performance for 20 time steps.

Figure 10: Aggregate final performance of an online decision transformer on
the Myriad environments, trained on 200000 environment steps. The Expert RTG
is conditioned on the optimal return for each environment, while the Overshot
RTG is conditioned on a return two times higher than the optimal return (10
seeds $\pm$ 95% C.I.).

###  Appendix G Extended Related Work on Sequence Models in RL

Sequence models in RL have primarily been used in one of three ways. First,
sequence models can be used as history encoders in RL algorithms to maximize
returns in partially observable MDPs (POMDPs) (Hausknecht & Stone, 2015; Ni et
al., 2021), sometimes through a history-dependent world model (Hafner et al.,
2022; 2023). Second, sequence models have recently shown promise in an
imitation learning or offline reinforcement learning setting by treating MDPs
as a sequence modeling problem (Chen et al., 2021; Zheng et al., 2022; Janner
et al., 2021), usually conditioning on returns to derive desired policies.
Lastly, a separate line of work has used sequence models to reshape the reward
landscape for improved temporal credit assignment (Hung et al., 2019; Arjona-
Medina et al., 2019; Liu et al., 2019). In contrast to all of these, our
framework is the only one to use an action-only conditioned sequence model to
directly improve long-term policy gradients in MDPs with no intermediate step.
Below, we go into a detailed comparison with each area.

Comparison with history-conditioned RL methods for POMDPs. History-conditioned
encoders, modeled as sequence models, are often used in POMDPs in both model-
free (Hausknecht & Stone, 2015; Ni et al., 2021) and model-based methods
(Hafner et al., 2022; 2023). The latter is typically also concerned with
predicting observations, they are conditioned on entire history information,
while our method includes only actions. The problem setup is also different;
such methods are usually concerned with memory in POMDPs (Ni et al., 2023),
while ours seeks to improve credit assignment in MDPs.

Comparison with decision and trajectory transformers.  Decision transformers
(Chen et al., 2021; Zheng et al., 2022) and trajectory transformers (Janner et
al., 2021) take a more extreme approach, casting the entire reinforcement
learning problem as a sequence modeling one. Conversely, we specifically draw
a parallel between policy gradients and sequence models. Their sequence models
are conditioned on entire trajectories, which include states, actions, rewards
and returns. While trajectory transformers predict states just like our
Actions World Models, both trajectory and decision transformers also predict
actions. In either case, their sequence models must first be trained on an
offline dataset in an imitation learning (Janner et al., 2021), offline RL
(Chen et al., 2021), or pre-training framework (Zheng et al., 2022), while
AWMs work for an online RL setting.

Comparison with reward reshaping methods. Perhaps more relevant, prior works
have already tried harnessing advanced sequence models for improved temporal
credit assignment in RL (Hung et al., 2019; Arjona-Medina et al., 2019; Liu et
al., 2019). In these cases, the predictive power of sequence models, either an
LSTM (Arjona-Medina et al., 2019) or transformer (Hung et al., 2019; Liu et
al., 2019) are used to redistribute or augment the given reward function to
produce a surrogate reward function. This surrogate reward is then used in a
more traditional model-free RL algorithm. Again, we stress a fundamental
difference in the inputs of the sequence models, which all include state
information. Our framework also establishes a more direct path between
sequence models and credit assignment, avoiding any intermediate steps like
reward reshaping.