summaries: This paper claims that one only needs a reward prediction model to learn a good latent representation for model-based reinforcement learning. They introduce a method that learns a latent dynamics model exclusively from multi-step reward prediction, then use MPC to plan directly in the latent space. They claim this is sample efficient in the model-based way, and is more useful than predicting full states. They learn a model that predicts only current and future rewards conditioned on action sequences, and that observation reconstruction is unnecessary to learn a good latent space. They provide planning performance guarantees for approximate latent reward prediction models.  This paper presents a technique for model based RL/planning with latent dynamics models, which learns the latent model only using reward prediction. This is in contrast to existing work which generally use a combination of reward prediction and state reconstruction to learn the latent model. The paper suggests that by removing the state reconstruction loss, the agent can learn to ignore irrelevant parts of the state, which should enable better performance in settings where state reconstruction is challenging.  This paper proposes a novel algorithm for planning on specific domains through latent reward prediction. The proposed model uses an encoder to learn embedding the state to the latent state, a forward dynamics function to learn dynamical system in latent state space, and a reward function to estimate the reward given a latent state and an action. Using these functions, the authors define the objective using the mean-squared error between true and multi-step prediction of rewards. To justify the proposed method, the authors provide a theoretical analysis and experimental results on specific RL domains, multi-pendulum and multi-cheetah, which contain irrelevant aspects of the state.
abstractText: Model-based reinforcement learning methods typically learn models for highdimensional state spaces by aiming to reconstruct and predict the original observations. However, drawing inspiration from model-free reinforcement learning, we propose learning a latent dynamics model directly from rewards. In this work, we introduce a model-based planning framework which learns a latent reward prediction model and then plan in the latent state-space. The latent representation is learned exclusively from multi-step reward prediction which we show to be the only necessary information for successful planning. With this framework, we are able to benefit from the concise model-free representation, while still enjoying the data-efficiency of model-based algorithms. We demonstrate our framework in multi-pendulum and multi-cheetah environments where several pendulums or cheetahs are shown to the agent but only one of them produces rewards. In these environments, it is important for the agent to construct a concise latent representation to filter out irrelevant observations. We find that our method can successfully learn an accurate latent reward prediction model in the presence of the irrelevant information while existing model-based methods fail. Planning in the learned latent state-space shows strong performance and high sample efficiency over model-free and model-based baselines.
references.authors: Ershad Banijamali Rui Shu Mohammad Ghavamzadeh Hung Bui Ali Ghodsi
references.title: Robust locallylinear controllable embedding
references.venue: arXiv preprint arXiv:1710.05373,
references.authors: Greg Brockman Vicki Cheung Ludwig Pettersson Jonas Schneider John Schulman Jie Tang Wojciech Zaremba
references.title: URL http://arxiv.org/ abs/1606.01540
references.venue: Openai gym. CoRR,
references.authors: Yuri Burda Harri Edwards Deepak Pathak Amos Storkey Trevor Darrell Alexei A Efros
references.title: Large-scale study of curiosity-driven learning
references.venue: arXiv preprint arXiv:1808.04355,
references.authors: Yuri Burda Harrison Edwards Deepak Pathak Amos J. Storkey Trevor Darrell Alexei A. Efros
references.title: Large-scale study of curiosity-driven learning
references.venue: CoRR, abs/1808.04355,
references.authors: Yevgen Chebotar Mrinal Kalakrishnan Ali Yahya Adrian Li Stefan Schaal Sergey Levine
references.title: Path integral guided policy search
references.venue: IEEE international conference on robotics and automation (ICRA),
references.authors: Marc Deisenroth Carl E Rasmussen
references.title: Pilco: A model-based and data-efficient approach to policy search
references.venue: In Proceedings of the 28th International Conference on machine learning
references.authors: Chelsea Finn Sergey Levine
references.title: Deep visual foresight for planning robot motion
references.venue: IEEE International Conference on Robotics and Automation (ICRA),
references.authors: Chelsea Finn Xin Yu Tan Yan Duan Trevor Darrell Sergey Levine Pieter Abbeel
references.title: Deep spatial autoencoders for visuomotor learning
references.venue: IEEE International Conference on Robotics and Automation (ICRA),
references.authors: Carles Gelada Saurabh Kumar Jacob Buckman Ofir Nachum Marc G Bellemare
references.title: Deepmdp: Learning continuous latent space models for representation learning
references.venue: In International Conference on Machine Learning,
references.authors: Tuomas Haarnoja Aurick Zhou Kristian Hartikainen George Tucker Sehoon Ha Jie Tan Vikash Kumar Henry Zhu Abhishek Gupta Pieter Abbeel
references.title: Soft actor-critic algorithms and applications
references.venue: arXiv preprint arXiv:1812.05905,
references.authors: Danijar Hafner Timothy Lillicrap Ian Fischer Ruben Villegas David Ha Honglak Lee James Davidson
references.title: Learning latent dynamics for planning from pixels
references.venue: In International Conference on Machine Learning,
references.authors: Brian Ichter Marco Pavone
references.title: Robot motion planning in learned latent spaces
references.venue: IEEE Robotics and Automation Letters,
references.authors: Diederik P. Kingma Max Welling
references.title: Auto-encoding variational bayes
references.venue: CoRR, abs/1312.6114,
references.authors: Ian Lenz Ross A Knepper Ashutosh Saxena
references.title: Deepmpc: Learning deep latent features for model predictive control
references.venue: In Robotics: Science and Systems. Rome, Italy,
references.authors: Sergey Levine Vladlen Koltun
references.title: Guided policy search
references.venue: In International Conference on Machine Learning, pp
references.authors: Volodymyr Mnih Koray Kavukcuoglu David Silver Andrei A Rusu Joel Veness Marc G Bellemare Alex Graves Martin Riedmiller Andreas K Fidjeland Georg Ostrovski
references.title: Human-level control through deep reinforcement learning
references.authors: Junhyuk Oh Satinder Singh Honglak Lee
references.title: Value prediction network
references.venue: In Advances in Neural Information Processing Systems,
references.authors: Reuven Y Rubinstein
references.title: Optimization of computer simulation models with rare events
references.venue: European Journal of Operational Research,
references.authors: David Silver Thomas Hubert Julian Schrittwieser Ioannis Antonoglou Matthew Lai Arthur Guez Marc Lanctot Laurent Sifre Dharshan Kumaran Thore Graepel
references.title: A general reinforcement learning algorithm that masters chess, shogi, and go through self-play
references.authors: Manuel Watter Jost Springenberg Joschka Boedecker Martin Riedmiller
references.title: Embed to control: A locally linear latent dynamics model for control from raw images
references.venue: In Advances in neural information processing systems,
references.authors: Grady Williams Nolan Wagener Brian Goldfain Paul Drews James M Rehg Byron Boots Evangelos A Theodorou
references.title: Information theoretic mpc for model-based reinforcement learning
references.venue: IEEE International Conference on Robotics and Automation (ICRA),
references.authors: Marvin Zhang Sharad Vikram Laura Smith Pieter Abbeel Matthew Johnson Sergey Levine
references.title: Solar: Deep structured representations for model-based reinforcement learning
references.venue: In International Conference on Machine Learning,
sections.heading: 1 INTRODUCTION
sections.text: Deep reinforcement learning (DRL) has demonstrated some of the most impressive results on several complex, high-dimensional sequential decision-making tasks such as video games (Mnih et al., 2015) and board games (Silver et al., 2018). Such algorithms are often model-free, i.e. they do not learn or exploit knowledge of the environment’s dynamics. Although these model-free schemes can achieve state-of-the-art performance, they often require millions of interactions with the environment. The high sample complexity restricts their applicability to real environments, where obtaining the required amount of interactions is prohibitively time-consuming. On the other hand, model-based approaches can achieve low sample complexity by first learning a model for the environment. The learned model can then be used for planning, thereby minimizing the number of environment interactions. However, model-based methods are typically limited to low-dimensional tasks or settings with certain assumptions imposed on the dynamics (Li & Todorov, 2004; Deisenroth & Rasmussen, 2011; Levine & Koltun, 2013; Chebotar et al., 2017)
  Several recent works have learned dynamics models parameterized by deep neural networks (Lenz et al., 2015; Finn & Levine, 2017; Williams et al., 2017), but such methods generally result in high-dimensional models which are unsuitable for planning. In order to efficiently apply model-based planning in high-dimensional spaces, a promising approach is to learn a low-dimensional latent state representation of the original high-dimensional space. A compact latent representation is often more conducive for planning, which can further reduce the number of environment interactions required to learn a good policy (Watter et al., 2015; Banijamali et al., 2017; Finn et al., 2016; Ichter & Pavone, 2019; Zhang et al., 2019; Hafner et al., 2019). Typically, such a model has a variational autoencoder (VAE) (Kingma & Welling, 2013)-like structure that learns a latent representation. In particular, it learns this representation through reconstruction and prediction of the full-state from the latent state. In other words, the learned low-dimensional latent state must retain a sufficient amount of information from the original state to accurately predict the state in the high-dimensional space. Although the ability to perform full-state prediction is sufficient for planning in latent space, it is not
  strictly necessary. Predicting the full-state may result in a latent state representation which contains information irrelevant to planning. For example, consider a task of maze navigation, with a TV placed in the maze, displaying images (Burda et al., 2018a). The TV content is irrelevant to the task, but the latent representation may learn features designed to predict the state of the TV.
  In RL, the objective is to maximize the cumulative sum of rewards. Given this objective, it may be excessive to learn to perform full-state prediction, if aspects of the full state have no influence on the reward. In examples like the aforementioned TV in a maze, it may be more fruitful to instead optimize a latent model to be able to perform reward prediction. In some sense, value-based modelfree methods can be seen as reward prediction models. However, they attempt to predict the optimal value which is policy dependent and requires a large number of samples to learn. Fortunately, modelbased planning methods such as model predictive control (MPC) only maximize the cumulative sum of future rewards over the choices of action sequences. This implies that what a planning method requires is a latent reward prediction model: a model which predicts current and future rewards over a horizon from a latent state under different action sequences. In this work, we introduce a model-based DRL framework where we learn a latent dynamics model exclusively from the multi-step reward prediction criterion and then use MPC to plan directly in the latent state-space. Learning the latent reward prediction model benefits from the sample-efficiency of model-based methods, and the learned latent model is more useful for planning than its full-state prediction counterparts as discussed above.
  The contributions of this paper are as follows. The proposed latent model is learned only from multi-step reward prediction. Optimizing exclusively for reward prediction allows us to circumvent the inefficiency stemming from learning irrelevant parts of the state space, while maintaining the model-based sample efficiency. We provide a performance guarantee on planning in the learned latent state-space without assumptions on the underlying task dynamics. Empirically, we demonstrate that our method outperforms model-free and model-based baselines on multi-pendulum and multi-cheetah environments where the agent must ignore a large amount of irrelevant information.
  The remainder of this paper is structured as follows. In Section 2, we present related work in modelbased DRL. In Section 3, we provide the requisite background on RL and MPC. In Section 4, we introduce our method. In Section 5, we present our experimental results, and conclude by indicating potential future directions of research.
