summaries: This paper takes a different approach for tackling the hierarchical RL problem. Their approach is to decompose the policy into a bunch of primitives. Each primitive acts according to its own interpretation of the state. All the primitives are competing with each other on a given state to take an action. It turns out that these primitive policies can be transferred to other tasks as they represent subtasks of a bigger task. The paper performs extensive experiments to show that this scheme improves over both flat and hierarchical policies in terms of generalization.

summaries: The paper draws upon the idea of information bottleneck to do task decomposition so as to learn policy primitives similar to hierarchical reinforcement learning that combine together in a competitive manner to specialize in different parts of the task's domain. These policy primitives don't need a higher-level meta-policy to stitch them together. Instead the decision is made in a decentralized manner balancing the cost of information acquisition with maximizing rewards.

summaries: This paper is about a policy design, where the policy is expressed as a mixture of policies called primitives.  Each primitive is made of an encoder and a decoder, mapping state to actions, rather than temporally extended actions (or options in RL).  The primitives compete with each other to be selected in each state and thus do away with the need for a meta-policy to select the primitives.  The selected primitive in each state trades between reward maximization and information content.

abstractText: Reinforcement learning agents that operate in diverse and complex environments can benefit from the structured decomposition of their behavior. Often, this is addressed in the context of hierarchical reinforcement learning, where the aim is to decompose a policy into lower-level primitives or options, and a higher-level meta-policy that triggers the appropriate behaviors for a given situation. However, the meta-policy must still produce appropriate decisions in all the states. In this work, we propose a policy design that decomposes into primitives, similarly to hierarchical reinforcement learning, but without an explicit high-level meta-policy. Instead, each primitive can decide for themselves whether they wish to act in the current state. We use an information-theoretic mechanism for enabling this decentralized decision: each primitive chooses how much information it needs about the current state to make a decision and the primitive that requests the most information about the current state acts in the environment. Regularizing the primitives to use as little information as possible leads to natural competition and specialization. We experimentally demonstrate that this policy architecture improves over both flat and hierarchical policies in terms of generalization.

authors:
Alessandro Achille
Stefano Soatto

title: Information dropout: learning optimal representations through noise

authors:
Alexander A. Alemi
Ian Fischer
Joshua V. Dillon
Kevin Murphy

title: Deep variational information

venue: bottleneck. CoRR,

authors:
Jacob Andreas
Dan Klein
Sergey Levine

title: Modular multitask reinforcement learning with policy sketches

venue: In Proceedings of the 34th International Conference on Machine Learning-Volume

authors:
Pierre-Luc Bacon
Jean Harb
Doina Precup

title: The option-critic architecture

venue: In AAAI,

authors:
Maxime Chevalier-Boisvert
Lucas Willems
Suman Pal

title: Minimalistic gridworld environment for openai gym

authors:
Kyunghyun Cho
Bart Van Merriënboer
Caglar Gulcehre
Dzmitry Bahdanau
Fethi Bougares
Holger Schwenk
Yoshua Bengio

title: Learning phrase representations using rnn encoder-decoder for statistical machine translation

venue: arXiv preprint arXiv:1406.1078,

authors:
Christian Daniel
Gerhard Neumann
Jan Peters

title: Hierarchical relative entropy policy search

venue: In Artificial Intelligence and Statistics,

authors:
Peter Dayan
Geoffrey E Hinton

title: Feudal reinforcement learning

title: Hierarchical reinforcement learning with the maxq value function decomposition

venue: Journal of Artificial Intelligence Research,

authors:
Benjamin Eysenbach
Abhishek Gupta
Julian Ibarz
Sergey Levine

title: Diversity is all you need: Learning skills without a reward function

venue: arXiv preprint arXiv:1802.06070,

authors:
Carlos Florensa
Yan Duan
Pieter Abbeel

title: Stochastic neural networks for hierarchical reinforcement learning

venue: arXiv preprint arXiv:1704.03012,

authors:
K. Frans
J. Ho
X. Chen
P. Abbeel
J. Schulman

title: Meta Learning Shared Hierarchies

authors:
Kevin Frans
Jonathan Ho
Xi Chen
Pieter Abbeel
John Schulman

title: Meta learning shared hierarchies

venue: arXiv preprint arXiv:1710.09767,

authors:
Tuomas Haarnoja
Kristian Hartikainen
Pieter Abbeel
Sergey Levine

title: Latent space policies for hierarchical reinforcement learning

venue: arXiv preprint arXiv:1804.02808,

authors:
Nicolas Heess
Srinivasan Sriram
Jay Lemmon
Josh Merel
Greg Wayne
Yuval Tassa
Tom Erez
Ziyu Wang
Ali Eslami
Martin Riedmiller

title: Emergence of locomotion behaviours in rich environments

venue: arXiv preprint arXiv:1707.02286,

authors:
Justin Johnson
Bharath Hariharan
Laurens van der Maaten
Judy Hoffman
Li Fei-Fei
C Lawrence Zitnick
Ross Girshick

title: Inferring and executing programs for visual reasoning

venue: In Proceedings of the IEEE International Conference on Computer Vision,

authors:
Tejas D Kulkarni
Karthik Narasimhan
Ardavan Saeedi
Josh Tenenbaum

title: Hierarchical deep reinforcement learning: Integrating temporal abstraction and intrinsic motivation

authors:
Libin Liu
Jessica Hodgins

title: Learning to schedule control fragments for physics-based characters using deep q-learning

venue: ACM Transactions on Graphics,

authors:
Marlos C Machado
Marc G Bellemare
Michael Bowling

title: A laplacian framework for option discovery in reinforcement learning

venue: arXiv preprint arXiv:1703.00956,

authors:
Josh Merel
Arun Ahuja
Vu Pham
Saran Tunyasuvunakool
Siqi Liu
Dhruva Tirumala
Nicolas Heess
Greg Wayne

title: Hierarchical visuomotor control of humanoids

authors:
Josh Merel
Leonard Hasenclever
Alexandre Galashov
Arun Ahuja
Vu Pham
Greg Wayne
Yee Whye Teh
Nicolas Heess

title: Neural probabilistic motor primitives for humanoid control

authors:
Volodymyr Mnih
Koray Kavukcuoglu
David Silver
Andrei A Rusu
Joel Veness
Marc G Bellemare
Alex Graves
Martin Riedmiller
Andreas K Fidjeland
Georg Ostrovski

title: Human-level control through deep reinforcement learning

authors:
Volodymyr Mnih
Adria Puigdomenech Badia
Mehdi Mirza
Alex Graves
Timothy Lillicrap
Tim Harley
David Silver
Koray Kavukcuoglu

title: Asynchronous methods for deep reinforcement learning

venue: In International conference on machine learning,

authors:
Giambattista Parascandolo
Niki Kilbertus
Mateo Rojas-Carulla
Bernhard Schölkopf

title: Learning independent causal mechanisms

venue: arXiv preprint arXiv:1712.00961,

