Model-Based Reinforcement Learning via Latent-Space Collocation

1 INTRODUCTION . In order for autonomous agents to perform complex tasks in open-world settings , they must be able to process high-dimensional sensory inputs , such as images , and reason over long horizons about the potential effects of their actions . Recent work in model-based reinforcement learning ( RL ) has shown impressive results in autonomous skill acquisition directly from image inputs , demonstrating benefits such as learning efficiency and improved generalization . While these advancements have been largely fueled by improvements on the modeling side – from better uncertainty estimates and incorporation of temporal information ( Ebert et al. , 2017 ) , to the explicit learning of latent representation spaces ( Hafner et al. , 2019 ) – they leave much room for improvement on the planning and optimization side . Most of the current best-performing deep learning approaches for vision-based planning use only gradient-free action sampling as the underlying optimizer , and are typically applied to settings where a dense and well-shaped reward signal is available . In this work , we argue that more powerful planners are necessary for longer-horizon reasoning . With this goal of long-horizon planning , we aim to extend the myopic planning behavior of existing visual planning methods . Whether it ’ s to avoid local minima due to short-sightedness , or to reason further into the future in order to solve multi-step or sparse-reward tasks , this ability to perform long-horizon planning is critical . Many of the current state-of-the-art visual planning approaches use gradient-free sampling-based optimization methods such as shooting ( Ebert et al. , 2018 ; Nagabandi et al. , 2020 ) ; as shown in Figure 6 , these approaches can get stuck when they must reason further into the future . Here , the curse of dimensionality in conjunction with a lack of shaped reward signal can prevent greedy planning methods from succeeding . Instead , we look to the gradient-based optimization approach of collocation , where crucially , dynamics satisfaction constraints are optimized jointly with the rewards . As shown in Figure 1 , this ability to violate dynamics constraints and imagine even impossible trajectories enables collocation to explore much more effectively and take shortcuts through the optimization landscape in order to learn about high-reward regions before figuring out how to get to those regions . In contrast to shooting approaches , or even to gradient-based approaches such as backpropagating rewards through model predictions , this ability to violate dynamics greatly helps to prevent getting stuck in local minima while still avoiding the need for special dense or shaped reward signals . Collocation , as introduced above , can provide many benefits over other optimization techniques , but it has thus far been demonstrated ( Ratliff et al. , 2009 ; Mordatch et al. , 2012 ; Schulman et al. , 2014 ) mostly in conjunction with known dynamics models and when performing optimization over states . In this work , we are interested in autonomous skill acquisition directly from image inputs , where both the underlying states as well as the underlying dynamics are unknown . Naı̈vely applying collocation to this visual planning setting would lead to intractable optimization , due to the high-dimensional as well as partially observed nature of images , in addition to the fact that only a thin manifold in the pixel space constitutes valid images . Instead , we draw from the representation learning literature and leverage latent dynamics models , which learn a latent representation of the observations that is not only Markovian but also compact and lends itself well to planning . In this learned latent space , we propose to perform collocation over states and actions with the joint objective of maximizing rewards as well as minimizing dynamics violations . Bridging control theory literature with modern deep-RL techniques for learning from images , we propose in this work to perform collocation in a learned latent space in order to enable effective long-horizon planning . The main contribution of this work is an algorithm for latent-space collocation ( LatCo ) , which is an efficient model-based RL approach for solving challenging planning tasks that works directly from image observations by leveraging collocation in a learned latent space . To the best of our knowledge , our paper is the first to scale collocation to visual observations , thus enabling longer-horizon reasoning in model-based RL . In our experiments , we analyze various aspects of our algorithm , and we demonstrate our approach significantly outperforming prior model-based approaches on visual control tasks that require longer-horizon planning with sparse rewards . 2 RELATED WORK . Model-based reinforcement learning . Recent work has scaled model-based reinforcement learning to complex systems leveraging powerful neural network dynamics models ( Chua et al. , 2018 ; Nagabandi et al. , 2020 ) , while showing significant data efficiency improvements over model-free agents . Further , these neural network models can be scaled to high-dimensional image observations using convolutional architectures ( Ebert et al. , 2018 ; Hafner et al. , 2019 ) . However , despite these successes in building better predictive models , planning with these black-box neural network dynamics remains a challenge . While this prior work used simple trajectory optimization techniques like derivative-free shooting , we propose to leverage the more powerful collocation methods . Other work explored more complex approaches based on mixed-integer linear programming ( Say et al. , 2017 ) or gradient descent with input-convex neural networks ( Chen et al. , 2019 ) , but it is unclear whether these approaches scale to visual observations . Latent Planning . Other works have considered different optimization methods such as iterative Linear-Quadratic Regulator ( iLQR ) ( Watter et al. , 2015 ; Zhang et al. , 2019 ) . However , these approaches require specialized locally-linear predictive models , and still rely on shooting and local search in the space of actions , which is prone to local minima . Instead , our collocation approach can be used with any latent state model , and is able to optimize in the state-space , which we show often enables us escape local minima and plan better trajectories . Another line of work relied on graph-based optimization ( Kurutach et al. , 2018 ; Savinov et al. , 2018 ; Eysenbach et al. , 2019 ; Liu et al. , 2020 ) tree search ( Schrittwieser et al. , 2019 ; Parascandolo et al. , 2020 ) , or other symbolic planners ( Asai & Fukunaga , 2018 ) , while we use continuous optimization , which is more suitable for continuous control . Recent work has designed hierarchical planning methods that plan over extended periods of time by considering intermediate subgoals ( Nair & Finn , 2019 ; Nasiriany et al. , 2019 ; Pertsch et al. , 2020 ) . These approaches are orthogonal to our trajectory optimization method , but they present many possible synergies as collocation can be naturally used in hierarchical approaches . Many of these methods ( Buesing et al. , 2018 ; Hafner et al. , 2019 ; 2020 ) used latent state-space models for improved prediction quality and reduced computational requirements . Our proposed method leverages this latent state-space design to construct an effective trajectory optimization method with collocation , and we design our method to be model-agnostic , such that it can benefit from improved latent variable models in the future . Collocation-based planning . Collocation is a powerful optimal control technique for trajectory optimization ( Hargraves & Paris , 1987 ; Witkin & Kass , 1988 ) that optimizes a sequence of states for the sum of expected reward , while eventually enforcing the constrained that the optimized trajectory conform to a dynamics model ( also see Kelly ( 2017 ) for a recent tutorial ) . Prior work in optimal control has explored many versions of collocation for complex motion planning tasks , including Hamiltonian optimization ( Ratliff et al. , 2009 ) , contact-invariant optimization ( Mordatch et al. , 2012 ) , sequential convex programming ( Schulman et al. , 2014 ) , as well as stochastic trajectory optimization ( Kalakrishnan et al. , 2011 ) . We note that in this paper we will use the terms ” trajectory optimization ” and ” planning ” synonymously referring to this type of approaches . These works have demonstrated good results in controlling complex simulated characters , such as humanoid robots , contact-heavy tasks , and tasks with complex constraints . Our work is most similar to that of Schulman et al . ( 2014 ) , however , all this prior work assumed availability of a ground truth model of the environment . Some recent works have attempted using collocation with learned neural network dynamics models ( Bansal et al. , 2016 ; Du et al. , 2019 ) , but it only considered simple low-dimensional dynamics . In this work , we address how to scale up collocation methods to high-dimensional image observations , where direct optimization over images is intractable . We propose to do this by utilizing a learned latent space representation . 3 BACKGROUND . Our method combines collocation methods for trajectory optimization with latent variable models , in order to make it possible to optimize paths over compact and Markovian latent states . We separately review collocation methods and latent state models in this section . 3.1 TRAJECTORY OPTIMIZATION WITH COLLOCATION . Given a dynamics model st+1 = f ( st , at ) that predicts the next state given the previous state and action , a reward function r ( st ) , and the current state s1 , the problem of trajectory optimization is to select actions that maximize the total reward : max a1 : T−1 ∑ t r ( st ) = [ r ( f ( s1 , a1 ) ) + r ( f ( f ( s1 , a1 ) , a2 ) ) + r ( f ( f ( f ( s1 , a1 ) , a2 ) , a3 ) + . . . ] . ( 1 ) Shooting methods optimize this objective directly with respect to the actions . However , this is known to be poorly conditioned due to recursive application of the dynamics function , which results in vanishing or exploding gradients . Instead , we can leverage the structure of the problem to construct an objective that only has pairwise dependencies between temporally adjacent states , and no recursive application of the model . To this end , collocation methods formulate the trajectory optimization problem in Equation 1 as a constrained optimization problem , optimizing over sequences of actions and states , while ensuring that the constraint imposed by the dynamics model f is satisfied : max s2 : T , a1 : T−1 ∑ t r ( st ) s.t . st+1 = f ( st , at ) . ( 2 ) This constrained optimization problem , under some regularity conditions , can be solved with a primal dual approach , which can be formulated as the following saddle point problem : min λ max s2 : T , a1 : T−1 ∑ t r ( st ) − λ||st+1 − f ( st , at ) ||2 . ( 3 ) In practice , we can address this problem with numerical optimization , taking alternating maximization and minimization steps . This collocation approach is better conditioned , as it is able to exploit the Markovian property of the state-space and remove the recursive application of the model . The dynamics constraint only needs to be evaluated on each pair of temporally adjacent states ( see Fig 2 , right ) . Further , the dynamics constraint can be relaxed in the beginning of the optimization , leading the method to rapidly discover high-reward state space regions ( potentially violating the dynamics ) , and then gradually modifying the trajectory to be more dynamically consistent , as illustrated in Figure 1 . This is in contrast to shooting-based approaches , which suffer from severe local optima in long horizon tasks , since the algorithm must simultaneously drive the states toward the high-reward regions and discover the actions that will get them there . The ability to disentangle these two stages by first finding the high-reward region and only then optimizing for the actions that achieve that reward allows collocation methods to solve more complex and temporally extended tasks while suffering less from local optima ( Ratliff et al. , 2009 ; Mordatch et al. , 2012 ; Schulman et al. , 2014 ) .

Summary: In this paper, the authors propose to replace commonly-used shooting-based methods for action sequence planning in learned latent-space dynamics models by a collocation-based method. They argue that shooting-based methods exhibit problematic behavior especially for sparse-reward and long-horizon tasks, as shooting methods do not allow for planning trajectories which (slightly) violate the learned dynamics. The authors propose a collocation method based on Levenberg-Marquard optimization with a scheduled Lagrange multiplier which outperforms two shooting methods (CEM and gradient-based) on a set of robotic tasks.