sections.heading: 2 RELATED WORK
sections.text: Model-based DRL Model-based works have aimed to increase sample efficiency and stability by utilizing a model during policy optimization. Methods like PILCO (Deisenroth & Rasmussen, 2011), DeepMPC (Lenz et al., 2015), and MPPI (Williams et al., 2017) have exhibited promising performance on some control tasks using only a handful of episodes, however they usually have difficulty scaling to high-dimensional observations. Guided-policy search methods (Levine & Koltun, 2013; Chebotar et al., 2017) make use of local models to update a global deep policy, however these local assumptions could fail to capture global properties and suffer at points of discontinuity. For high-dimensional observations such as images, deep neural networks have been trained to directly predict the observations (Finn & Levine, 2017), but the resulting high-dimensional models limits their applicability in efficient planning.
  Learning latent state spaces for planning To make model-based methods more scalable, several works have adopted an auto-encoder scheme to learn a reduced-dimensional latent state representation. Then one may plan directly in the latent space using some conventional planning methods such as linear quadratic regulator (LQR), or rapidly exploring random tree (RRT) (Watter et al., 2015; Banijamali et al., 2017; Finn et al., 2016; Ichter & Pavone, 2019; Zhang et al., 2019). However, these methods construct latent spaces based on state reconstruction/prediction which may not be entirely relevant to the planning problem as discussed in the previous section. A recent work PlaNet (Hafner et al., 2019) offered a latent space planning approach by predicting multi-step observations and rewards. The PlaNet has shown to achieve good sample complexity and performance compared to the state-of-the-art model-free methods, but its latent model might also suffer from the aforementioned issues with irrelevant observation predictions. In this work we use a similar planning algorithm to PlaNet, however, removing the dependency of observation prediction provides a more concise latent representation for planning.
  Reward-based representations learning for RL The work of VPN (Oh et al., 2017) proposed a RL framework which learns a model to predict the future values rather than the full-states of the environment. Although VPN shares a very similar idea as our framework to not predict irrelevant information in full-state, it suffers the sample complexity issue of learning the optimal value as in model-free methods. Performance guarantees of latent models were analyzed under the DeepMDP framework in a recent paper (Gelada et al., 2019) without full-state reconstruction or prediction. With Lipschitz assumptions, their work provides performance bounds for latent representations when the two prediction losses of the current reward and the next latent state are optimized. They also show improvements in model-free DRL by adding the DeepMDP losses as auxiliary objectives. However, with only the single-step prediction loss, DeepMDP may not be suitable for making long term reward predictions necessary for planning algorithms. In addition, the latent state prediction loss may be unnecessary and could result in local minima observed in their paper. Our framework shares the same spirit of DeepMDP by learning latent dynamics only from the reward sequences, but focusing on multi-step reward prediction allows us to achieve higher performance in model-based planning.