authors:
Adam Paszke
Sam Gross
Soumith Chintala
Gregory Chanan
Edward Yang
Zachary DeVito
Zeming Lin
Alban Desmaison
Luca Antiga
Adam Lerer

title: Automatic differentiation in PyTorch

venue: In NIPS Autodiff Workshop,

authors:
Xue Bin Peng
Glen Berseth
Kangkang Yin
Michiel Van De Panne

title: Deeploco: Dynamic locomotion skills using hierarchical deep reinforcement learning

venue: ACM Trans. Graph.,

authors:
Xue Bin Peng
Pieter Abbeel
Sergey Levine
Michiel van de Panne

title: Deepmimic: Exampleguided deep reinforcement learning of physics-based character skills

authors:
Clemens Rosenbaum
Ignacio Cases
Matthew Riemer
Tim Klinger

title: Routing networks and the challenges of modular and compositional computation

authors:
Alexander Sasha Vezhnevets
Simon Osindero
Tom Schaul
Nicolas Heess
Max Jaderberg
David Silver
Koray Kavukcuoglu

title: Feudal networks for hierarchical reinforcement learning

venue: arXiv preprint arXiv:1703.01161,

authors:
John Schulman
Philipp Moritz
Sergey Levine
Michael Jordan
Pieter Abbeel

title: High-dimensional continuous control using generalized advantage estimation

venue: arXiv preprint arXiv:1506.02438,

authors:
John Schulman
Filip Wolski
Prafulla Dhariwal
Alec Radford
Oleg Klimov

title: Proximal policy optimization algorithms

venue: arXiv preprint arXiv:1707.06347,

authors:
Richard S Sutton
Andrew G Barto

title: Reinforcement learning: An introduction

venue: MIT press,

authors:
Richard S. Sutton
David McAllester
Satinder Singh
Yishay Mansour

title: Policy gradient methods for reinforcement learning with function approximation

venue: In Proceedings of the 12th International Conference on Neural Information Processing Systems,

authors:
Richard S Sutton
Doina Precup
Satinder Singh

title: Between mdps and semi-mdps: A framework for temporal abstraction in reinforcement learning

venue: Artificial intelligence,

authors:
Tijmen Tieleman
Geoffrey Hinton

title: Lecture 6.5-rmsprop, coursera: Neural networks for machine learning

venue: University of Toronto, Technical Report,

authors:
Naftali Tishby
Fernando C.N. Pereira
William Bialek

title: The information bottleneck method

venue: CoRR, physics/0004057,

authors:
Emanuel Todorov
Tom Erez
Yuval Tassa

title: Mujoco: A physics engine for model-based control

venue: In 2012 IEEE/RSJ International Conference on Intelligent Robots and Systems,

authors:
Laurens van der Maaten
Geoffrey Hinton

title: Visualizing data using t-SNE

venue: Journal of Machine Learning Research,

title: Simple Statistical Gradient-Following Algorithms for Connectionist Reinforcement Learning

venue: Machine Learning,

authors:
Bohan Wu
Jayesh K Gupta
Mykel J Kochenderfer

title: Model primitive hierarchical lifelong reinforcement learning

venue: In Proceedings of the 18th International Conference on Autonomous Agents and MultiAgent Systems, pp. 34–42. International Foundation for Autonomous Agents and Multiagent Systems,

authors:
Yuhuai Wu
Elman Mansimov
Roger B Grosse
Shun Liao
Jimmy Ba

title: Scalable trust-region method for deep reinforcement learning using kronecker-factored approximation

heading: 1 INTRODUCTION

text:
Learning policies that generalize to new environments or tasks is a fundamental challenge in reinforcement learning. While deep reinforcement learning has enabled training powerful policies, which outperform humans on specific, well-defined tasks (Mnih et al., 2015), their performance often diminishes when the properties of the environment or the task change to regimes not encountered during training.
This is in stark contrast to how humans learn, plan, and act: humans can seamlessly switch between different aspects of a task, transfer knowledge to new tasks from remotely related but essentially distinct prior experience, and combine primitives (or skills) used for distinct aspects of different tasks in meaningful ways to solve new problems. A hypothesis hinting at the reasons for this discrepancy is that the world is inherently compositional, such that its features can be described by compositions of small sets of primitive mechanisms (Parascandolo et al., 2017). Since humans seem to benefit from learning skills and learning to combine skills, it might be a useful inductive bias for the learning models as well.
This is addressed to some extent by the hierarchical reinforcement learning (HRL) methods, which focus on learning representations at multiple spatial and temporal scales, thus enabling better exploration strategies and improved generalization performance (Dayan & Hinton, 1993; Sutton et al., 1999b; Dietterich, 2000; Kulkarni et al., 2016). However, hierarchical approaches rely on some form of learned high-level controller, which decides when to activate different components in the hierarchy. While low-level sub-policies can specialize to smaller portions of the state space, the top-level controller (or master policy) needs to know how to deal with any given state. That is, it should provide optimal behavior for the entire accessible state space. As the master policy is trained
1 Mila, University of Montreal; 2 Facebook AI Research; work done while the author was at Mila, University of Montreal; 3 University of California, Berkeley. anirudhgoyal9119@gmail.com
on a particular state distribution, learning it in a way that generalizes to new environments effectively becomes the bottleneck for such approaches (Sasha Vezhnevets et al., 2017; Andreas et al., 2017).

heading: Action selection

heading: Competition mechanism

text:
We argue, and empirically show, that in order to achieve better generalization, the interaction between the low-level primitives and the selection thereof should itself be performed without requiring a single centralized network that understands the entire state space. We, therefore, propose a decentralized approach as an alternative to standard HRL, where we only learn a set of low-level primitives without learning an explicit high-level controller. In particular, we construct a factorized representation of the policy by learning simple primitive policies, which focus on distinct regions of the state space. Rather than being gated by a single meta-policy, the primitives directly compete with one another to determine which one should be active at any given time, based on the degree to which their state encoders “recognize” the current state input. While, technically, the competition between primitives implicitly realizes a global selection mechanism, we consider our model decentralized in the sense that individual primitives can function on their own, and can be combined in new ways, without relying on an explicit high-level controller.
We frame the problem as one of information transfer between the current state and a dynamically selected primitive policy. Each policy can, by itself, decide to request information about the current state, and the amount of information requested is used to determine which primitive acts in the current state. Since the amount of state information that a single primitive can access is limited, each primitive is encouraged to use its resources wisely. Constraining the amount of accessible information in this way naturally leads to a decentralized competition and decision mechanism where individual primitives specialize in smaller regions of the state space. We formalize this information-driven objective based on the variational information bottleneck. The resulting set of competing primitives achieves both a meaningful factorization of the policy and an effective decision mechanism for which primitives to use. Importantly, not relying on a centralized meta-policy enables the individual primitive mechanisms can be recombined in a plug-and-play fashion, and the primitives can be transferred seamlessly to new environments.
Contributions: In summary, the contributions of our work are as follows: (1) We propose a method for learning and operating a set of functional primitives in a decentralized way, without requiring an explicit high-level meta-controller to select the active primitives (see Fig. 1 for illustration). (2) We introduce an information-theoretic objective, the effects of which are twofold: a) it leads to the specialization of individual primitives to distinct regions of the state space, and b) it enables a competition mechanism, which is used to select active primitives in a decentralized manner. (3) We demonstrate the superior transfer learning performance of our model, which is due to the flexibility of the proposed framework regarding the dynamic addition, removal, and recombination of primitives. Decentralized primitives can be successfully transferred to larger or previously unseen environments, and outperform models with an explicit meta-controller for primitive selection.

