Title: LEARNING DYNAMIC STATE ABSTRACTIONS FOR MODEL-BASED REINFORCEMENT LEARNING Anonymous authors Paper under double-blind review

Abstract: A key challenge in model-based reinforcement learning (RL) is to synthesize computationally efficient and accurate environment models. We show that carefully designed models that learn predictive and compact state representations, also called state-space models, substantially reduce the computational costs for predicting outcomes of sequences of actions. Extensive experiments establish that state-space models accurately capture the dynamics of Atari games from the Arcade Learning Environment (ALE) from raw pixels. Furthermore, RL agents that use Monte-Carlo rollouts of these models as features for decision making outperform strong model-free baselines on the game MS PACMAN, demonstrating the benefits of planning using learned dynamic state abstractions.

1. INTRODUCTION:

Deep reinforcement learning has demonstrated remarkable progress in recent years, achieving high levels of performance across a wide array of challenging tasks, including Atari games (Mnih et al., 2015) , locomotion (Schulman et al., 2015) , and 3D navigation (Mnih et al., 2016) . Many of these advances have relied on combining deep learning methods with model-free RL algorithms. A critical drawback of this approach is the vast amount of experience required to achieve good performance, as only weak prior knowledge is encoded in the agents' networks (e.g., spatial translation invariance via convolutions). The promise of model-based reinforcement learning is to improve sample-efficiency by making use of explicit models of the environment. The idea is that given a model of the environment (which can possibly be learned in the absence of rewards or from observational data only), an agent can learn task-specific policies rapidly by leveraging this model e.g., by trajectory optimization (Betts, 1998) , search (Browne et al., 2012; Silver et al., 2016a) , dynamic programming (Bertsekas et al., 1995) or generating synthetic experiences (Sutton, 1991) . However, model-based RL algorithms typically pose strong requirements on the environment models, namely that they make predictions about the future state of the environment efficiently and accurately. Recent innovations combining model-free and model-based methods have helped to increase robustness to model imperfections. In the framework of Weber et al. (2017) , the Imagination-Augmented Agent (I2A) queries its internal, pre-trained model via Monte-Carlo rollouts under a rollout policy. It then uses features computed from these rollouts to anticipate the outcomes of taking different actions, thereby informing its decision-making. RL is used to learn to interpret the model's predictions; this was shown to greatly diminish the susceptibility of planning to model imperfections. Nevertheless, training of efficient models and integrating them into the I2A architecture remains difficult. In fact, this challenge prevented the application of I2As to environments with complex transitions (Weber et al., 2017) . In this paper we address, the I2A framework, the main challenge posed by model-based RL: training accurate, computationally efficient models on more complex domains and using them with agents. First, we consider computationally efficient state-space environment models that make predictions at a higher level of abstraction, both spatially and temporally, than at the level of raw pixel observations. Such models substantially reduce the amount of computation required to perform rollouts, as future states can be represented much more compactly. Second, in order to increase model accuracy, we examine the benefits of explicitly modeling uncertainty in the transitions between these abstract states. Finally, we explore different strategies of learning rollout policies that define the interface between agent and environment model: We consider the possibility of learning to query the internal model, for guiding the Monte-Carlo rollouts of the model towards informative outcomes.

The main contributions of the paper are as follows: 1) we provide a comprehensive comparison of deterministic and stochastic, pixel-space and state-space models on a number of challenging environments from the Arcade Learning Environment (ALE, Bellemare et al., 2013) ; 2) we demonstrate state-of-the-art environment modeling accuracy (as measured by log-likelihoods) with stochastic, state-space models that efficiently produce diverse yet consistent rollouts; 3) using state-space models, we show model-based RL results on MS PACMAN, and obtain significantly improved performance compared to strong model-free baselines, and 4) we show that learning to query the model further increases policy performance.

2. ENVIRONMENT MODELS:

In the following, for any sequence of variables x, we use x <t (or x ≤t ) to denote all elements of the sequences up to t, excluding (respectively including) x t . We write subsequences (x t , x t+1 , . . . , x s ) as x t:s . We denote the observations and rewards from the environment with o t and r t . We also refer to the observations o t as pixels or frames, to give the intuition that they are high-dimensional and highly redundant in many domains of interest. To ease the notation, in the following we will also write o t for the observations and rewards (o t , r t ) unless explicitly stated otherwise. )] (assuming undiscounted reward over a finite horizon up to T ), where E p is the expectation under the environment model p. In practice, however, this optimization is costly and brittle. Given unavoidable imperfections of p when modelling complex environments, it often leads to catastrophic planning outcomes (Talvitie, 2015) . Therefore, we focus here on using models to just predict at time t any future statistics x t+1:t+τ over a horizon τ that are useful features for decision making. We call x t+1:t+τ imaginations and τ the rollout horizon or depth. Concretely, we will assume that we are interested at every time step t in generating samples x t+1:t+τ by doing Monte-Carlo rollouts from the model p given an arbitrary sequence of actions a t:t+τ -1 (which will later be sampled from the rollout policy).

2.1. MODEL TAXONOMY:

In the following, we present different choices for the model p and for which imagination statistic x they compute (and subsequently pass to the agent). In particular, we focus on how these design choices trade-off computational efficiency with providing useful information to the agent. The structure of the models we consider are illustrated in Fig. 1 .

AUTO-REGRESSIVE MODELS:

A straight-forward choice is the family of causal, temporally auto-regressive models over the observations o t+1:t+τ , which we write in the following way:

p(o t+1:t+τ |o ≤t , a <t+τ ) = t+τ r=t+1 p(o r |f (o <r , a <r )).