sections.heading: 3 PRELIMINARIES
sections.heading: 3.1 PROBLEM SETUP
sections.text: In this paper we consider a discrete time nonlinear dynamical system f : S×A → S with continuous state and action spaces S ⊆ Rdn , A ⊆ Rdm , and a reward function R : S ×A → R. Then, given an admissible action at ∈ A at time t, the system state st evolves according to the dynamics
  st+1 = f(st, at) (1)
  where st+1 is the next state at time t+ 1.
  Our goal is to find a policy π : S → A that selects actions {at = π(st), t = 0, 1, . . .} so as to maximize the cumulative discounted rewards
  ∞∑ t=0 γtR(st, at) (2)
  where the initial state s0 is assumed to be fixed, and γ is the discount factor.
  Note that Equations (1) and (2) formulate an infinite horizon optimal control problem. This problem can also be viewed as a deterministic Markov decision process (MDP) by the tuple (S,A, f, R, γ).
sections.heading: 3.2 MPC PLANNING
sections.text: When a dynamics model is available, model predictive control (MPC) is a powerful planning framework for optimal control problems. A MPC agent chooses its action by online optimization over a finite planning horizon H . More specifically, at each time t, the agent solves a H-horizon trajectory optimization problem:
  max ât:(t+H−1) J(st, ât:(t+H−1)) (3a)
  subject to J(st, ât:(t+H−1)) = t+H−1∑ τ=t γτ−tR(ŝτ , âτ ) (3b)
  ŝt = st, and ŝτ+1 = f(ŝτ , âτ ) for τ = t, t+ 1, . . . , (t+H − 2) (3c) Suppose the optimal control sequence is â∗t:(t+H−1), then the MPC agent will select â ∗ t as the control
  action at time t. This method is termed model predictive control due to its use of the model f to predict the future trajectory ŝ(t+1):(t+H−1) from state st and the intended action sequence ât:(t+H−1). Let πMPC(st|f,R) denote the MPC policy using the dynamics model f and reward R. Then the MPC agent selects its action by πMPC(st|f,R) = â∗t .