heading: 2 PRELIMINARIES

text:
We consider a Markov decision process (MDP) defined by the tuple (S,A, P, r, γ), where the state space S and the action space A may be discrete or continuous. The environment emits a bounded reward r : S × A → [rmin, rmax] on each transition and γ ∈ [0, 1) is the discount factor. π(.|s) denotes a policy over the actions given the current state s. R(π) = Eπ[ ∑ t γ
tr(st)] denotes the expected total return when an agent follows the policy π. The standard objective in reinforcement learning is to maximize the expected total return R(π). We use the concept of the information bottleneck (Tishby et al., 2000) to learn compressed representations. The information bottleneck objective is formalized as minimizing the mutual information of a bottleneck representation layer with the input while maximizing its mutual information with the corresponding output. This type of input compression has been shown to improve generalization (Achille & Soatto, 2016; Alemi et al., 2016).

heading: 3 INFORMATION-THEORETIC LEARNING OF DISTINCT PRIMITIVES

text:
Our goal is to learn a policy, composed of multiple primitive sub-policies, to maximize the expected reward over T -step interactions for a distribution of tasks. Simple primitives which focus on solving a part of the given task (and not the complete task) should generalize more effectively, as they can be applied to similar aspects of different tasks (subtasks) even if the overall objective of the tasks are drastically different. Learning primitives in this way can also be viewed as learning a factorized representation of a policy, which is composed of several independent policies.
Our proposed approach consists of three mechanisms: 1) a mechanism for restricting a particular primitive to a subset of the state space; 2) a competition mechanism between primitives to select the most effective primitive for a given state; 3) a regularization mechanism to improve the generalization performance of the policy as a whole. We consider experiments with both fixed and variable sets of primitives and show that our method allows for primitives to be added or removed during training, or recombined in new ways. Each primitive is represented by a differentiable, parameterized function approximator, such as a neural network.

heading: 3.1 PRIMITIVES WITH AN INFORMATION BOTTLENECK

text:
To encourage each primitive to encode information from a particular part of state space, we limit the amount of information each primitive can access from the state. In particular, each primitive has an information bottleneck with respect to the input state, preventing it from using all the information from the state.
We define the overall policy as a mixture of primitives, π(a | s) = ∑ k ckπ k(a | s) ,
where πk(a | s) denotes the kth primitive and ck = δkk′ for k′ ∼ p(k′ | s). We denote the probability of selecting the kth primitive as αk(s) := p(k | s). Rather than learning an explicit model for p(k | s), however, we impose an information-based mechanism for selecting primitives, wherein we limit the amount of information each primitive can contain and select the ones that request the most information about the state. To implement an information bottleneck, we design each of the K primitives to be composed of an encoder penc(zk | s) and a decoder pdec(a | zk), together forming the primitive policy,
πkθ (a | s) = ∫ z penc(zk | s) pdec(a | zk) dzk .
The encoder output zk is meant to represent the information about the current state s that an individual primitive k believes is important to access in order to perform well. The decoder takes this encoded information and produces a distribution over the actions a. Following the variational information bottleneck objective (Alemi et al., 2016), we penalize the KL divergence of penc(zk|s) and a prior p(z),
Lk = DKL (penc(zk | s)||p(z)) . (1)
In practice, we estimate the marginalization over z using a single sample throughout our experiments.
In other words, a primitive pays an “information cost” proportional toLk for accessing the information about the current state.
In the experiments below, we fix the prior to be a unit Gaussian. In the general case, we can learn the prior as well and include its parameters in θ. The information bottleneck encourages each primitive to limit its knowledge about the current state, but it will not prevent multiple primitives from specializing to similar parts of the state space. To mitigate this redundancy, and to make individual primitives focus on different regions of the state space, we introduce an information-based competition mechanism to encourage diversity among the primitives.

heading: 3.2 COMPETING INFORMATION-CONSTRAINED PRIMITIVES

text:
We can use the information measure from equation 1 to define a selection mechanism for the primitives without having to learn a centralized meta-policy. The intuition is that the information content of an individual primitive encodes its effectiveness in a given state s such that the primitive with the highest value Lk should be activated in that particular state. In particular, we set αk = Z−1 exp(βLk) to obtain a distribution over k as a function of the information content, activating the primitives with the highest information content. Here, Z =∑ k exp(βLk) is a normalization constant. This mechanism enables competition between primitives, leading them to focus on parts of the state space that they “understand” well and letting others act in other parts.
Trading reward and information. To perform proper credit assignment, the environment reward is distributed to primitives according to their participation in the global decision, i.e. the reward rk given to the kth primitive is weighted by its selection coefficient, such that rk = αkr, with r = ∑ k rk. Hence, a primitive can potentially get a higher reward when deciding to act, but it also pays a higher price for accessing more information about the current state. The information bottleneck and the competition mechanism, when combined with the overall reward maximization objective, will lead to specialization of individual primitives to distinct regions in the state space. That is, each primitive should specialize in a part of the state space that it can reliably associate rewards with. Since the entire ensemble still needs to understand all of the state space for the given task, different primitives need to encode and focus on different parts of the state space.

heading: 3.3 REGULARIZING PRIMITIVE SELECTION

text:
The objective described above will optimize the expected return while minimizing the information content of individual primitives. This is not sufficient, however, as it might lead to highly unbalanced outcomes: some primitives might be more active initially and learn to become even more active, completely disabling other primitives.
Thus, in addition to minimizing each primitive’s absolute information content, we need to normalize their activity with respect to each other. To do so, we penalize their information content in proportion to their activation by adding a regularization term of the form
Lreg = ∑ k αkLk . (2)
Note that this can be rewritten (see Appendix A) as Lreg = −H(α) + LSE(L1, . . . ,LK) , where H(α) is the entropy of α, and LSE is the LogSumExp function, LSE(x) = log( ∑ j e xj ). Thus, minimizing Lreg increases the entropy of α, leading to a diverse set of primitive selections, in turn, ensuring that different combinations of the primitives are used. Similarly, LSE approximates the maximum of its arguments, LSE(x) ≈ maxj xj , and, therefore, penalizes the dominating Lk terms, thus equalizing their magnitudes.

heading: 3.4 OBJECTIVE AND ALGORITHM SUMMARY