If f is given by a first-in-first-out (FIFO) buffer of the last K observations and actions o r-K:r-1 , a r-K:r-1 , the above definition is a regular auto-regressive model (of order K), which we denote by AR. Rolling out AR models is slow for two reasons: 1) we have to sequentially sample, or "render", all pixels o t+1:t+τ explicitly, which is particularly computationally demanding for high-dimensional observations, and 2) vanilla AR models without any additional structure do not reuse any computations from evaluating p(o r |f (o <r , a <r )) for evaluating p(o r+1 |f (o ≤r , a ≤r )).

To speed-up AR models, we address the latter concern by considering the following model variant: we allow f to be a recurrent mapping that recursively updates sufficient statistics h r = f (h r-1 , a r-1 , o r-1 ), therefore reusing the previously computed statistics h r-1 . We call these models recurrent auto-regressive models (RAR); if f is parameterized as a neural network, RARs are equivalent to recurrent neural networks (RNNs). Although faster, we still expect Monte-Carlo rollouts of RARs to be slow, as they still need to explicitly render observations o t+1:t+τ . For both autoregressive models, the natural choice for imaginations x t+1:t+τ = o t+1:t+τ are predicted frames.

STATE-SPACE MODELS: ABSTRACTION IN SPACE:

As discussed above, rolling out ARs is computationally demanding as it requires sampling, or "rendering" all observations o t+1:t+τ . Causal state-space models (SSM) circumvent this by positing that there is a compact state representation s t that captures essential aspects of the environment on an abstract level: it is assumed that s t+1 can be "rolled out", i.e. predicted, from the previous state s t and action a t alone, without the help of previous pixels o ≤t or any action other than a t : p(s t+1 |s ≤t , a <t+τ , o ≤t+τ ) = p(s t+1 |s t , a t ). Hence SSMs allow for the following factorization of the predictive distribution: p(o t+1:t+τ |o ≤t , a <t+τ ) = t+τ r=t+1 p(s r |s r-1 , a r-1 )p(o r |s r ) p init (s t |o ≤t , a <t )ds t:t+τ .

For this model class, we choose the imaginations x t+1:t+τ ∼ p(s t+1:t+τ |o ≤t , a <t+τ ) to be sampled from the model distribution over the state representation. This choice implies that imaginations are, by construction, sufficient to generate all possible predictions, do not require sampling pixel observations and that they live in a conveniently lower-dimensional space.

Transition model We consider two flavors of SSMs: deterministic SSMs (dSSMs) and stochastic SSMs (sSSMs). For dSSMs, the latent transition s t+1 = g(s t , a t ) is a deterministic function of the past, whereas for sSSMs we transition distributions p(s t+1 |s t , a t ) that explicitly model uncertainty over the state s t+1 . sSSMs are a strictly larger model class than dSSMs, and we illustrate their difference in capacity for modelling stochastic time-series in the Appendix. We parameterize sSSMs by introducing for every t a latent variable z t whose distribution depends on s t-1 and a t-1 , and by making the state a deterministic function of the past state, action, and latent variable:

z t+1 ∼ p(z t+1 |s t , a t ), s t+1 = g(s t , a t , z t+1 ).

Observation ). We "chunk" the actions by concatenating them into a vector a t ← (a t , . . . , a t+c-1 ) , and sum the rewards r t ← c-1 s=0 r t+s . We refer to models trained on data pre-processed in this way as jumpy models. Jumpy training is a convenient way to inject temporal abstraction over at a time scale c into environment models. This approach allows us to further reduce the computational load for Monte-Carlo rollouts roughly by a factor of c.

2.3. MODEL ARCHITECTURES, INFERENCE AND TRAINING:

Here, we describe the parametric architectures for the models described above. We discuss the architecture of the sSSM in detail, and then briefly explain the modifications of this model used to implement RARs and dSSMs. The states s t , latent variables z t and observations o t are all shaped as convolutional feature maps and are generated by transition modules z t ∼ p(z t |s t-1 , a t-1 ), s t = g(s t-1 , z t , a t-1 ), and the decoder o t ∼ p(o t |s t , z t ) respectively. All latent variables are constrained to be normal with diagonal covariances. All modules consist of stacks of convolutional neural networks with ReLU nonlinearities. The transition modules use size-preserving convolutions, the decoder, size-expanding ones. To overcome the limitations of small receptive fields associated with convolutions, for modelling global effects of the environment dynamics, we use pool-and-inject layers: they perform max-pooling over their input feature maps, tile the results and concatenate them back to the inputs. Using these layers we can induce long-range spatial dependencies in the state s t . All modules are illustrated in detail in the Appendix. We train the AR, RAR and dSSM-DET models by maximum likelihood estimation (MLE), i.e. by maximizing L(θ) = log p θ (o 1:T |a 0:T -1 , ô0 ) over model parameters θ, where T = 10 and ô0 denotes some initial context (in our experiments ô0 := o -2:0 ). We initialize the state p init (s 0 |ô 0 ) with a convolutional network including an observation encoder e. This encoder e uses convolutions that reduce the size of the feature maps from the size of the observation to the size of the state. For the models containing latent variables, i.e. dSSM-VAE and sSSM, we cannot evaluate L(θ) in closed form in general. We maximize instead the evidence-based lower bound ELBO q (θ) ≤ L(θ), where q denotes an approximate posterior distribution, which reads as follows for the sSSM:

ELBO q (θ) = T t=1 E q [log p(o t |s t ) + log p(z t |s t-1 , a t-1 ) -log q(z t |s t-1 , a t-1 , o t:T )],