sections.heading: 3.3 SAMPLE-BASED TRAJECTORY OPTIMIZATION
sections.text: From the description in Section 3.2, an MPC agent’s need to solve the trajectory optimization problem specified by Equation equation 3. In this work, we use the cross entropy method (CEM) (Rubinstein,
  1997) to solve problem equation 3. CEM is a very simple sampling-based planning method which only relies on a forward model and trajectory returns which may make direct use of our learned latent mappings somewhat like a simulator. The advantage to using a method like CEM, is that it’s sampling procedure can be easily parallelized directly in the learned latent space.
sections.heading: 4 LATENT REWARD PREDICTION MODEL FOR PLANNING
sections.text: In order to perform model-based planning directly from observations without a model, we need to learn a dynamics model. For high-dimensional problems, a common model-based DRL approach is to learn a latent representation and then conduct planning in the low-dimensional latent space. In this work, rather than predict the full state/observation, we learn a latent state-space model which predicts only the current and future rewards conditioned on action sequences. We will discuss that observation reconstruction is not necessary if we can predict the rewards well over the planning horizon.
sections.heading: 4.1 LEARNING A LATENT REWARD PREDICTION MODEL
sections.text: Our model consist of three distinct components which can be trained in an end-to-end fashion. To do so, we first embed the state observation st to the latent state zt with the parameterized function φθ : S → Z ⊆ Rdz . In order to propagate the latent state forward in latent space, we define the forward dynamics function which maps fzψ : Z × A → Z , corresponding to the discrete-time dynamical system in latent state space. Finally, we require a reward function Rzζ : Z ×A → R which provides us an estimate of the reward given a latent state zt and control action at. The full latent reward prediction model is depicted in Figure 1.
  zt = φθ(st), ẑt+1 = f z ψ(zt, at), r̂t = R z ζ(zt, at) (4)
  Note that we do not require a “decoding” function since planning algorithms often only need to evaluate the rewards along a trajectory given an action sequence.
  To achieve the aforementioned reward prediction ability, we learn these functions φθ, fzψ, R z ζ over multi-step reward prediction losses. Formally, given a H-step sequence of states and actions {(sk, ak), k = t, . . . , t + H − 1}, the training objective is defined to be the mean-squared error between the original rewards and the multi-step latent reward predictions.
  LH = 1
  H t+H−1∑ τ=t ||γτ−tR(sτ , aτ )− γτ−tRzζ(ẑτ , aτ )||22 (5) where ẑτ = fzψ(. . . f z ψ(φθ(st), at), . . . , aτ−1) (6)
  We later show that this H-step reward prediction loss is sufficient to bound the planning performance over the same horizon.
  Note that we do not impose other losses such as the observation prediction loss or latent state prediction loss as in some previous latent model frameworks (Gelada et al., 2019; Hafner et al., 2019).
  The reason is that other loss functions are not necessary for planning performance. Having those additional losses would increase the training complexity and may induce undesirable local minima.