text:
Our overall objective function consists of 3 terms,
1. The expected return from the standard RL objective, R(π) which is distributed to the primitives according to their participation,
2. The individual bottleneck terms leading the individual primitives to focus on specific parts of the state space, Lk for k = 1, . . . ,K,
3. The regularization term applied to the combined model, Lreg.
The overall objective for the kth primitive thus takes the form:
Jk(θ) ≡ Eπθ [rk]− βindLk − βregLreg , (3)
where Eπθ denotes an expectation over the state trajectories generated by the agent’s policy, rk = αkr is the reward given to the kth primitive, and βind, βreg are the parameters controlling the impact of the respective terms.
Implementation: In our experiments, the encoders penc(zk | s) and decoders pdec(a | zk) (see. Fig. 1) are represented by neural networks, the parameters of which we denote by θ. Actions are sampled through each primitive every step. While our approach is compatible with any RL method, we maximize J(θ) computed on-policy from the sampled trajectories using a score function estimator (Williams, 1992; Sutton et al., 1999a) specifically A2C (Mnih et al., 2016) (unless otherwise noted). Every experimental result reported has been averaged over 5 random seeds. Our model introduces 2 extra hyper-parameters βind, βreg.

heading: 4 RELATED WORK

text:
There are a wide variety of hierarchical reinforcement learning approaches(Sutton et al., 1998; Dayan & Hinton, 1993; Dietterich, 2000). One of the most widely applied HRL framework is the Options framework ((Sutton et al., 1999b)). An option can be thought of as an action that extends over multiple timesteps, thus providing the notion of temporal abstraction or subroutines in an MDP. Each option has its own policy (which is followed if the option is selected) and the termination condition (to stop the execution of that option). Many strategies are proposed for discovering options using task-specific hierarchies, such as pre-defined sub-goals (Heess et al., 2017), hand-designed features (Florensa et al., 2017), or diversity-promoting priors (Daniel et al., 2012; Eysenbach et al., 2018). These approaches do not generalize well to new tasks. Bacon et al. (2017) proposed an approach to learn options in an end-to-end manner by parameterizing the intra-option policy as well as the policy and termination condition for all the options. Eigen-options (Machado et al., 2017) use the eigenvalues of the Laplacian (for the transition graph induced by the MDP) to derive an intrinsic reward for discovering options as well as learning an intra-option policy.
In this work, we consider a sparse reward setup with high dimensional action spaces. In such a scenario, performing unsupervised pretraining or using auxiliary rewards leads to much better performance (Frans et al., 2017; Florensa et al., 2017; Heess et al., 2017). Auxiliary tasks such as motion imitation have been applied to learn motor primitives that are capable of performing a variety of sophisticated skills (Liu & Hodgins, 2017; Peng et al., 2017; Merel et al., 2019b;a). Our work is also related to the Neural Module Network family of architectures (Andreas et al., 2017; Johnson et al., 2017; Rosenbaum et al., 2019) where the idea is to learn modules that can perform some useful computation like solving a subtask and a controller that can learn to combine these modules for solving novel tasks. More recently, Wu et al. (2019) proposed a framework for using diverse suboptimal world models to learn primitive policies. The key difference between our approach and all the works mentioned above is that we learn functional primitives without requiring any explicit high-level meta-controller or master policy.

heading: 5 EXPERIMENTAL RESULTS

text:
In this section, we briefly outline the tasks that we used to evaluate our proposed method and direct the reader to the appendix for the complete details of each task along with the hyperparameters used for the model. We designed experiments to address the following questions: a) Learning primitives – Can an ensemble of primitives be learned over a distribution of tasks? b) Transfer Learning using primitives – Can the learned primitives be transferred to unseen/unsolvable sparse environments? c) Comparison to centralized methods – How does our method compare to approaches where the primitives are trained using an explicit meta-controller, in a centralized way?
Baselines. We compare our proposed method to the following baselines: a) Option Critic (Bacon et al., 2017) – We extended the author’s implementation of the Option Critic architecture and experimented with multiple variations in terms of hyperparameters and state/goal encoding. None of these yielded reasonable performance in partially observed tasks, so we omit it from the results. b) MLSH (Meta-Learning Shared Hierarchy) (Frans et al., 2017) – This method uses meta-learning to learn sub-policies that are shared across tasks along with learning a task-specific high-level master. It also requires a phase-wise training schedule between the master and the sub-policies to stabilize training. We use the MLSH implementation provided by the authors. c) Transfer A2C: In this method, we first learn a single policy on the one task and then transfer the policy to another task, followed by fine-tuning in the second task.

heading: 5.1 LEARNING ENSEMBLES OF FUNCTIONAL PRIMITIVES

text:
We evaluate our approach on a number of RL environments to demonstrate that we can indeed learn sets of primitive policies focusing on different aspects of a task and collectively solving it.
Four Room Maze: We consider the Four-rooms gridworld environment (Sutton et al., 1999c) where the agent has to navigate its way through a grid of four interconnected rooms to reach a goal position within the grid. We consider the scenario where the starting position of the agent is fixed, but the goal is sampled from a discrete set. Fig. 3 shows that the proposed algorithm can learn four primitives. Refer to Appendix F for more details.
Motion Imitation. To evaluate the proposed method in terms of scalability, we present a series of tasks from the motion imitation domain, showing that we can use a set of distinct primitives for imitation learning. In these tasks, we train a simulated 2D biped character to perform a variety of highly dynamic skills by imitating motion capture clips recorded from human actors. Each mocap
https://github.com/jeanharb/option_critic https://github.com/openai/mlsh
clip is represented by a target state trajectory τ∗ = {s∗0, s∗1, ..., s∗T }, where s∗t denotes the target state at timestep t. The input to the policy is augmented with a goal gt = {s∗t+1, s∗t+2}, which specifies the the target states for the next two timesteps. Both the state st and goal gt are then processed by the encoder penc(zt|st, gt). The repertoire of skills consists of 8 clips depicting different types of walks, runs, jumps, and flips. The motion imitation approach closely follows Peng et al. (2018). To analyze the specialization of the various primitives, we computed 2D embeddings of states and goals which each primitive is active in, and the actions proposed by the primitives. Fig. 4 illustrates the embeddings computed with t-SNE (van der Maaten & Hinton, 2008). The embeddings show distinct clusters for the primitives, suggesting a degree of specialization of each primitive to certain states, goals, and actions.
5.2 MULTI-TASK TRAINING
We evaluate our model in a partiallyobservable 2D multi-task environment called Minigrid, similar to the one introduced in (Chevalier-Boisvert et al., 2018). The environment is a two-dimensional grid with a single agent, impassable walls, and many objects scattered in the environment. The agent is provided with a natural language string that specifies the task that the agent needs to complete. The setup is partially observable, and the agent only gets the small, egocentric view of the grid (along with the natural language task description). We consider three tasks here: the Pickup task (A), where the agent is re-
quired to pick up an object specified by the goal string, the Unlock task (B) where the agent needs to unlock the door (there could be multiple keys in the environment, and the agent needs to use the key which matches the color of the door) and the UnlockPickup task (C), where the agent first needs to unlock a door that leads to another room. In this room, the agent needs to find and pick up the object specified by the goal string. Additional implementation details of the environment are provided in appendix D. Details on the agent model can be found in appendix D.3.
We train agents with varying numbers of primitives on various tasks – concurrently, as well as in transfer settings. The different experiments are summarized in Figs. 5 and 7. An advantage of the multi-task setting is that it allows for quantitative interpretability as to when and which primitives are being used. The results indicate that a system composed of multiple primitives generalizes more easily to a new task, as compared to a single policy. We further demonstrate that several primitives can be combined dynamically and that the individual primitives respond to stimuli from new environments when trained on related environments.

