On Trade-offs of Image Prediction in Visual Model-Based Reinforcement Learning

1 INTRODUCTION . The key component of any model-based reinforcement learning ( MBRL ) methods is the predictive model . In visual MBRL , this model predicts the future observations ( i.e. , images ) that will result from taking different actions , enabling the agent to select the actions that will lead to the most desirable outcomes . These features enable MBRL agents to perform successfully with high dataefficiency ( Deisenroth & Rasmussen , 2011 ) in many tasks ranging from healthcare ( Steyerberg et al. , 2019 ) , to robotics ( Ebert et al. , 2018 ) , and playing board games ( Schrittwieser et al. , 2019 ) . More recently , MBRL methods have been extended to settings with high-dimensional observations ( i.e. , images ) , where these methods have demonstrated good performance while requiring substantially less data than model-free methods without explicit representation learning ( Watter et al. , 2015 ; Finn & Levine , 2017 ; Zhang et al. , 2018 ; Hafner et al. , 2018 ; Kaiser et al. , 2020 ) . However , the models used by these methods , also commonly known as World Models ( Ha & Schmidhuber , 2018 ) , vary in their fundamental design . For example , some recent works only predict the expected reward ( Oh et al. , 2017 ) or other low-dimensional task-relevant signals ( Kahn et al. , 2018 ) , while others predict the images as well ( Hafner et al. , 2019 ) . Along a different axis , some methods model the dynamics of the environment in the latent space ( Hafner et al. , 2018 ) , while some other approaches model autoregressive dynamics in the observation space ( Kaiser et al. , 2020 ) . Unfortunately , there is little comparative analysis of how these design decisions affect performance and efficiency , making it difficult to understand the relative importance of the design decisions that have been put forward in prior work . The goal of this paper is to understand the trade-offs between the design choices of model-based agents . One basic question that we ask is : does predicting images actually provide a benefit for MBRL methods ? A tempting alternative to predicting observations is to simply predict future rewards , which , in principle , gives a sufficient signal to infer all task-relevant information . However , as we will see , predicting images has clear and quantifiable benefits – in fact , we observe that accuracy in predicting observations correlates more strongly with control performance than accuracy of predicting rewards . Our goal is to specifically analyze the design trade-offs in the models themselves , decoupling this as much as possible from the confounding differences in the algorithm . While a wide range of different algorithms have been put forward in the literature , we restrict our analysis to arguably the simplest class of MBRL methods , which train a model and then use it for planning without any explicit policy . While this limits the scope of our conclusions , it allows us to draw substantially clearer comparisons . The main contributions of this work are two-fold . First , we provide a coherent conceptual framework for high-level design decisions in creating models . Second , we investigate how each one of these choices and their variations can affect the performance across multiple tasks . We find that : 1 . Predicting future observations ( i.e . images ) leads to significant task performance improvement compared to only predicting rewards . And , somewhat suprisingly , image prediction accuracy correlates more strongly with downstream task performance than reward prediction accuracy . 2 . We show how this phenomenon is related to exploration : - Some of the lower-scoring models on standard benchmarks that require exploration will perform the same as the best-performing models when trained on the same training data . - In the absence of exploration , models that fit the data better usually perform better on the downstream task , but surprisingly , these are often not the same models that perform the best when learning and exploring from scratch . 3 . A range of design decisions that are often considered crucial , such as the use of latent spaces , have little effect on task performance . These findings suggest that performance and exploration place important and potentially contradictory requirements on the model . We will open-source our implementation that can be used to reproduce the experiments and can be further extended to other environments and models . 2 RELATED WORK . MBRL is commonly used for applications where sample efficiency is essential , such as real physical systems ( Deisenroth & Rasmussen , 2011 ; Deisenroth et al. , 2013 ; Levine et al. , 2016 ) or healthcare ( Raghu et al. , 2018 ; Yu et al. , 2019 ; Steyerberg et al. , 2019 ) . In this work , our focus is on settings with a high-dimensional observation space ( i.e. , images ) . Scaling MBRL to this setting has proven to be challenging ( Zhang et al. , 2018 ) , but has shown recent successes ( Hafner et al. , 2019 ; Watter et al. , 2015 ; Levine et al. , 2016 ; Finn et al. , 2016b ; Banijamali et al. , 2017 ; Oh et al. , 2017 ; Zhang et al. , 2018 ; Ebert et al. , 2018 ; Hafner et al. , 2018 ; Dasari et al. , 2019 ; Kaiser et al. , 2020 ) . Besides these methods , a number of works have studied MBRL methods that do not predict pixels , and instead directly predict future rewards ( Oh et al. , 2017 ; Liu et al. , 2017 ; Schrittwieser et al. , 2019 ; Sekar et al. , 2020 ) , other reward-based quantities ( Gelada et al. , 2019 ) , or features that correlate with the reward or task ( Dosovitskiy & Koltun , 2016 ; Kahn et al. , 2018 ) . It might appear that predicting rewards is sufficient to perform the task , and a reasonable question to ask is whether image prediction accuracy actually correlates with better task performance or not . One of our key findings is that predicting images improves the performance of the agent , suggesting a way to improve the task performance of these methods . Visual MBRL methods must make a number of architecture choices in structuring the predictive model . Some methods investigate how to make the sequential high-dimensional prediction problem easier by transforming pixels ( Finn et al. , 2016a ; De Brabandere et al. , 2016 ; Liu et al. , 2017 ) or decomposing motion and content ( Tulyakov et al. , 2017 ; Denton et al. , 2017 ; Hsieh et al. , 2018 ; Wichers et al. , 2018 ; Amiranashvili et al. , 2019 ) . Other methods investigate how to incorporate stochasticity through latent variables ( Xue et al. , 2016 ; Babaeizadeh et al. , 2018 ; Denton & Fergus , 2018 ; Lee et al. , 2018 ; Villegas et al. , 2019 ) , autoregressive models ( Kalchbrenner et al. , 2017 ; Reed et al. , 2017 ; Weissenborn et al. , 2019 ) , flow-based approaches ( Kumar et al. , 2019 ) and adversarial methods ( Lee et al. , 2018 ) . However , whether prediction accuracy actually contributes to MBRL performance has not been verified in detail on image-based tasks . We find a strong correlation between image prediction accuracy and downstream task performance , suggesting video prediction is likely a fruitful area of research for improving visual MBRL . 3 A FRAMEWORK FOR VISUAL MODEL-BASED RL . MBRL methods must model the dynamics of the environment conditioned on the future actions of the agent , as well as current and preceding observations . The data for this model comes from trials in the environment . In an online setting , collecting more data and modeling the environment happens iteratively , and in an offline setting , data collection happens once using a predefined policy . We are interested in cases where the observation state is pixels , which is usually not enough to reveal the exact state of the environment . Since our goal is specifically to analyze the design trade-offs in the models themselves , we focus our analysis only on MBRL methods which train a model and use it for planning , rather than MBRL methods that learn policies or value functions . We defer discussion and analysis of policies and value functions to future work . Therefore , we consider the problem as a partially observable Markov decision process ( POMDP ) with a discrete time step t ∈ [ 1 , T ] , hidden state st , observed image ot , continuous action at and scalar immediate reward rt . The dynamics are defined as a transition function st ∼ p ( st | st−1 , at−1 ) , observation function ot ∼ p ( ot | st ) and reward function rt ∼ p ( rt | st ) . In RL , the goal is to find a policy p ( at|o≤t , r≤t , a≤t ) that maximizes the expected discounted sum of future rewards Ep [ ∑T t=1 γrt ] , where γ is the discount factor , T is the task horizon and the expectation is over actions sampled from the policy . In MBRL , we approximate the expected reward by predicting its distribution conditioned on the previous observations and the future actions p ( rt|o≤t , a≥t ) , and then search for a high-reward action sequences via a policy optimization method such as proximal policy optimization ( PPO ) ( Schulman et al. , 2017 ) or a planning algorithm like the cross-entropy method ( CEM ) ( Rubinstein , 1997 ; Chua et al. , 2018 ) . In this paper we focus on the latter . Given the assumption that the environment is partially observable and the ground-truth states are not accessible , the models in visual MBRL can be divided into five main categories , depending on training signals and learned representations ( Figure 1 ) . The first category approximates the expected rewards conditioned on future actions and previous observations p ( rt | o≤t , a≥t ) without explicitly predicting images ( Oh et al. , 2017 ; Dosovitskiy & Koltun , 2016 ; Racanière et al. , 2017 ; Liu et al. , 2017 ; Kahn et al. , 2018 ; Schrittwieser et al. , 2019 ) . In contrast , the next four categories learn to predict the next observations ôt+1 , in addition to the reward , which results in a learned latent representation . Given this latent space , it is possible to : ( 1 ) Model the transition function of the environment in the observation space : ôt+1 ∼ p ( ot | o≤t , a≥t ) or directly in the latent space : ht ∼ p ( ht | ht−1 , at ) where ht is the hidden state of the model at time step t ( 2 ) Predict the future reward using the learned latent space : rt ∼ p ( rt | ht ) or from the predicted future observation : rt ∼ p ( rt | ôt+1 ) . 4 EXPERIMENTS AND ANALYSIS . Our goal in this paper is to study how each axis of variation in the generalized MBRL framework , described in Section 3 , impacts the performance of the agent . To this end , we will not only study endto-end performance of each model on benchmark tasks , but also analyze how these models compare when there is no exploration ( i.e . trained from the same static dataset ) , evaluate the importance of predicting images , and analyze a variety of other model design choices .

Summary: This paper empirically investigates the effect of different reward-prediction models for model-based reinforcement learning (MBRL) in the particular context of decision-time planning with the cross-entropy method and visual tasks. To that end, five models are evaluated: four models that predict rewards and observations jointly, while one model only learns to predict rewards but without observation prediction. The authors claim that predicting observations in addition to rewards improves performance of MBRL in their setting (in terms of cumulative reward maximization). This hypothesis is tested on a bunch of tasks from the DeepMind Control Suite, both in an online and offline setting. The authors conclude, based on their experiments, that their hypothesis can be confirmed.