sections.heading: 4.2 MPC WITH THE LATENT REWARD PREDICTION MODEL
sections.text: Once a latent state-space model is learned, we can perform MPC using the latent embedding. Let πzMPC(st|φθ, fzψ, Rzζ) be the latent MPC policy with the latent reward model (φθ, fzψ, Rzζ). Similar to Equation (3), the latent MPC agent at each time solves a H-step trajectory planning problem.
  max ât:(t+H−1)
  Jz(st, ât:(t+H−1)) (7a)
  subject to Jz(st, ât:(t+H−1)) = t+H−1∑ τ=t γτ−tRzζ(ẑτ , âτ ) (7b)
  ẑt = φθ(st), and ẑτ+1 = fzψ(ẑτ , âτ ) for τ = t, t+ 1, . . . , (t+H − 2) (7c)
  After solving the optimization problem, the latent MPC agent will choose its action by
  πzMPC(st|φθ, fzψ, Rzζ) = â∗t (8) where â∗t:(t+H−1) = argmax
  ât:(t+H−1)
  Jz(st, ât:(t+H−1)) (9)
  Similar to other model-based methods, the latent reward prediction model can be used in an offline setting with a pre-collected dataset. Together with an exploration policy, e.g. -greedy, our model can also be used in an online iterative scheme. See appendix for the descriptions of Algorithm 1 (offline) and Algorithm 2 (online iterative).