heading: 5.3 DO LEARNED PRIMITIVES HELP IN TRANSFER LEARNING?

text:
We evaluate our approach in the settings where the adaptation to the changes in the task is vital. The argument in favor of modularity is that it enables better knowledge transfer between related tasks. Naturally, the transfer is easier when the tasks are closely related, as the model will only need to learn how to compose the already-learned primitives. In general, it is difficult to determine how closely related two tasks are, however, and the inductive bias of modularity could even be harmful if the two tasks are very different. In such cases, we could add new primitives (which would need to be learned) and still obtain a sample-efficient transfer, as some part of the task structure would already have been captured by the pretrained primitives. This approach can be extended towards adding primitives during training, providing a seamless way to combine knowledge about different tasks to solve more complex tasks. We investigate here the transfer properties of a primitive trained in one environment and transferred to a different one. Results are shown in Fig. 5.
Continuous control for ant maze We evaluate the transfer performance of pretrained primitives on the cross maze environment (Haarnoja et al., 2018). Here, a quadrupedal ant robot must walk to the different goals along the different paths (see Appendix G for details). The goal is randomly
chosen from a set of available goals at the start of each episode. We pretrain a policy (see model details in Appendix G.1) with a motion reward in an environment which does not have any walls (similar to Haarnoja et al. (2018)), and then transfer the policy to the second task where the ant has to navigate to a random goal chosen from one of the 3 (or 10) available goal options. For our model, we make four copies of the pretrained policies and then finetune the model using the pretrained policies as primitives. We compare to both MLSH (Frans et al., 2017) and option-critic (Bacon et al., 2017). All these baselines have been pretrained in the same manner. As evident from Fig. 7, our method outperforms the other approaches. The fact that the initial policies successfully adapt to the transfer environment underlines the flexibility of our approach.
Zero Shot Generalization: The purpose of this experiment is to show that the model consisting of multiple primitives is somewhat able to decompose the task C into its subtasks, A and B. The better this decomposition is the better should the model transfer to the individual subtasks. In order to test this, we trained a set of 4 primitives on task C, and then evaluate them (without finetuning) on tasks A and B. We note that the ensemble is able to solve the transfer tasks, A and B, successfully 72% of the time, while a monolithic policy’s success rate is 38%. This further shows that the primitives learn meaningful decompositions.
Continual Learning: 4 Rooms Scneario. We consider a continual learning scenario where we train two primitives for two-goal positions ie the goal position is selected randomly from one of the two positions at the start of the episode. The primitives are then transfer (and finetuned) on four-goal
positions then transfer (and finetune) on eight-goal positions. The results are shown in fig. 6. The proposed method achieves better sample efficiency as compared to training a single monolithic policy.

heading: 6 SUMMARY AND DISCUSSION

text: We present a framework for learning an ensemble of primitive policies that can collectively solve tasks without learning an explicit master policy. Rather than relying on a centralized, learned meta-controller, the selection of active primitives is implemented through an information-theoretic mechanism. The learned primitives can be flexibly recombined to solve more complex tasks. Our experiments show that, on a partially observed “Minigrid” task and a continuous control “Ant Maze” walking task, our method can enable better transfer than flat policies and hierarchical RL baselines, including the Meta-learning Shared Hierarchies model and the Option-Critic framework. On Minigrid, we show how primitives trained with our method can transfer much more successfully to new tasks. On the Ant Maze, we show that primitives initialized from a pretrained walking control can learn to walk to different goals in a stochastic, multi-modal environment with nearly twice the success rate of a more conventional hierarchical RL approach, which uses the same pretraining but a centralized high-level policy. The proposed framework could be very attractive for continual learning settings, where one could add more primitive policies over time. Thereby, the already learned primitives would keep their focus on particular aspects of the task, and newly added ones could specialize on novel aspects.

heading: 7 ACKNOWLEDGEMENTS

text: The authors acknowledge the important role played by their colleagues at Mila throughout the duration of this work. AG would like to thank Greg Wayne, Mike Mozer, Matthew Botvinick, Bernhard Schölkopf for very useful discussions. The authors would also like to thank Nasim Rahaman, Samarth Sinha, Nithin Vasisth, Hugo Larochelle, Jordan Hoffman, Ankesh Anand, Michael Chang for feedback on the draft. The authors are grateful to NSERC, CIFAR, Google, Samsung, Nuance, IBM, Canada Research Chairs, Canada Graduate Scholarship Program, Nvidia for funding, and Compute Canada for computing resources. We are very grateful to Google for giving Google Cloud credits used in this project.

heading: B ADDITIONAL RESULTS

text:
B.1 2D BANDITS ENVIRONMENT
In order to test if our approach can learn distinct primitives, we used the 2D moving bandits tasks (introduced in Frans et al. (2017)). In this task, the agent is placed in a 2D world and is shown the position of two randomly placed points. One of these points is the goal point but the agent does not know which. We use the sparse reward setup where the agent receives the reward of 1 if it is within a certain distance of the goal point and 0 at all other times. Each episode lasts for 50 steps and to get the reward, the learning agent must reach near the goal point in those 50 steps. The agent’s action space consists of 5 actions - moving in one of the four cardinal directions (top, down, left, right) and staying still.

heading: B.1.1 RESULTS FOR 2D BANDITS

text:
We want to answer the following questions:
1. Can our proposed approach learn primitives which remain active throughout training? 2. Can our proposed approach learn primitives which can solve the task?
We train two primitives on the 2D Bandits tasks and evaluate the relative frequency of activation of the primitives throughout the training. It is important that both the primitives remain active. If only 1 primitive is acting most of the time, its effect would be the same as training a flat policy. We evaluate the effectiveness of our model by comparing the success rate with a flat A2C baseline. Fig. 8 shows that not only do both the primitives remain active throughout training, our approach also outperforms the baseline approach.

heading: B.2 FOUR-ROOMS ENVIRONMENT

text: We consider the Four-rooms gridworld environment (Sutton et al., 1999c) where the agent has to navigate its way through a grid of four interconnected rooms to reach a goal position within the grid. The agent can perform one of the following four actions: move up, move down, move left, move right. The environment is stochastic and with 1/3 probability, the agent’s chosen action is ignored and a new action (randomly selected from the remaining 3 actions) is executed ie the agent’s selected action is executed with a probability of only 2/3 and the agent takes any of the 3 remaining actions with a probability of 1/9 each.

heading: B.2.1 TASK DISTRIBUTION FOR THE FOUR-ROOM ENVIRONMENT