where θ now denotes the union of the model parameters and the parameters of q. Here, we used that the structure of the sSSM to assume without loss of generality that q is Markovian in (z t , s t ) (see Krishnan et al. (2015) for an in depth discussion). Furthermore, we restrict ourselves to the filtering distribution q(z t |s t-1 , a t-1 , o t ), which we model as normal distribution with diagonal covariance matrix. We did not observe improvements in experiments by using the full smoothing distribution q(z t |s t-1 , a t-1 , o t:T ). We share parameters between the prior and the posterior by making the posterior a function of the state s t computed by the transition module g, as follows:

z t+1 ∼ q(z t |s t , a t , o t+1 ), s t+1 = g(s t , z t+1 , a t ).

The posterior uses the observation encoder e on o t+1 ; the resulting feature maps are then concatenated to s t , and a number of additional convolutions compute the posterior mean and standard deviation of z t+1 . For all latent variable models, we use the reparametrization trick (Kingma & Welling, 2013; Rezende et al., 2014) and a single posterior sample to obtain unbiased gradient estimators of the ELBO. We can restrict the above sSSM to a dSSM-VAE, by not feeding samples of z t into the transition model g. To ensure a fair model comparison (identical number of parameters and same amount of computation), we numerically implement this by feeding the mean µ t of p(z t |s t-1 , a t-1 ) into the transition function g instead. If we also do not feed z t (but the mean µ t ) into the decoder for rendering o t , we arrive at the dSSM-DET, which does not contain any samples of z t . We implement the RAR based on the dSSM-DET by modifiying the transition model to s t+1 = g(s t , µ t+1 , a t , e(o t )), where e(•) denotes an encoder with the same architecture as the one of sSSM and dSSM-VAE.

o t . . . . . . s t|t s t+1|t • • • s t+τ |t π r π r π r π t a t V t a t | t a t + 1 | t a t + τ - 1 | t summarize model-free path

Figure 2 : The architecture of the Imagination-Augmented Agent, which computes its policy π t and value function V t , by combining information from a model-free path with information from Monte-Carlo rollouts of its environment model.

3. RL AGENTS WITH STATE-SPACE MODELS:

We now use the environment models described above for model-based RL by integrating them into the Imagination-Augmented Agents (I2A). We briefly describe the agent, for details see Weber et al. (2017) . The I2A is an RL agent with an actor-critic architecture, i.e. at each time step t, it explicitly computes its policy π(a t |o ≤t , a <t ) over the next action to take a t and an approximate value function V (o ≤t , a <t ), and it is trained using standard policy gradient methods (Mnih et al., 2016) . Its policy and value function are informed by the outputs of two separate pathways: 1) a model-free path, that tries to estimate the value and which action to take directly from the latest observation o t using a convolutional neural network (CNN); and 2) a model-based path, which we describe in the next paragraph.

The model-based path of an I2A is designed in the following way. The I2A is endowed with a pretrained, fixed environment model p. At every time t, conditioned on past observations and actions o ≤t , a <t , it uses the model to simulate possible futures ("rollouts") represented by imagnations x t+1:t+τ over a horizon τ , under a rollout policy π r . It then extracts informative features from the rollout imaginations x, and uses these, together with the results from the model-free path, to compute π and V . It has been shown that I2As are robust to model imperfections: they learn to interpret imaginations produced from the internal models in order to inform decision making as part of standard return maximization. More precisely, the model-based path is computed by executing the following steps (also see Fig. 2 ):

• The I2A updates the state s t of its internal model by sampling from the initial model distribution s t|t ∼ p init (s t |o ≤t ). We denote this sample s t|t to clearly indicate the real environment information is contained in that sample up to time t. • The I2A draws K samples x 1:K t+1:t+τ |t from the distribution p πr (x t+1:t+τ |s t|t , a ≤t ). Here, p πr denotes the model distribution with internal actions a t:t+τ |t being sampled from the rollout policy π r . For SSMs, we require the rollout policy to only depend on the state so that rollouts can be computed purely in abstract space.

• The imaginations x 1:K t+1:t+τ |t are summarized by a "summarizer" module (e.g. an LSTM), then combined with the model-free output and finally used to compute π(a t |o ≤t , a <t ) and V (o ≤t , a <t ).

Which imaginations x the model predicts and passes to the agent is a design choice, which strongly depends on the model itself. As described above, for auto-regressive models (AR, RAR), we choose the imaginations to be rendered pixel predictions o k t+1:t+τ |t . For SSM, we are free to use predicted pixels or predicted abstract states s k t+1:t+τ |t as imaginations, the latter being much cheaper to compute. Apart from the choice of environment model, a key ingredient to I2As is the choice of internal actions applied to the model. How to best design a rollout policy π r that extracts useful information from a given environment model remains an open question, which also depends on the choice of model itself. In the following and we investigate in the following different possibilities.

3.1. DISTILLATION:

In Weber et al. (2017) , the authors propose to train the rollout policy π r to imitate the agent's model-based behavioral policy π. We call the resulting agent the distillation agent. Concretely, we minimize the Kullback-Leibler divergence between π(•|o ≤t , s ≤t ) and π r (•|s t|t ):

L D [π r ] = λ D KL(π π r ) = -λ D E π [log π r (a t |s t|t )] + const,

where E π is the expectation over states and actions when following policy π. λ D is a hyperparameter that trades off reward maximization with the distillation loss.

3.2. LEARNING TO QUERY BY BACKPROPAGATION:

An obvious alternative to distillation is to learn the parameters of π r jointly with the other parameters of the agents by policy gradient methods. As the rollout actions sampled from π r are discrete random variables, this optimization would require "internal" RL -i.e. redefining the action space to include the internal actions and learning a joint policy over external and internal actions. However, we expect the credit assignment of the rewards to the internal actions to be a difficult problem, resulting in slow learning. Therefore, we take a heurisitic approach similar to Henaff et al. (2017) (and related to Bengio et al., 2013) : Instead of feeding the sampled one-hot environment action to the model, we can instead directly feed the probability vector π r (a t |s t |t ) into the environment model during rollouts. This can be considered as a relaxation of the discrete internal RL optimization problem. Concretely, we back-propagate the RL policy gradients through the entire rollout into π r . This is possible since the environment model is fully differentiable thanks to the reparametrization trick, and the simulation policy is differentiable thanks to the relaxation of discrete actions. Parameters of the environment model p are not optimized but kept constant, however. As the model was only trained on one-hot representation a t ∈ {0, 1} N , and not on continuous actions probabilities, it is not guaranteed a-priori that the model generalizes appropriately. We explore promoting rollout probabilities π r (•|s t |t ) to be close to one-hot action vectors, and therefore are numerically closer to the training data of the model, by introducing an entropy penalty.

3.3. MODULATION AGENT:

When learning the rollout policy (either by distillation or back-propagation), we learn to choose internal actions such that the simulated rollouts provide useful information to the agent. In these approaches, we do not change the environment model itself, which, by construction, aims to capture the true frequencies of possible outcomes. We can, however, go even one step further based on the following consideration: It might be beneficial for the agent to preferentially "imagine" extreme outcomes, e.g. rare (or even impossible) but highly rewarding or catastrophic transitions for a sequence of actions; hence to change the environment model itself in an informative way. For instance, in the game of MS PACMAN, an agent might profit form imagining the ghosts moving in a particularly adversarial way, in order to choose actions safely. We can combine this consideration with the learning-to-query approach above, by learning an informative joint "imagination" distribution over actions and outcomes. We implement this in the following way. First, we train an unconditional sSSM on environment transitions, i.e. a model that does not depend on the executed actions a <t (this can simply be done by not providing the actions as inputs to the components of our state-space models). As a result, the sSSM has to jointly capture the uncertainty over the environment and the policy π data (the policy under which the training data was collected) in the latent variables z. This latent space is hence a compact, distributed representation over possible futures, i.e. trajectories, under π data . We then let the I2A search over z for informative trajectories, by replacing the learned prior module p(z t |s t-1 ) with a different distribution p imag (z t |s t-1 ). The model is fully differentiable and we simply backpropagate the policy gradients through the entire model; the remaining weights of the model are left unchanged, except for those of p imag . In our experiments, we simply replace the neural network parameterizing p(z t |s t-1 ) with a new one of the same size for p imag , but with freshly initialized weights.

4. RESULTS:

Here, we apply the above models and agents to domains from the Arcade Learning Environment (ALE, Bellemare et al., 2013) . In spite of significant progress (Hessel et al., 2017) 

4.1. COMPARISON OF ENVIRONMENT MODELS:

We applied auto-regressive and state-space models to four games of the ALE, namely BOWLING, CENTIPEDE, MS PACMAN and SURROUND. These environment where chose to cover a broad range of environment dynamics. The data was obtained by a running a pre-trained baseline policy p data and collecting sequences of actions, observations and rewards a 1:T , o 1:T , r 1:T of length T = 10. Results are computed on held-out test data. We optimized model hyper-parameters (learning rate, weight decay, mini-batch size) on one game (MS PACMAN) for each model separately and report mean likelihoods over five runs with the best hyper-parameter settings. In Tab. 1, we report likelihood improvements over a baseline model, being a Variational Autoencoder (VAE) that models frames as independent (conditioned on three initial frames).

In general, we found that, although operating on an abstract level, SSMs are competitive with, or even outperform, auto-regressive models. The sSSM, which take uncertainty into account, achieves consistently higher likelihoods in all games compared to models with deterministic state transitions, namely dSSM-DET and dSSM-VAE, in spite of having the same number of parameters and operations. An example illustrating the differences in modelling capacity is shown in Fig. 17 (in the Appendix) on MS PACMAN: the prediction of dSSM-DET exhibits "sprite splitting" (and eventually, "sprite melting") at corridors, whereas multiple samples from the sSSM show that the model has a reasonable and consistent representation of uncertainty in this situation. We also report the relative computation time of rolling out, i.e. sampling from, the models. We observe that SSMs, which avoid computing pixel renderings at each rollout step, exhibit a speedup of > 5 over the standard AR model1 . We want to point out that our AR implementation is already quite efficient compared to a naive one, as it reuses costly vision pre-processing for rollouts where possible. Furthermore, we show that a jumpy sSSM, which learns a temporally and spatially abstracted state representation, is faster than the AR model by more than an order of magnitude, while exhibiting comparable performance as shown in Tab. 1. This shows that using an appropriate model architecture, we can learn highly predictive and compact dynamic state abstractions. Qualitatively, we observe that the best models capture the dynamics of ALE games well, even faithfully predicting global, yet subtle effects such as pixel representation of games scores over tens of steps (see Fig. 17 in the Appendix).

4.2. RL AGENTS WITH STATE-SPACE MODELS ON MS PACMAN:

Here, we apply the I2A to a slightly simplified variant of the MS PACMAN domain with five instead of eighteen actions. As environment models we use jumpy SSMs, since they exhibit a very favourable speed-accuracy trade-off as shown in the previous section; in fact I2As with AR models proved too expensive to run. In the following we compare the performance of I2As with different variants of SSMs, as well as various baselines. All agents we trained with an action repeat of 4 ( Mnih et al., 2015) . We report results in terms of averaged episode returns as a function of experience (in number of environment steps), averaged over the best hyper-parameter settings. All I2As do K = 5 (equal to the number of actions) rollouts per time step. Rollout depth τ was treated as a hyper-parameter and varied over τ ∈ {2, 3, 4}; this corresponds to 24, 36 and 48 environment steps (due to action repeats and jumpy training), allowing I2As to plan over a substantial horizon.

Learning curves for all agents with deterministic dSSMs are shown in Fig. 3 . Results and detailed discussion for agents with sSSMs can be found in section A.3 the Appendix. We first establish that all I2A agents, irrespective of the models they use, perform better than the model-free baseline agent; the latter is equivalent to an I2A without a model-based pathway. The improved performance of I2As is not simply due to having access to a larger number of input features: an I2A agent with an untrained environment model performs substantially worse (data not shown). A final baseline consists in using an I2A agent for which all imaginations s t+1:t+τ |t are set to the initial state representation s t|t . The agent has the exact same architecture, number of weights (forward model excluded), and operations as the I2A agent (denoted as "baseline copy model" in the figure legend). This agent performs substantially worse than the I2A agent, showing that environment rollouts lead to better decisions. It performs better however than the random model agent, which suggests that simply providing the initial state representation to the agent is already beneficial, emphasizing the usefulness of abstract dynamic state representations.

A surprising result is that I2As with the deterministic state-space models dSSM outperform their stochastic counterparts with sSSMs by a large margin. Although sSSMs capture the environment dynamics better than dSSM, learning from their outputs seems to be more challenging for the agents. We hypothesize that this could be due to the fact that we only produce only a small number of samples (5 in our simulations), resulting in highly variable features that are passed to the I2As.

For the agents with deterministic models, we find that a uniform random rollout policy is a strong baseline. It is outperformed by the distillation strategy, itself narrowly outperformed by the learningto-query strategy. This demonstrates that "imagining" behaviors different from the agents' policy can be beneficial for planning. Furthermore, we found that in general deeper rollouts with τ = 4 proved to outperfrom more shallow rollouts τ = 2, 3 for all I2As with deterministic SSMs.

A final experiment consists of running the I2A agent with distillation, but instead of providing the abstract state features s t+1:t+τ |t to the agent, we provide rendered pixel observations o t+1:t+τ |t instead (as was done in Weber et al., 2017) , and strengthen the summarizer (by adding convolutions). This model has to decode and re-encode observations at every imagination step, which makes it our slowest agent. We find that reasoning at pixel level eventually outperforms the copy and model-free baselines. It is however significantly outperformed by all variants of I2A which work at the abstract level, showing that the dynamics state abstractions, learned in an unsupervised way by a state-space model, are highly informative features about future outcomes, while being cheap to compute at the same time.

5. RELATED WORK:

Generative sequence models We build directly on a plethora of recent work exploring the continuum of models ranging from standard recurrent neural networks (RNNs) to fully stochastic models with uncertainty (Chung et al., 2015; Archer et al., 2015; Fraccaro et al., 2016; Krishnan et al., 2015; Gu et al., 2015) . Chung et al. (2015) explore a model class equivalent to what we called RARs here. Archer et al. (2015) ; Fraccaro et al. (2016) train stochastic state-space models, without however investigating their computational efficiency and their applicability to model-based RL. Most of the above work focuses on modelling music, speech or other low-dimensional data, whereas here we present stochastic sequence models on high-dimensional pixel-based observations; noteworthy exception are found in Watter et al. (2015) ; Wahlström et al. (2015) . There, the authors chose a twostage approach by first learning a latent representation and then learning a transition model in this representation. Multiple studies investigate the graphical-model structure of the prior and posterior graphs and stress the possible importance of smoothing over filtering inference distributions (e.g. Krishnan et al., 2015) ; in our investigations we did not find a difference between these distributions. Furthermore, to the best our knowledge, this is the first study applying stochastic state-space models as action-conditional environment models. Most previous work on learning simulators for ALE games apply deterministic models, and do not consider learning state-space models for efficient Monte-Carlo rollouts (Oh et al., 2015) . Chiappa et al. (2017) successfully train deterministic statespace models for ALE modelling (largely equivalent to the considered dSSMs here); they however do not explore the computational complexity advantage of SSMs, and do not study RL applications of their models.

Model-based reinforcement learning Most model-based RL with neural network models has previously focused on training the models on a given, compact state-representations. Directly learning models from pixels for RL is still an under-explored topic due to high demands on model accuracy and computational budget, but see Finn & Levine (2017) ; Watter et al. (2015) ; Wahlström et al. (2015) . Finn & Levine (2017) train an action-conditional video-prediction network and use it for model-predictive control (MPC) of a robot arm. The applied model requires explicit pixel rendering for long-term predictions and does not operate in abstract space. Agrawal et al. (2016) propose to learn a forward and inverse dynamics model from pixels with applications to robotics. Our work is related to multiple approaches in RL which aim to implicitly learn a model on the environment using model-free methods. Tamar et al. (2016) propose an architecture that is designed to learn the value-iteration algorithm which requires knowledge about environment transitions. The Predictron is another implicit planning network, trained in a supervised way directly from raw pixels, mimicking Bellman updates / iterations (Silver et al., 2016b) . A generalization of the Predictron to the controlled setting was introduced by Oh et al. (2017) . Another approach, presented by Jaderberg et al. (2016) , is to add auxiliary prediction losses to the RL training criterion in order to encourage implicit learning of environment dynamics. van Seijen et al. (2017) obtain state of the art performance on MS PACMAN with a model free architecture, but they however rely on privileged information (object identity and positions, and decomposition of the reward function).