sections.heading: 4.3 PLANNING PERFORMANCE GUARANTEE
sections.text: To evaluate the performance of MPC planning, define recursively the optimal n-step Q-function by
  Q∗1(s, a) = R(s, a), Q∗n(s, a) = R(s, a) + γmax a′∈A Q∗n−1(f(s, a), a ′) for n = 2, 3, . . . ,H
  Then the Bellman optimality shows that Q∗n(s, a) gives the optimal value for the MPC planning problem (3), i.e., maxât:(t+H−1) J(st, ât:(t+H−1)) = maxaQ ∗ H(st, a).
  Similarly, define the optimal n-step Q-functions for the latent state-space model (φθ, fzψ, R z ζ):
  Q∗,z1 (z, a) = R z ζ(z, a), Q∗,zn (z, a) = R z ζ(z, a) + γmax
  a′∈A Q∗,zn−1(f z ψ(z, a), a ′)) for n = 2, 3, . . . ,H
  We have the following result connecting theH-step reward prediction loss and theH-step Q-functions.
  Theorem 1. Suppose the H-step prediction losses LH ≤ 2 for any H-step trajectory. Then the H-step Q-functions satisfy
  Q∗,zH (φθ(s), a) ≥ Q ∗ H(s, a)− H1/2
  for all s ∈ S, a ∈ A. Moreover, if πzMPC(s|φθ, fzψ, Rzζ) is optimal for the latent MPC planning problem equation 7, then it is also a (2 H1/2)-optimal policy for the MPC planning problem (3).
  This result provide a planning performance guarantee for the latent reward prediction model. It also confirms that we don’t need other unnecessary losses to achieve the desired performance from the latent model. Note that this result does not require additional assumptions such as Lipschitz continuity of the environment. See the appendix for a proof of the theorem.
sections.heading: 5 EXPERIMENTS
sections.text: In this section we introduce several experiments and their results which are aimed to investigate the representation ability, long-term prediction ability and sample efficiency. We implement the mappings φθ, fzψ, R z ζ by deterministic feed-forward neural networks. The extensions to stochastic and/or recurrent neural networks are potential future directions. In order to do MPC, we choose to use CEM with a planning horizon of H = 12 and K = 1000 trajectory samples. We use the k = 100 best trajectories to compute the control, re-planning after every action. CEM will use the model somewhat like a simulator to collect samples in latent space (in our case this is a batch forward pass in a deep network). We consider two sets of environments with high-dimensional irrelevant observations. Results for an additional image-based environment is also available in the appendix.
sections.heading: 5.1 BASELINES
sections.text: Before introducing the experiments, we briefly introduce the baseline algorithms used for comparison. Throughout our experiments, we have several baselines, each serving to demonstrate the attributes of our reward model: concise representation, long-term prediction and sample efficiency.
  State prediction model To demonstrate the representation ability of our reward prediction model, we consider a latent model where the latent space is learned from state prediction only. We train a reward predictor separately for the model. This model works well for the concise representations, however, like many model-based methods it fails to scale to high-dimensional observations.
  DeepMDP We use the DeepMDP model (Gelada et al., 2019), which is essentially a one-step variant of the reward prediction model with an additional latent state prediction loss. We aim to show that our multi-step reward prediction-only loss improves the planning performance.
  SAC Finally to demonstrate sample efficiency, we compare against the model free Soft Actor-Critic algorithm (SAC) (Haarnoja et al., 2018). SAC is currently a popular DRL algorithm known for its sample efficiency and performance.
sections.heading: 5.2 MULTI-PENDULUM
sections.text: In several RL tasks, there may be irrelevant aspects of the state, e.g. the noisy TV problem (Burda et al., 2018b), so reward prediction might be extremely crucial to learn a more conducive and useful model for planning. To verify this claim, we develop a multi-pendulum environment, an adaptation of the classic pendulum control task (Brockman et al., 2016), where the agent must swing up and stabilize a pendulum at the up-right state with torque constraints. Although this task is simple, it encompasses nonlinear dynamics and long planning horizons in order to accumulate enough energy to swing up. In the multi-pendulum task, the agent operates in an environment with N pendulums, but is only rewarded for swinging up a single pendulum. The other N − 1 pendulums serving as determined observational noise being driven by uniform random actions. The state observation for the N -pendulum environment is given as:
  s = [sin(θ1), cos(θ1), θ̇1, sin(θ2), cos(θ2), θ̇2, . . . , sin(θN ), cos(θN ), θ̇N ] ∈ R3N
  Since reward prediction model can be used either completely offline or iterative online setting, we present two sets of experiments. First we show the average final performances of Algorithm 1 and baselines after training completely offline from 20000 steps of the environment under time-correlated noised control. We vary the number of pendulums N ∈ {1, 2, 5, 10, 15} where we maintain the latent space to be 3-dimensional throughout. In the second experiment we train the model and baselines under the iterative scheme described in Algorithm 2, consuming approximately the same number of samples. We also compare against the SAC algorithm as a state-of-the-art model-free baseline.
  As shown in Figure 2, the reward prediction loss does not degrade quickly for our reward prediction model as the number of irrelevant pendulums increases, but degrades almost immediately for the state prediction model. DeepMDP struggles to predict well in all cases, due to only training on a single-step latent plus reward reconstruction loss. In Table 1, we can further see that performance of the reward prediction model (Algorithm 1) is not very sensitive to the number of excess pendulum
  environments, where the state prediction model immediately fails to solve the task. Note that the SAC method requires a significantly greater number of samples to achieve similar performance.
  For the iterative scheme (Algorithm 2), the reward prediction model clearly outperforms all baselines as can be seen in Figure 3. Note that the DeepMDP model has a much higher variance than other methods. This high variance could be due the local optima from the latent state prediction loss in DeepMDP. On the other hand, our reward prediction model provides a stable low-variance result.