text: In the Four-room environment, the agent has to navigate to a goal position which is randomly selected from a set of goal positions. We can use the size of this set of goal positions to define a curriculum of task distributions. Since the environment does not provide any information about the goal state, the larger the goal set, harder is the task as the now goal could be any element from a larger set. The choice of the set of goal states and the choice of curriculum does not affect the environment dynamics. Specifically, we consider three tasks - Fourroom-v0, Fourroom-v1 and Fourroom-v2 with the set of 2, 4 and 8 goal positions respectively. The set of goal positions for each task is fixed but not known to the learning agent. We expect, and empirically verify, that the Fourroom-v0 environment requires the least number of samples to be learned, followed by the Fourroom-v1 and the Fourroom-v2 environment (Fig. 6 in the paper).

heading: B.2.2 RESULTS FOR FOUR-ROOMS ENVIRONMENT

text:
We want to answer the following questions:
1. Can our proposed approach learn primitives that remain active when training the agent over a sequence of tasks?
2. Can our proposed approach be used to improve the sample efficiency of the agent over a sequence of tasks?
To answer these questions, we consider two setups. In the baseline setup, we train a flat A2C policy on Fourrooms-v0 till it achieves a 100 % success rate during evaluation. Then we transfer this policy to Fourrooms-v1 and continue to train till it achieves a 100 % success rate during the evaluation on Fourrooms-v1. We transfer the policy one more time to Fourrooms-v2 and continue to train the policy until it reaches a 60% success rate. In the last task(Fourrooms-v2), we do not use 100% as the threshold as the models do not achieve 100% for training even after training for 10M frames. We use 60% as the baseline models generally converge around this value.
In the second setup, we repeat this exercise of training on one task and transferring to the next task with our proposed model. Note that even though our proposed model converges to a higher value than 60% in the last task(Fourrooms-v2), we compare the number of samples required to reach 60% success rate to provide a fair comparison with the baseline.
C IMPLEMENTATION DETAILS
In this section, we describe the implementation details which are common for all the models. Other task-specific details are covered in the respective task sections.
1. All the models (proposed as well as the baselines) are implemented in Pytorch 1.1 unless stated otherwise. (Paszke et al., 2017).
2. For Meta-Learning Shared Hierarchies (Frans et al., 2017) and Option-Critic (Bacon et al., 2017), we adapted the author’s implementations for our environments.
3. During the evaluation, we use 10 processes in parallel to run 500 episodes and compute the percentage of times the agent solves the task within the prescribed time limit. This metric is referred to as the “success rate”.
4. The default time limit is 500 steps for all the tasks unless specified otherwise.
5. All the feedforward networks are initialized with the orthogonal initialization where the input tensor is filled with a (semi) orthogonal matrix.
6. For all the embedding layers, the weights are initialized using the unit-Gaussian distribution.
7. The weights and biases for all the GRU model are initialized using the uniform distribution from U(− √ k, √ k) where k = 1hidden_size .
8. During training, we perform 64 rollouts in parallel to collect 5-step trajectories.
9. The βind and βreg parameters are both selected from the set {0.001, 0.005, 0.009} by performing validation.
In section D.4.2, we explain all the components of the model architecture along with the implementation details in the context of the MiniGrid Environment. For the subsequent environments, we describe only those components and implementation details which are different than their counterpart in the MiniGrid setup and do not describe the components which work identically.

heading: D MINIGRID ENVIRONMENT

text: We use the MiniGrid environment (Chevalier-Boisvert et al., 2018) which is an open-source, gridworld environment package . It provides a family of customizable reinforcement learning environments that are compatible with the OpenAI Gym framework (Brockman et al., 2016). Since the environments can be easily extended and modified, it is straightforward to control the complexity of the task (eg controlling the size of the grid, the number of rooms or the number of objects in the grid, etc). Such flexibility is very useful when experimenting with curriculum learning or testing for generalization.

heading: D.1 THE WORLD

text: In MiniGrid, the world (environment for the learning agent) is a rectangular grid of size say MxN . Each tile in the grid contains either zero or one object. The possible object types are wall, floor, lava, door, key, ball, box and goal. Each object has an associated string (which denote the object type) and an associated discrete color (could be red, green, blue, purple, yellow and grey). By default, walls are always grey and goal squares are always green. Certain objects have special effects. For example, a key can unlock a door of the same color.

heading: D.1.1 REWARD FUNCTION

text:
We consider the sparse reward setup where the agent gets a reward (of 1) only if it completes the task and 0 at all other time steps. We also apply a time limit of 500 steps on all the tasks ie the agent must complete the task in 500 steps. A task is terminated either when the agent solves the task or when the time limit is reached - whichever happens first.
https://github.com/openai/mlsh, https://github.com/jeanharb/option_critic https://github.com/maximecb/gym-minigrid

heading: D.1.2 ACTION SPACE

text:
The agent can perform one of the following seven actions per timestep: turn left, turn right, move forward, pick up an object, drop the object being carried, toggle, done (optional action).
The agent can use the turn left and turn right actions to rotate around and face one of the 4 possible directions (north, south, east, west). The move forward action makes the agent move from its current tile onto the tile in the direction it is currently facing, provided there is nothing on that tile, or that the tile contains an open door. The toggle actions enable the agent to interact with other objects in the world. For example, the agent can use the toggle action to open the door if they are right in front of it and have the key of matching color.

heading: D.1.3 OBSERVATION SPACE

text:
The MiniGrid environment provides partial and egocentric observations. For all our experiments, the agent sees the view of a square of 4x4 tiles in the direction it is facing. The view includes the tile on which the agent is standing. The observations are provided as a tensor of shape 4x4x3. However, note that this tensor does not represent RGB images. The first two channels denote the view size and the third channel encodes three integer values. The first integer value describes the type of the object, the second value describes the color of the object and the third value describes if the doors are open or closed. The benefit of using this encoding over the RGB encoding is that this encoding is more space-efficient and enables faster training. For human viewing, the fully observable, RGB image view of the environments is also provided and we use that view as an example in the paper.
Additionally, the environment also provides a natural language description of the goal. An example of the goal description is: “Unlock the door and pick up the red ball”. The learning agent and the environment use a shared vocabulary where different words are assigned numbers and the environment provides a number-encoded goal description along with each observation. Since different instructions can be of different lengths, the environment pads the goal description with <unk> tokens to ensure that the sequence length is the same. When encoding the instruction, the agent ignores the padded sub-sequence in the instruction.

heading: D.2 TASKS IN MINIGRID ENVIRONMENT