6. DISCUSSION:

We have shown that state-space models directly learned from raw pixel observations are good candidates for model-based RL: 1) they are powerful enough to capture complex environment dynamics, exhibiting similar accuracy to frame-auto-regressive models; 2) they allow for computationally efficient Monte-Carlo rollouts; 3) their learned dynamic state-representations are excellent features for evaluating and anticipating future outcomes compared to raw pixels. This enabled Imagination-Augemented Agents to outperform strong model-free baselines. On a conceptual level, we present (to the best of our knowledge) the first results on what we termed learning-to-query: We show a learning a rollout policy by backpropagating policy gradients leads to consistent (if modest) improvements.

Here, we adopted the I2A assumption of having access to a pre-trained envronment model. In future work, we plan to drop this assumption and jointly learn the model and the agent. Also, further speeding up environment models is a major direction of research; we think that learning models with the capacity of learning adaptive temporal abstractions is a particularly promising direction for achieving agents that plan to react flexibly to their environment.

A APPENDIX:

A.1 DETAILS ON ENVIRONMENT MODELS

A.1.1 ARCHITECTURES:

We show the structures the inference distributions of the models with latent variables in Fig. 4 and Fig. 5 . .

a t-2 a t-1 s t-1 s t z t-1 z t o t-1 o t Figure 4:

The architecture of the inference model q for the dSSM-VAE.

a t-2 a t-1 s t-1 s t z t-1 z t o t-1 o t Figure 5:

The architecture of the inference model q for the sSSM.

A.1.2 DETAIL IN THE OBSERVATION MODEL:

For all models (auto-regressive and state-space), we interpret the three color channels of each pixel in the observation o t ∈ [0, 1] H×W ×3 (with frame height H and width W ) as pseudo-probabilities; we score these using their KL divergence with model predictions. We model the reward with a separate distribution: we first compute a binary representation of the reward N -1 n=0 b t,n 2 n = r t and model the coefficients b t,n as independent Bernoulli variables (conditioned on s t , z t ). We also add two extra binary variables: one for the sign of the reward, and the indicator function of the reward being equal to 0.

A.1.3 DETAILS OF NEURAL NETWORK IMPLEMENTATIONS:

Here we show the concrete neural network layouts used to implement the sSSM. We first introduce three higher level build blocks:

• a three layer deep convolutional stack conv stack : (k i , c i ) i=1,2,3 , with kernel sizes k 1 , k 2 , k 3 and channels sizes c 1 , c 2 , c 3 , shown in Fig. 6 ;

• a three layer deep residual convolutional stack res conv with fixed sizes, shown in Fig. 7 ;

• the Pool & Inject layer, shown in Fig. 8 .

Based on these building blocks, we define all modules in Fig. 9 to Fig. 14 . dSSM (trained with an approximate maximum likelihood criterion, see above) will "hedge its bets" by producing a blurry prediction. A similar result can be observed in rollouts from models trained on ALE games, see Fig. 

input k 1 × k 1 , c 1 relu k 2 × k 2 , c 2 + relu k 3 × k 3 ,

A.2.2 ARCHITECTURE:

We re-implemented closely the agent architecture presented by Weber et al. (2017) . In the following we list the changes in the architecture necessitated by the different environments and environment models.

Model-free baseline The model-free baseline consisted of a four-layer CNN operating on o t with sizes (4, 2, 16), (8, 4, 32), (4, 2, 64) and (3, 1, 64) , where (k, s, c) donates a CNN layer with square kernel size k, stride s and output channels s; each CNN layer is followed by a relu nonlinearity. The output of the CNN is flatten and passed trough a fully-connected (FC) layer with 512 hidden units; the final output is a value function approximation and the logits of the policy at time t.

Imagination-Augmented Agent (I2A) The model-free path consists of a CNN with the same size as the one of the model-free agent (including the FC layer with 512 units). The model-based path is designed as follows: The rollout outputs for each imagined time step s are encoded with a two layer CNN with sizes (4, 1, 32) and (4, 1, 16), then flattened and passed to a fully-connected (FC) layer with 128 outputs. These rollout statistics are then summarized (in reversed order) with an LSTM with 256 hidden units and concatenated with the outputs of the model-free path.

Rollout policies Trainable rollout policies that operate on the state s t are given by a two layer CNN with sizes (4, 1, 32) and (4, 1, 32), followed by an FC layer with 128 units. Pixel-based rollout policies have the same neural network sizes as the model-free baseline, except that the last two CNN layers have 32 feature maps each.

A.3 RESULTS I2A WITH STOCHASTIC STATE-SPACE MODELS:

Learning curves for I2As with sSSMs are shown in Fig. 18 . Both, I2As with learing-to-query and distillation rollout policies outperform a uniform random rollout policy. The learning-to-query agent shows weak initial performance, but eventually outperforms the other agents. This shows that learning-to-sample informative outcomes is beneficial for agent performance. 

:

For lack of time, we could not collect performance of the AR model on all games; this will be fixed in a later revision.

Figures and Tables:

Figure 3: Learning curves of different agents on the MS PACMAN environment. Model-based Imagination-Augmented Agents (I2As) outperform the model-free baseline by a large margin. Furthermore, learning the rollout policy π r , either by back-propagation or distillation provides the best results.

Figure 14: Initial state module for computing the first initial state s 0 as a function of the embedding e(o i ) for i = -2, -1, 0 of three previous observations.

Figure 15: Learning curves of the environment models on MS PACMAN.

17

.

APPENDIX: DETAILS ON AGENTSA.2.1 MS PACMAN ENVIRONMENT VARIANTFor the RL experiments in the paper, we consider a slightly simplified version of the MS PACMAN environment with only five actions (UP, LEFT, DOWN, RIGHT, NOOP). Furthermore, all agents have an action-repeat of four, and only observe every fourth frame from the environment.

Figure 16: Rollouts from a deterministic (dSSM, above) and a stochastic (sSSM, below) state-space model trained on a bouncing ball dataset with diffusion.

Figure 18: Model results.

Figure 1: The graphical models representing the architectures of different environment models. Boxes are deterministic nodes, circles are random variables and filled circles represent variables observed during training.

Table content: <table><tr><td>RAR</td><td>a t-2</td><td>a t-1</td><td>dSSM-DET</td><td>a t-2</td><td>a t-1</td></tr><tr><td/><td>s t-1</td><td>s t</td><td/><td>s t-1</td><td>s t</td></tr><tr><td/><td>o t-1</td><td>o t</td><td/><td>o t-1</td><td>o t</td></tr><tr><td>dSSM-VAE</td><td>a t-2</td><td>a t-1</td><td>sSSM</td><td>a t-2</td><td>a t-1</td></tr><tr><td/><td>s t-1</td><td>s t</td><td/><td>s t-1</td><td>s t</td></tr><tr><td/><td>z t-1</td><td>z t</td><td/><td>z t-1</td><td>z t</td></tr><tr><td/><td>o t-1</td><td>o t</td><td/><td>o t-1</td><td>o t</td></tr></table>