sections.heading: 5.3 MULTI-CHEETAH
sections.text: In order to demonstrate that our method scales to higher dimensional environment, we use the Deepmind control suite “cheetah-run” environment (Tassa et al., 2018) to construct “multi-cheetah”. The cheetah-run environment has a continuous 17-dimensional observation space and 6-dimensional action space. Similar to multi-pendulum, we concatenate each cheetah observation vector together into one observation, where only the first cheetah environment return the reward signal. Because a random policy is not sufficient to explore the state space in cheetah, the remaining N − 1 cheetahs act according to an expert SAC policy. In multi-cheetah, we mainly compare our method with the model-free baseline SAC to highlight the strong performance and sample efficiency. Shown in figure 4, the reward prediction model is able to learn and plan near-optimally using far less samples than SAC in the 1-cheetah environment (i.e., the regular cheetah-run). Within 500 episodes, the reward model is able to achieve a near-optimal running behavior (503 and 501 final return in evaluation for 1-cheetah and 5-cheetah respectively), which takes SAC 3000 episodes to achieve. After 1000 episodes of 1-cheetah (not shown in the figure), the reward model achieves 570 return, approaching SAC’s maximum performance of 590 after 4000 episodes. When there are several irrelevant running cheetahs, the performance SAC further drops. In comparison, our reward prediction model is rather agnostic to the irrelevant observations and maintains high performance and sample efficiency.
sections.heading: 6 CONCLUSION
sections.text: In this paper we introduced a method for learning a latent state representation from rewards. By constructing a latent space based on multi-step reward predictions, we obtained a concise representation similar to that of model-free DRL algorithms while maintaining high sample efficiency by planning with model-base algorithms. In this work we have used a sample-based CEM algorithm for MPC planning, but the use of other planning algorithms is also plausible. We demonstrated in the multi-pendulum and multi-cheetah environments that the latent reward prediction model is able to succeed in the presence of high-dimensional irrelevant information with both offline and iterative online schemes.
  Although, there are several aspects which still need to be investigated such as more complex environments and sparse reward settings, where typically a finite-horizon planner would struggle. We would also like to investigate how to incorporate model uncertainty through either stochastic dynamics or ensemble methods. One interesting question we have is what is the optimal observer or what the optimal combination of reward and state reconstruction to do well in the task. We believe our minimal formulation is a good starting point for this question.
sections.heading: A ALGORITHM DESCRIPTIONS
sections.text: Algorithm 1 Offline Training Inputs: planning horizon H , dataset D, latent model (φθ, fzψ, R z ζ).
  Train (φθ, fzψ, R z ζ) using the loss function LH in equation 5 and data D. Output: policy πzMPC(st|φθ, fzψ, Rzζ)
  Algorithm 2 Online Iterative Scheme Inputs: planning horizon H , number of iterations N , latent model (φθ, fzψ, R z ζ).
  Initialize dataset D with random trajectories. for i = 1 to N do
  Train (φθ, fzψ, R z ζ) using the loss function LH in equation 5 and data D. Run latent MPC πzMPC(st|φθ, fzψ, Rzζ) in the environment with -greedy and collect data Di. D ← D ∪Di
  end for Output: policy πzMPC(st|φθ, fzψ, Rzζ)