text:
We consider the following tasks in the MiniGrid environment:
1. Pickup: In the Pickup task, the agent spawns at an arbitrary position in a 8× 8 grid (Fig. 9 ). It is provided with a natural language goal description of the form “go pickup a yellow box”. The agent has to navigate to the object being referred to in the goal description and pick it up.
2. Unlock: In the Unlock task, the agent spawns at an arbitrary position in a two-room grid environment. Each room is 8× 8 square (Fig. 10 ). It is provided with a natural language goal description of the form “open the door”. The agent has to find the key that corresponds to the color of the door, navigate to that key and use that key to open the door.
3. UnlockPickup: This task is basically a union of the Unlock and the Pickup tasks. The agent spawns at an arbitrary position in a two-room grid environment. Each room is 8× 8 square (Fig. 11 ). It is provided with a natural language goal description of the form “open the door and pick up the yellow box”. The agent has to find the key that corresponds to the color of the door, navigate to that key, use that key to open the door, enter the other room and pick up the object mentioned in the goal description.

heading: D.3 MODEL ARCHITECTURE

heading: D.3.1 TRAINING SETUP

text: Consider an agent training on any task in the MiniGrid suite of environments. At the beginning of an episode, the learning agent spawns at a random position. At each step, the environment provides observations in two modalities - a 4× 4× 3 tensor xt (an egocentric view of the state of the environment) and a variable length goal description g. We describe the design of the learning agent in terms of an encoder-decoder architecture.

heading: D.3.2 ENCODER ARCHITECTURE

text:
The agent’s encoder network consists of two models - a CNN+GRU based observation encoder and a GRU (Cho et al., 2014) based goal encoder
Observation Encoder:
It is a three layer CNN with the output channel sizes set to 16, 16 and 32 respectively (with ReLU layers in between) and kernel size set to 2× 2 for all the layers. The output of the CNN is flattened and fed to a GRU model (referred to as the observation-rnn) with 128-dimensional hidden state. The output from the observation-rnn represents the encoding of the observation.
Goal Encoder:
It comprises of an embedding layer followed by a unidirectional GRU model. The dimension of the embedding layer and the hidden and the output layer of the GRU model are all set to 128.
The concatenated output of the observation encoder and the goal encoder represents the output of the encoder.

heading: D.3.3 DECODER

text:
The decoder network comprises the action network and the critic network - both of which are implemented as feedforward networks. We now describe the design of these networks.
D.3.4 VALUE NETWORK
1. Two-layer feedforward network with the tanh non-linearity.
2. Input: Concatenation of z and the current hidden state of the observation-rnn.
3. Size of the input to the first layer and the second layer of the policy network are 320 and 64 respectively.
4. Produces a scalar output.

heading: D.4 COMPONENTS SPECIFIC TO THE PROPOSED MODEL

text:
The components that we described so far are used by both the baselines as well as our proposed model. We now describe the components that are specific to our proposed model. Our proposed model consists of an ensemble of primitives and the components we describe apply to each of those primitives.
D.4.1 INFORMATION BOTTLENECK
Given that we want to control and regularize the amount of information that the encoder encodes, we compute the KL divergence between the output of the action-feature encoder network and a diagonal unit Gaussian distribution. More is the KL divergence, more is the information that is being encoded with respect to the Gaussian prior and vice-versa. Thus we regularize the primitives to minimize the KL divergence.

heading: D.4.2 HYPERPARAMETERS

text:
Table 1 lists the different hyperparameters for the MiniGrid tasks.
E 2D BANDITS ENVIRONMENT

heading: E.0.1 OBSERVATION SPACE

text: The 2D bandits task provides a 6-dimensional flat observation. The first two dimensions correspond to the (x, y) coordinates of the current position of the agent and the remaining four dimensions correspond to the (x, y) coordinates of the two randomly chosen points.

heading: E.1 MODEL ARCHITECTURE

heading: E.1.1 TRAINING SETUP

text: Consider an agent training on the 2D bandits tasks. The learning agent spawns at a fixed position and is randomly assigned two points. At each step, the environmental observation provides the current position of the agent as well the position of the two points. We describe the design of the learning agent in terms of an encoder-decoder architecture.

heading: E.1.2 ENCODER ARCHITECTURE

text: The agent’s encoder network consists of a GRU-based recurrent model (referred as the observationrnn) with a hidden state size of 128. The 6-dimensional observation from the environment is the input to the GRU model. The output from the observation-rnn represents the encoding of the observation.

heading: E.2 HYPERPARAMETERS

text: The implementation details for the 2D Bandits environment are the same as that for MiniGrid environment and are described in detail in section D.4.2. In the table below, we list the values of the task-specific hyperparameters.

heading: F FOUR-ROOMS ENVIRONMENT

heading: F.1 THE WORLD

text: In the Four-rooms setup, the world (environment for the learning agent) is a square grid of say 11×11. The grid is divided into 4 rooms such that each room is connected with two other rooms via hallways. The layout of the rooms is shown in Fig. 12. The agent spawns at a random position and has to navigate to a goal position within 500 steps.

heading: F.1.1 REWARD FUNCTION

text: We consider the sparse reward setup where the agent gets a reward (of 1) only if it completes the task (and reaches the goal position) and 0 at all other time steps. We also apply a time limit of 300 steps on all the tasks ie the agent must complete the task in 300 steps. A task is terminate either when the agent solves the task or when the time limit is reached - whichever happens first.

heading: F.1.2 OBSERVATION SPACE

text:
The environment is a 11× 11 grid divided into 4 interconnected rooms. As such, the environment has a total of 104 states (or cells) that can be occupied. These states are mapped to integer identifiers. At any time t, the environment observation is a one-hot representation of the identifier corresponding to
the state (or the cell) the agent is in right now. ie the environment returns a vectors of zeros with only one entry being 1 and the index of this entry gives the current position of the agent. The environment does not return any information about the goal state.

heading: F.2 MODEL ARCHITECTURE FOR FOUR-ROOM ENVIRONMENT

heading: F.2.1 TRAINING SETUP

text: Consider an agent training on any task in the Four-room suite of environments. At the beginning of an episode, the learning agent spawns at a random position and the environment selects a goal position for the agent. At each step, the environment provides a one-hot representation of the agent’s current position (without including any information about the goal state). We describe the design of the learning agent in terms of an encoder-decoder architecture.

heading: F.3 ENCODER ARCHITECTURE