Given action a t , the environment transitions into a new state and returns a sample of the observation and reward of the next time step with probability p * (o t+1 |o ≤t , a ≤t ). A main challenge in model-based RL is to learn an accurate and efficient model p of the environment p * . Given a perfect model p ≈ p * and unlimited computational resources, an agent could e.g. perform in principle a brute-force search for the optimal open-loop policy a * t:T -1 in any state o ≤t by computing argmax a t:T -1 E p [

Table content: <table/>

model The observation model, or decoder, computes the conditional distribution p(o t |•). It either takes as input the state s t (deterministic decoder), or the state s t and latent z t (stochastic decoder). For sSSMs, we always use the stochastic decoder. For dSSMs, we can use either the deterministic decoder (dSSM-DET), or the stochastic decoder (dSSM-VAE). The latter can capture joint uncertainty over pixels in a given observation o t , but not across time steps. Further details can be found in section A.1.2 in the Appendix.2.2 JUMPY MODELS: ABSTRACTION IN TIMETo further reduce the computational time required for sampling a rollout of horizon τ , we also consider modelling environment transitions at a coarser time scale. To this end, we sub-sample observations by a factor of c, i.e. for τ = τ /c , we replace sequences (o t , o t+1 , . . . , o t+τ ), by the subsampled sequence (o t , o t+c , o t+2c , . . . , o t+τ c

Table content: <table/>

, some games are still considered very challenging environments for RL agents, e.g. MS PACMAN, especially when Improvement of test likelihoods of environment models over a baseline model (standard variational autoencoder, VAE), on 4 different ALE domains. Stochastic models with state uncertainty (RAR, sSSM) outperform models without uncertainty representation. Furthermore, statespace models (dSSM, sSSM) show a substantial speed-up over auto-regressive models. Results are given as mean ± standard deviation, in units of 10 -3 • nats • pixel -1 . not using any privileged information. All results are based on slightly cropped but full resolution ALE observations, i.e. o t ∈ [0, 1] 200×160×3 .

Table content: <table><tr><td>Model</td><td colspan="5">BOWLING CENTIPEDE MS PACMAN SURROUND rel. speed</td></tr><tr><td>AR</td><td>-</td><td>-</td><td>1.9 ± --</td><td>-</td><td>1.0×</td></tr><tr><td>RAR</td><td>-0.9 ± 3.4</td><td>5.6 ± 0.3</td><td>4.3 ± 0.5</td><td>-4.7 ± 12.2</td><td>2.0×</td></tr><tr><td>dSSM-DET</td><td>0.4 ± 0.0</td><td>3.5 ± 0.2</td><td>0.4 ± 0.3</td><td>-0.4 ± 0.1</td><td>5.2×</td></tr><tr><td>dSSM-VAE</td><td>0.5 ± 0.0</td><td>5.0 ± 1.3</td><td>2.4 ± 3.0</td><td>0.7 ± 0.0</td><td>5.2×</td></tr><tr><td>sSSM</td><td>0.6 ± 0.0</td><td>5.6 ± 1.0</td><td>4.3 ± 0.3</td><td>0.9 ± 0.2</td><td>5.2×</td></tr><tr><td>sSSM (jumpy)</td><td>-</td><td>-</td><td>3.0 ± 2.0</td><td>-</td><td>13.6×</td></tr></table>

c 3 output Figure6: Definition of the basic convolutional stack conv stack : (k i , c i ) i=1,2,3 with kernel size parameters k 1,2,3 and channel parameters c 1,2,3 . Here, a box with the label k i × k i , c i denotes a convolution with a square kernel of size k i with c i output channels; strides are always 1 × 1.We train a standard DQN agents on the four games BOWLING, CENTIPEDE, MS PACMAN and SURROUND from the ALE as detailed byMnih et al. (2015) using the original action space of 18 actions. After training, we collect a training set of 10 8 and a test set of 10 7 environment transitions for each game by executing the learned policies. Actions are represented by one-hot vectors and are tiled to yield convolutional feature maps of appropriate size. Pixel observations o t were cropped to 200 × 160 pixels and normalized by 255 to lie in the unit cube [0, 1] 3 . Because the DQN agent were trained with an action-repeat of four, we only model every fourth frame. Decoder module for computing the log-odds statistics of the Bernoulli distributions over the pixels o t and the binary coefficients of the reward r t = Figure12: Prior module for computing mean µ zt and diagonal variance σ zt of the normal distribution p(z t |s t-1 , a t-1 ). Fig.16, after training, pixels rendered from the rollouts of a sSSM depict a plausible realization of a trajectory of the ball, whereas the dSSM produces blurry samples, as conditioned on any number of previously observed frames, the state of the ball is not entirely predictable due to diffusion. A Figure13: Posterior module for computing mean μzt and diagonal variance σzt of the normal distribution q(z t |s t-1 , a t-1 , o t ). The posterior gets as additional inputs the prior statistics µ zt , σ zt .

Table content: <table><tr><td colspan="2">at-1</td><td/><td/><td/><td/><td>st</td><td>zt</td><td>μz t</td></tr><tr><td colspan="2">st-1</td><td>concat</td><td colspan="4">concat conv stack: (1, 32), (3, 32), (3, 64)</td><td>log(1 + exp(•))</td><td>σz t</td></tr><tr><td colspan="2">e(ot)</td><td/><td>3 × 3, 24</td><td colspan="3">conv stack: (1, 32), (5, 32), (3, 64)</td></tr><tr><td colspan="2">µz t , σz t</td><td/><td>relu</td><td/><td colspan="2">depth to space (2)</td></tr><tr><td/><td/><td/><td>reshape</td><td colspan="3">conv stack: (3, 64), (3, 64), (1, 48)</td></tr><tr><td/><td colspan="2">e(o0)</td><td>linear</td><td/><td colspan="2">depth to space (4)</td></tr><tr><td>input</td><td colspan="2">3 × 3, 32 e(o-1)</td><td colspan="2">relu log-odds(bt) 5 × 5, 32</td><td colspan="2">relu log-odds(ot) 3 × 3, 64</td><td>+</td></tr><tr><td colspan="7">Figure 10: N -1 n=0 b t,n 2 n . e(o-2)</td></tr><tr><td/><td/><td/><td/><td colspan="2">e(ot)</td></tr><tr><td/><td>input</td><td colspan="2">3 × 3, 32</td><td>max pool</td><td/><td>tile</td><td>concat</td></tr><tr><td/><td/><td/><td/><td>relu</td><td/></tr><tr><td colspan="2">at-1</td><td/><td colspan="4">conv stack: (3, 32), (5, 32), (3, 64)</td></tr><tr><td colspan="2">st-1</td><td>concat</td><td>res conv</td><td colspan="2">relu space to depth (2)</td><td>pool &amp; inject</td><td>res conv</td></tr><tr><td>zt</td><td/><td/><td colspan="4">conv stack: (3, 16), (5, 16), (3, 64)</td></tr><tr><td/><td/><td/><td/><td colspan="2">space to depth (4)</td></tr><tr><td colspan="4">A.1.4 COLLECTION OF TRAINING DATA</td><td/><td/></tr><tr><td colspan="2">at-1</td><td/><td/><td/><td/><td>µz t</td></tr><tr><td colspan="3">A.1.5 TRAINING DETAILS st-1 concat</td><td colspan="4">conv stack: (1, 32), (3, 32), (3, 64)</td><td>log(1 + exp(•))</td><td>σz t</td></tr><tr><td colspan="7">All models were optimized using Adam (Kingma &amp; Ba, 2014) with a mini-batch size of 16.</td></tr></table>

Peer review (iclr2018):
 The paper proposes a method for inferring dynamical models from partial observations, that can later be used in model-based RL algorithms such as I2A. The essence of the method is to perform variational inference of the latent state, representing its distribution as Gaussian, and to use it in an ELBO to infer the dynamics of state and observation. While this is an interesting approach, many of the architectural choices involved seem arbitrary and unjustified. This wouldn*t be so bad if they were justified by empirical success rather than principled design, but I*m also a bit skeptical of the strength of the results. A few examples of such architectural choices: 1. What*s the significance of separating the stochastic state transition into a stochastic choice of z and a transition g deterministic in z? 2. How is the representational power affected by having only the observations depend on z? What*s the intuition behind calling this model VAE, when sSSM is also trained with variational inference? 3. What is gained by using pool-and-inject layers? By the way, is this a novel component? If so please elaborate, if not please cite. As for the strength of the results, in Table 1 the proposed methods don*t seem to outperform "RAR" (i.e., RNN) in expected value. They do seem to have lower variance, and the authors would do well to underline the importance of this. In Figure 3, it*s curious that the model-free baseline remains unnamed, as it also does in the text and appendix. This makes it hard to evaluate whether the significant wins are indicative of the strength of the proposed method. Finally, a notational point that the authors should really get right, is that conditioning on future actions naively changes the distribution of the current state or observation in ways they didn*t intend. The authors intended for actions to be used as "interventions", i.e. a-causally, and should denote this conditioning by some sort of "do" operator.