sections.heading: B PROOF OF THEOREM 1
sections.text: Proof. Note that the H-step Q-function can be written as
  Q∗H(s, a) = H−1∑ t=0 γtR(st, a ∗ t )
  where s0 = s, a∗0 = a, and for t > 0
  st = f(st−1, a ∗ t−1) = ft(s, a ∗ 0:t−1) a∗t = π ∗(st)
  Here π∗ denotes the optimal policy for the MPC problem equation 3. Because Q∗,zH (φθ(s), a) is the optimal Q-function for the latent MPC problem equation 7, we have
  Q∗,zH (φθ(s), a) ≥ H−1∑ t=0 γtRzζ(f z ψ,t(φθ(s), a ∗ 0:t−1), a ∗ t )
  Adding and subtracting Q∗H(s, a) we obtain
  Q∗,zH (φθ(s), a) ≥Q ∗ H(s, a) + H−1∑ t=0 γtRzζ(f z ψ,t(φθ(s), a ∗ 0:t−1), a ∗ t )−Q∗H(s, a)
  =Q∗H(s, a) + H−1∑ t=0 γt(Rzζ(f z ψ,t(φθ(s), a ∗ 0:t−1), a ∗ t )−R(st, a∗t ))
  ≥Q∗H(s, a)− H−1∑ t=0 ||γt(Rzζ(fzψ,t(φθ(s), a∗0:t−1), a∗t )−R(st, a∗t ))||2 =Q∗H(s, a)− (HLH)1/2 ≥ Q∗H(s, a)−H1/2ε (10)
  Now let π∗,z(z) be the optimal policy for the latent MPC problem in the latent space, then from the Q-function definition we have
  Qz,∗H (φθ(s), a) = H−1∑ t=0 γtRzζ(st, a ∗,z t )
  where z0 = φθ(s), a ∗,z 0 = a, and for t > 0
  zt = f z ψ(st−1, a ∗,z t−1) = f z ψ,t(φθ(s), a ∗,z 0:k−1) a∗,zt = π ∗,z(zt)
  The H-step Q-function of policy π∗,z ◦ φθ in problem equation 3 is then given by
  Qπ ∗,z◦φθ H (s, a) = H−1∑ t=0 γtR(ft(s, a ∗,z 0:t−1), a ∗,z k )
  =Q∗,zH (φ(s), a) + H−1∑ t=0 γtR(ft(s, a ∗,z 0:t−1), a ∗,z k )−Q ∗,z H (φ(s), a)
  =Q∗,zH (φ(s), a) + H−1∑ t=0 γt(R(ft(s, a ∗,z 0:t−1), a ∗,z k )−R z ζ(f z ψ,t(φθ(s), a ∗,z 0:t−1), a ∗,z t )) ≥Q∗,zH (φ(s), a)−H 1/2ε (11)
  Now combining equation 10 and equation 11 we get
  Qπ ∗,z◦φθ H (s, a) ≥ Q ∗ H(s, a)− 2H1/2ε
  which means that π∗,z ◦ φ is a 2H1/2 -optimal policy.
sections.heading: C EXPERIMENTS WITH IMAGE OBSERVATIONS
sections.text: To demonstrate that this method can solve tasks with rather high-dimensional state observations, we solve the single Pendulum task from two-gray-scale images (2× 100× 100 pixels). Similar to the Mutli-Pendulum experiments, we let the latent space be 3-dimensional. In provide two images per state to insure that velocity can be inferred and fully-observable assumption in maintained. We find that the prediction results are fairly similar to the kinematic state representation, but more noisy due to the course image observations. We train the model using only mostly offline, using only 1 iteration. We initialize the model with 10000 samples under random control trained for 100 epochs and then trained once more with another 10000 samples using the new control policy. This policy is able to solve the task efficiently (142 return) similarly to the kinematic state representation, producing similar prediction results. Its interesting to note that the open-loop model predicts that the pendulum will continue to oscillate and not stabilize. This could be evidence that the model over fits to the random training data, however it would also require very accurate model to predict stabilizing behavior far into the future.