text:
The agent’s encoder network consists of a GRU-based recurrent model (referred as the observationrnn with a hidden state size of 128. The 104-dimensional one-hot input from the environment is fed to the GRU model. The output from the observation-rnn represents the encoding of the observation.
The implementation details for the Four-rooms environment are the same as that for MiniGrid environment and are described in detail in section D.4.2.

heading: G ANT MAZE ENVIRONMENT

text: We use the Mujoco-based quadruple ant (Todorov et al., 2012) to evaluate the transfer performance of our approach on the cross maze environment (Haarnoja et al., 2018). The training happens in two phases. In the first phase, we train the ant to walk on a surface using a motion reward and using just 1 primitive. In the second phase, we make 4 copies of this trained policy and train the agent to navigate to a goal position in a maze (Figure 13). The goal position is chosen from a set of 3 (or 10) goals. The environment is a continuous control environment and the agent can directly manipulate the movement of joints and limbs.

heading: G.0.1 OBSERVATION SPACE

text: In the first phase (training the ant to walk), the observations from the environment correspond to the state-space representation ie a real-valued vector that describes the state of the ant in mechanical terms - position, velocity, acceleration, angle, etc of the joints and limbs. In the second phase (training the ant to navigate the maze), the observation from the environment also contains the location of the goal position along with the mechanical state of the ant.

heading: G.1 MODEL ARCHITECTURE FOR ANT MAZE ENVIRONMENT

heading: G.1.1 TRAINING SETUP

text: We describe the design of the learning agent in terms of an encoder-decoder architecture.

heading: G.1.2 ENCODER ARCHITECTURE

text: The agent’s encoder network consists of a GRU-based recurrent model (referred as the observationrnn with a hidden state size of 128. The real-valued state vector from the environment is fed to the GRU model. The output from the observation-rnn represents the encoding of the observation. Note that in the case of phase 1 vs phase 2, only the size of the input to the observation-rnn changes and the encoder architecture remains the same.

heading: G.1.3 DECODER

text: The decoder network comprises the action network and the critic network. All these networks are implemented as feedforward networks. The design of these networks is very similar to that of the decoder model for the MiniGrid environment as described in section D.3.3 with just one difference. In this case, the action space is continuous so the action-feature decoder network produces the mean and log-standard-deviation for a diagonal Gaussian policy. This is used to sample a real-valued action to execute in the environment.

title: REINFORCEMENT LEARNING WITH COMPETITIVE ENSEMBLES OF INFORMATION-CONSTRAINED PRIMITIVES


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.

summaries: 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.

summaries: 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.

authors:
Ershad Banijamali
Rui Shu
Mohammad Ghavamzadeh
Hung Bui
Ali Ghodsi

title: Robust locallylinear controllable embedding

venue: arXiv preprint arXiv:1710.05373,

authors:
Greg Brockman
Vicki Cheung
Ludwig Pettersson
Jonas Schneider
John Schulman
Jie Tang
Wojciech Zaremba

title: URL http://arxiv.org/ abs/1606.01540

venue: Openai gym. CoRR,

authors:
Yuri Burda
Harri Edwards
Deepak Pathak
Amos Storkey
Trevor Darrell
Alexei A Efros

title: Large-scale study of curiosity-driven learning

venue: arXiv preprint arXiv:1808.04355,

authors:
Yuri Burda
Harrison Edwards
Deepak Pathak
Amos J. Storkey
Trevor Darrell
Alexei A. Efros

authors:
Yevgen Chebotar
Mrinal Kalakrishnan
Ali Yahya
Adrian Li
Stefan Schaal
Sergey Levine

title: Path integral guided policy search

venue: IEEE international conference on robotics and automation (ICRA),

authors:
Marc Deisenroth
Carl E Rasmussen

title: Pilco: A model-based and data-efficient approach to policy search

venue: In Proceedings of the 28th International Conference on machine learning

authors:
Chelsea Finn
Sergey Levine

title: Deep visual foresight for planning robot motion

venue: IEEE International Conference on Robotics and Automation (ICRA),

authors:
Chelsea Finn
Xin Yu Tan
Yan Duan
Trevor Darrell
Sergey Levine
Pieter Abbeel

title: Deep spatial autoencoders for visuomotor learning

authors:
Carles Gelada
Saurabh Kumar
Jacob Buckman
Ofir Nachum
Marc G Bellemare

title: Deepmdp: Learning continuous latent space models for representation learning

authors:
Tuomas Haarnoja
Aurick Zhou
Kristian Hartikainen
George Tucker
Sehoon Ha
Jie Tan
Vikash Kumar
Henry Zhu
Abhishek Gupta
Pieter Abbeel

title: Soft actor-critic algorithms and applications

venue: arXiv preprint arXiv:1812.05905,

authors:
Danijar Hafner
Timothy Lillicrap
Ian Fischer
Ruben Villegas
David Ha
Honglak Lee
James Davidson

title: Learning latent dynamics for planning from pixels

authors:
Brian Ichter
Marco Pavone

title: Robot motion planning in learned latent spaces

venue: IEEE Robotics and Automation Letters,

authors:
Diederik P. Kingma
Max Welling

title: Auto-encoding variational bayes

authors:
Ian Lenz
Ross A Knepper
Ashutosh Saxena

title: Deepmpc: Learning deep latent features for model predictive control

venue: In Robotics: Science and Systems. Rome, Italy,

authors:
Sergey Levine
Vladlen Koltun

title: Guided policy search

venue: In International Conference on Machine Learning, pp

authors:
Volodymyr Mnih
Koray Kavukcuoglu
David Silver
Andrei A Rusu
Joel Veness
Marc G Bellemare
Alex Graves
Martin Riedmiller
Andreas K Fidjeland
Georg Ostrovski

authors:
Junhyuk Oh
Satinder Singh
Honglak Lee

title: Value prediction network

venue: In Advances in Neural Information Processing Systems,

title: Optimization of computer simulation models with rare events

venue: European Journal of Operational Research,

authors:
David Silver
Thomas Hubert
Julian Schrittwieser
Ioannis Antonoglou
Matthew Lai
Arthur Guez
Marc Lanctot
Laurent Sifre
Dharshan Kumaran
Thore Graepel

title: A general reinforcement learning algorithm that masters chess, shogi, and go through self-play

authors:
Manuel Watter
Jost Springenberg
Joschka Boedecker
Martin Riedmiller

title: Embed to control: A locally linear latent dynamics model for control from raw images

venue: In Advances in neural information processing systems,

authors:
Grady Williams
Nolan Wagener
Brian Goldfain
Paul Drews
James M Rehg
Byron Boots
Evangelos A Theodorou

title: Information theoretic mpc for model-based reinforcement learning

authors:
Marvin Zhang
Sharad Vikram
Laura Smith
Pieter Abbeel
Matthew Johnson
Sergey Levine

title: Solar: Deep structured representations for model-based reinforcement learning

heading: 1 INTRODUCTION

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.

heading: 2 RELATED WORK

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.

heading: 3 PRELIMINARIES

heading: 3.1 PROBLEM SETUP

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, γ).

heading: 3.2 MPC PLANNING

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 .

heading: 3.3 SAMPLE-BASED TRAJECTORY OPTIMIZATION

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.

heading: 4 LATENT REWARD PREDICTION MODEL FOR PLANNING

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.

heading: 4.1 LEARNING A LATENT REWARD PREDICTION MODEL

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.

heading: 4.2 MPC WITH THE LATENT REWARD PREDICTION MODEL

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).

heading: 4.3 PLANNING PERFORMANCE GUARANTEE

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.

heading: 5 EXPERIMENTS

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.

heading: 5.1 BASELINES

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.

heading: 5.2 MULTI-PENDULUM

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.

heading: 5.3 MULTI-CHEETAH

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.

heading: 6 CONCLUSION

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.

heading: A ALGORITHM DESCRIPTIONS

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ζ)

heading: B PROOF OF THEOREM 1

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.

heading: C EXPERIMENTS WITH IMAGE OBSERVATIONS

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.