summaries: 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. This paper introduces a vision-based motion planning approach using collocation. Many existing approaches to vision-based control rely on computationally expensive planning approaches using shooting to perform model-based control, which is often only useful in simple control tasks. Collocation approaches are effective in settings with difficult path constraints, and thus exploited by this work to dramatically improve model-based reinforcement learning. The paper studies the problem of planning in domains with sparse rewards where observations are in the form of images. It focuses on solving this problem using model-based RL with emphasis on better trajectory optimization. The proposed solution uses latent models to extract latent representations of the planning problem that is optimized using the Levenberg-Marquardt algorithm (over a horizon). The experimental results show improvements over a) zeroth-order CEM optimization, b)  PlaNet (Hafner et al., 2019) and c)  gradient-based method that optimizes the objective in Eq. 1.
abstractText: The ability to plan into the future while utilizing only raw high-dimensional observations, such as images, can provide autonomous agents with broad and general capabilities. However, realistic tasks may require handling sparse rewards and performing temporally extended reasoning, and cannot be solved with only myopic, short-sighted planning. Recent work in model-based reinforcement learning (RL) has shown impressive results using heavily shaped reward functions that require only short-horizon reasoning. In this work, we study how techniques trajectory optimization can enable more effective long-horizon reasoning. We draw on the idea of collocation-based planning and adapt it to the image-based setting by leveraging probabilistic latent variable models, resulting in an algorithm that optimizes trajectories over latent variables. Our latent collocation method (LatCo) provides a general and effective approach to longer-horizon visual planning. Empirically, our approach significantly outperforms prior model-based approaches on challenging visual control tasks with sparse rewards and long-term goals.
references.authors: Masataro Asai Alex Fukunaga
references.title: Classical planning in deep latent space: Bridging the subsymbolicsymbolic boundary
references.authors: Somil Bansal Anayo K Akametalu Frank J Jiang Forrest Laine Claire J Tomlin
references.title: Learning quadrotor dynamics using neural network for flight control
references.venue: In 2016 IEEE 55th Conference on Decision and Control (CDC),
references.authors: Lars Buesing Theophane Weber Sébastien Racanière S.M. Ali Eslami Danilo Jimenez Rezende David P. Reichert Fabio Viola Frederic Besse Karol Gregor Demis Hassabis Daan Wierstra
references.title: Learning and querying fast generative models for reinforcement learning
references.authors: Yize Chen Yuanyuan Shi Baosen Zhang
references.title: Optimal control via neural networks: A convex approach
references.authors: Kurtland Chua Roberto Calandra Rowan McAllister Sergey Levine
references.title: Deep reinforcement learning in a handful of trials using probabilistic dynamics models
references.venue: In Advances in Neural Information Processing Systems,
references.authors: Junyoung Chung Kyle Kastner Laurent Dinh Kratarth Goel Aaron C Courville Yoshua Bengio
references.title: A recurrent latent variable model for sequential data
references.authors: E. Denton R. Fergus
references.title: Stochastic video generation with a learned prior
references.authors: Yilun Du Toru Lin Igor Mordatch
references.title: Model based planning with energy based models
references.venue: arXiv preprint arXiv:1909.06878,
references.authors: Frederik Ebert Chelsea Finn Alex X Lee Sergey Levine
references.title: Self-supervised visual planning with temporal skip connections
references.authors: Frederik Ebert Chelsea Finn Sudeep Dasari Annie Xie Alex Lee Sergey Levine
references.title: Visual foresight: Model-based deep reinforcement learning for vision-based robotic control
references.authors: Ben Eysenbach Russ R Salakhutdinov Sergey Levine
references.title: Search on the replay buffer: Bridging planning and reinforcement learning
references.venue: In Advances in Neural Information Processing Systems,
references.authors: Marco Fraccaro Søren Kaae Sønderby Ulrich Paquet Ole Winther
references.title: Sequential neural models with stochastic layers
references.venue: In Advances in neural information processing systems,
references.authors: Danijar Hafner Timothy Lillicrap Ian Fischer Ruben Villegas David Ha Honglak Lee James Davidson
references.title: Learning latent dynamics for planning
references.authors: Danijar Hafner Timothy Lillicrap Jimmy Ba Mohammad Norouzi
references.title: Dream to control: Learning behaviors by latent imagination
references.authors: Charles R Hargraves Stephen W Paris
references.title: Direct trajectory optimization using nonlinear programming and collocation
references.venue: Journal of guidance, control, and dynamics,
references.authors: Mrinal Kalakrishnan Sachin Chitta Evangelos Theodorou Peter Pastor Stefan Schaal
references.title: Stomp: Stochastic trajectory optimization for motion planning
references.venue: IEEE international conference on robotics and automation,
references.authors: Matthew Kelly
references.title: An introduction to trajectory optimization: How to do your own direct collocation
references.venue: SIAM Review,
references.authors: Diederik P Kingma Jimmy Ba
references.title: Adam: A method for stochastic optimization
references.authors: Thanard Kurutach Aviv Tamar Ge Yang Stuart J Russell Pieter Abbeel
references.title: Learning plannable representations with causal infogan
references.venue: In Advances in Neural Information Processing Systems,
references.authors: Kara Liu Thanard Kurutach Christine Tung Pieter Abbeel Aviv Tamar
references.title: Hallucinative topological memory for zero-shot visual planning
references.venue: arXiv preprint arXiv:2002.12336,
references.authors: Igor Mordatch Emanuel Todorov Zoran Popović
references.title: Discovery of complex behaviors through contact-invariant optimization
references.venue: ACM Transactions on Graphics (TOG),
references.authors: Anusha Nagabandi Kurt Konolige Sergey Levine Vikash Kumar
references.title: Deep dynamics models for learning dexterous manipulation
references.venue: In Conference on Robot Learning,
references.authors: Suraj Nair Chelsea Finn
references.title: Hierarchical foresight: Self-supervised learning of long-horizon tasks via visual subgoal generation
references.venue: arXiv preprint arXiv:1909.05829,
references.authors: Soroush Nasiriany Vitchyr Pong Steven Lin Sergey Levine
references.title: Planning with goal-conditioned policies
references.venue: In Advances in Neural Information Processing Systems,
references.authors: Giambattista Parascandolo Lars Buesing Josh Merel Leonard Hasenclever John Aslanides Jessica B Hamrick Nicolas Heess Alexander Neitz Theophane Weber
references.title: Divide-and-conquer monte carlo tree search for goal-directed planning
references.authors: Karl Pertsch Oleh Rybkin Frederik Ebert Chelsea Finn Dinesh Jayaraman Sergey Levine
references.title: Long-horizon visual planning with goal-conditioned hierarchical predictors, 2020
references.authors: Nathan Ratliff Matt Zucker J Andrew Bagnell Siddhartha Srinivasa
references.title: Chomp: Gradient optimization techniques for efficient motion planning
references.venue: IEEE International Conference on Robotics and Automation,
references.authors: Nikolay Savinov Alexey Dosovitskiy Vladlen Koltun
references.title: Semi-parametric topological memory for navigation
references.venue: arXiv preprint arXiv:1803.00653,
references.authors: Buser Say Ga Wu Yu Qing Zhou Scott Sanner
references.title: Nonlinear hybrid planning with deep net learned transition models and mixed-integer linear programming
references.venue: In IJCAI,
references.authors: Julian Schrittwieser Ioannis Antonoglou Thomas Hubert Karen Simonyan Laurent Sifre Simon Schmitt Arthur Guez Edward Lockhart Demis Hassabis Thore Graepel
references.title: Mastering atari, go, chess and shogi by planning with a learned model
references.venue: arXiv preprint arXiv:1911.08265,
references.authors: John Schulman Yan Duan Jonathan Ho Alex Lee Ibrahim Awwal Henry Bradlow Jia Pan Sachin Patil Ken Goldberg Pieter Abbeel
references.title: Motion planning with sequential convex optimization and convex collision checking
references.venue: The International Journal of Robotics Research,
references.authors: Manuel Watter Jost Springenberg Joschka Boedecker Martin Riedmiller
references.title: Embed to control: A locally linear latent dynamics model for control from raw images
references.venue: In Advances in neural information processing systems,
references.authors: Tianhe Yu Deirdre Quillen Zhanpeng He Ryan Julian Karol Hausman Chelsea Finn Sergey Levine
references.title: Meta-world: A benchmark and evaluation for multi-task and meta reinforcement learning
references.venue: In Conference on Robot Learning,
references.authors: Marvin Zhang Sharad Vikram Laura Smith Pieter Abbeel Matthew Johnson Sergey Levine
references.title: Solar: Deep structured representations for model-based reinforcement learning
references.venue: In International Conference on Machine Learning,
sections.heading: 1 INTRODUCTION
sections.text: 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.
sections.heading: 2 RELATED WORK
sections.text: 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.
sections.heading: 3 BACKGROUND
sections.text: 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.
sections.heading: 3.1 TRAJECTORY OPTIMIZATION WITH COLLOCATION
sections.text: 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).
sections.heading: 3.2 LATENT VARIABLE MODELS
sections.text: The design of predictive models for high-dimensional visual observations is challenging. Recent work has proposed learning latent state-space models that represent the observations in a compact latent space. Specifically, this work learns a latent space zt, a latent dynamics model pθ(zt+1|zt, at), and decoding models that can be used to decode the observation pθ(ot|zt) and the reward r(zt) from the latent variable (see Fig. 2, left). This approach is powerful due to the use of high-capacity neural network latent dynamics models, and is computationally efficient as the latent space is compact. Importantly, the Markov property of the latent variable can be enforced, allowing a convenient interpretation of the latent state-space as the belief state (Hafner et al., 2019; Buesing et al., 2018).
  A probabilistic latent variable model can then be trained, by maximizing the likelihood of the observed rewards r1:T and images o1:T . While maximizing the exact likelihood is often intractable, we can construct and optimize a variational lower bound on it using a variational distribution qθ(zt+1|ot+1, at, zt) (Chung et al., 2015; Fraccaro et al., 2016):
  ln pθ(o2:T , r1:T |o1, a1:T ) ≥ LELBO(o1:T , a1:T , r1:T ) = Eqθ(z1:T |o1:T ,a1:T ,z0) ∑ t [ln pθ(ot+1|zt+1)− KL (qθ(zt+1|ot+1, at, zt) || pθ(zt+1|zt, at))] . (4)
sections.heading: 4 LATENT COLLOCATION (LATCO)
sections.text: Our aim is to design a collocation method that can be used to produce near-optimal trajectories from raw image observations. A naı̈ve approach would learn a dynamics model over images, and directly optimize an image sequence using Eq (2). However, such as method would be impractical for several
  reasons. First, the optimization over images would make the problem more difficult due to the high dimensionality of the images and the fact that valid images lie on a thin manifold. Second, images typically do not constitute a Markovian state space, violating the assumptions of the method. We propose to instead learn a Markovian and compact state space by means of a latent variable model, and then use this learned state space for collocation.
  Latent state models. To make visual collocation scalable, we leverage the latent state-space dynamics model described in Section 3.2. Following Hafner et al. (2019); Denton & Fergus (2018), we implement this model with convolutional neural networks for the encoder and decoder, and a recurrent neural network for the transition dynamics. The latent state of this model includes a stochastic component with a conditional Gaussian transition function, and the hidden state of the recurrent neural network with deterministic transitions. The model is shown in Fig. 2 (left).
  Latent space collocation. We adapt the collocation procedure described in Section 3.1 into the probabilistic latent state setting. Instead of optimizing a sequence of observations, we optimize a sequence of latent variables z2:T and actions a1:T . Since the latent state space model predicts the reward directly from the latent state zt using an approximation r(zt), we never need to decode images during the optimization procedure, which makes it memory-efficient. We adapt the collocation algorithm to use probabilistic dynamics by enforcing that the next latent state is the mean of the distribution predicted by the model: zt+1 = Epθ(zt+1|zt,at) [zt+1], where the expectation is given simply as the mean of the Gaussian zt+1 = µ(zt, at). Other approaches are possible, such as maximizing the likelihood of the latent state, but these require introduction of additional balance hyperparameters, while our constrained optimization approach automatically tunes this balance. We visualize the collocation procedure in Fig 2 and provide the detailed algorithm as Algorithm 1.
  Algorithm 1 Latent Collocation (LatCo) 1: Start with any available data D 2: while not converged do 3: for each environment step t = 1 . . . Ttot with step Tcache do 4: Infer latent state: zt ∼ q(zt|ot) 5: Define the Lagrangian:
  L(zt+1:t+H , at:t+H , λ) = ∑ t [ r(zt)− λdynt (||zt+1 − µ(zt, at)||2 − )− λactt max(0, |at| − am)2 ] 6: for each optimization step k = 1 . . .K do 7: Update plan: zt+1:t+H , at:t+H := zt+1:t+H , at:t+H + ∇̃L . Eq (5) 8: Update dual variables: λt:t+H := UPDATE(L, λt:t+H) . Eq (6) 9: Execute at:t+Tcache in environment: ot:t+Tcache , rt:t+Tcache ∼ penv
  10: Add episode to replay buffer: D := D ∪ (o1:Ttot , a1:Ttot , r1:Ttot) 11: for training iteration i = 1 . . . It do 12: Sample minibatch from replay buffer: (o1:T , a1:T , r1:T )1:b ∼ D 13: Train dynamics model: θ := θ + α∇LELBO(o1:T , a1:T , r1:T )1:b . Eq (4)
sections.heading: 5 OPTIMIZATION FOR LATENT COLLOCATION
sections.text: The latent collocation framework described in the previous section can be used with any optimization procedure, such as gradient descent or Adam (Kingma & Ba, 2015). However, using appropriate optimization methods is important when time constraints are a factor. We found that the choice of the optimizer for both the latent states and the Lagrange multipliers has a large influence on runtime performance. We detail our specific implementation below.
  Levenberg-Marquardt optimization. We use the Levenberg-Marquardt optimizer for the states and actions, which pre-conditions the gradient direction with the matrix (JTJ)−1, where J is the Jacobian of the objective with respect to the states and actions. This preconditioner approximates the Hessian inverse, significantly improving convergence speed:
  ∇̃ = (JTJ + λI)−1JTρ. (5)
  The Levenberg-Marquardt optimizer has cubic complexity in the number of optimized dimensions. However, by noting that the Jacobian of the problem has a block-tridiagonal structure, we can implement a more efficient optimizer that scales linearly with the planning horizon (Mordatch et al., 2012). This efficient version of the optimizer converges 10-100 times faster than naı̈ve gradient descent on wall clock time in our experiments.
  The Levenberg-Marquardt algorithm optimizes the sum of squares ∑ i ρ 2 i , defined in terms of residuals ρ. Any bounded objective can be expressed in terms of residuals by a combination shifting and square root operations. For the dynamics objective, we use the zt+1 − µ(zt, at) differences as residuals directly, with one residual per state dimension. For the action objective, we use max(0, |at|−am). For the reward objective, we form residuals by passing the negative reward through a softplus operation: ρ = ln(1 + e−r).
  Constrained optimization. The Lagrange multipliers control the balance between the strength of the dynamics constraint and the reward. For fast optimization, it is crucial that the Lagrange multipliers be updated fast when the dynamics constraint is not yet optimized, but the update magnitude should decrease over time to ensure smooth convergence as the dynamic violation converges to zero. We therefore design an update rule that scales with the relative difference between the current dynamics constraint ||zt+1 − µ(zt, at)||2 and the target constraint , such that the update is large when the constraint violation is large and vice versa. We found the following rule to work well in our experiments:
  λdyn += 0.1 log
  ( ||zt+1 − µ(zt, at)||2
  + η ) λdyn, (6)
  where η = 0.01 is used to ensure numerical stability. We found that using a small non-zero is beneficial for the optimization and ensures fast convergence, as the exact constraint might be hard to reach. We additionally constrain the planned actions to be within the environment range using the same constrained optimization procedure.
sections.heading: 6 EXPERIMENTS
sections.text: Our evaluation aims to test the effectiveness of our method on performing long-horizon MBRL from images by answering the following questions: (1) Does LatCo produce effective and feasible plans from visual inputs? (2) Can LatCo solve complex simulated robotic tasks with sparse rewards? (3) How does LatCo compare to trajectory optimization methods used in state-of-the-art visual planning algorithms? (4) What is the effect of LatCo allowing dynamics violations during optimization?
sections.heading: 6.1 EXPERIMENTAL SETUP
sections.text: We evaluate our method on three visual tasks (see Fig. 3): Obstacle, a navigation environment with an obstacle, Sawyer Reach, a robotic reaching task, Sawyer Push, a robotic pushing task. The first
  of these tasks is a custom environment, while the last two are from the Meta-World benchmark (Yu et al., 2020). All tasks only provide the agent with image inputs and no access to underlying state of the simulator. The Obstacle environment consists of an image-based 2D navigation task where the agent must reach a goal on the other side of a wall. The reward defined is the negative distance to the goal. The agent plans over a planning horizon H = 30 steps, executes every action Tcache = 30, Ttot = 30, and results are averaged over 100 runs. In the Sawyer Reach task, a 7-DoF Sawyer arm needs to reach the goal. For this task, H = 30, Tcache = 30, Ttot = 150. In the Sawyer Push task, a 7-DoF Sawyer arm needs to push an object (puck) on the table to the goal. For this task, H = 40, Tcache = 20, Ttot = 120. In both Sawyer tasks, we use a sparse reward of 1 for reaching the goal and 0 otherwise. In particular, for the push task, there is no special reward shaping to encourage the arm to go toward the puck. The arm is initialized to its default position, and results are averaged over 20 runs. We train all methods online according to Algorithm 1. The hyperparameters used in the experiments are detailed in Appendix A.
sections.heading: 6.2 DOES LATCO PRODUCE EFFECTIVE PLANS ON SPARSE-REWARD VISUAL CONTROL TASKS?
sections.text: In our first set of experiments, we evaluate how well LatCo can solve image-based control tasks from sparse rewards, where effectively reasoning into the future is critical for success. As shown in Figure 3, LatCo is indeed able to perform collocation in a learned latent space in order to successfully produce plans that solve these tasks, with effective longerhorizon reasoning. In these results, we see LatCo (top) avoiding local minima that would arise from myopic planners, (middle) extending to higherdimensional control tasks while still operating only from image inputs, and (bottom) reasoning into the future and succeeding even when the sparse task reward provides no direct incentive to go toward the object.
sections.heading: 6.3 HOW DOES LATCO COMPARE TO ALTERNATIVE SHOOTING-BASED APPROACHES?
sections.text: To specifically ascertain the benefits of collocation, we must determine whether the benefits of LatCo stem from gradient-based optimization, from optimizing over states, or both. Therefore, we include a prior method based on zeroth-order CEM optimization, PlaNet (Hafner et al.,
  2019), which is representative of a broader class of shooting-based methods (Ebert et al., 2018; Nagabandi et al., 2020), as well as a gradient-based method that optimizes actions directly using the objective in Eq. 1, which we denote GD actions. To provide a fair comparison that isolates the effects of different planning methods, we use the same dynamics model architecture for all agents.
  From Table 1, we observe that LatCo exhibits superior performance to shooting-based methods on all three tasks. On the Sawyer Reach task, shooting fails to find a successful trajectory most of the time due to a lack of local reward signal. LatCo, on the other hand, solves this task consistently. The Sawyer Push task is challenging for all of the methods, as the dynamics are more complex and requires accurate long-term predictions to correctly optimize the sparse signal. Specifically,
  there is no partial reward for reaching down to the object, so the planner has to look ahead into the high reward state and reason backwards that it needs to approach the object. As shown in Fig. 4, shooting-based methods altogether fail to solve this task, not reaching for the object in any run. LatCo outperforms them by a considerable margin, achieving a 29% success rate. It is worth noting that even in the unsuccessful runs, LatCo plans often display goal-reaching behaviors, as can be seen in Fig. 4. However, significant room for improvement for MBRL still remains on this task. In addition, to evaluate the different methods in a more controlled setup, we test them on the same model trained with offline data in Table 2 and App. C
sections.heading: 6.4 THE EFFECT OF ALLOWING DYNAMICS VIOLATIONS
sections.text: In this section, we analyze the ability of our method to temporarily violate dynamics in order to effectively plan for long-term reward. We show the dynamics violation costs and the predicted rewards associated with the planned trajectories over the course of optimization quantitatively in Fig. 5 and qualitatively in Figs. 1, 6. Since the dynamics constraint is only gradually enforced with the increase of the Lagrange multipliers, the first few steps of optimization allow for dynamics violations in order to focus on discovering the high-reward region (corresponding to steps 0 to 20 in Fig. 5). In the later steps of the optimization, the reward is relaxed and the constraints are optimized until the trajectory becomes feasible.
sections.heading: 7 CONCLUSION
sections.text: We presented LatCo, an algorithm for latent collocation that improves performance of visual-model based reinforcement learning agents. In contrast to the commonly used shooting-based methods, LatCo performs powerful trajectory optimization and is able to plan for tasks where prior work fails, such as long-horizon and sparse reward tasks. By improving the planning capabilities of visual model-based reinforcement learning agents and removing the need for reward shaping, LatCo enables these agents to scale to complex tasks more easily and with less manual instrumentation.
  Future work will examine improvements to the collocation procedure, such as Hamiltonian (Ratliff et al., 2009) or stochastic (Kalakrishnan et al., 2011) optimization. Further, collocation can be applied with variety of latent variable models, including specialized models with more structure. Finally, we believe the ability to perform reinforcement learning in state space as opposed to action space opens up many new avenues for algorithmic development, such as in imitation learning or hierarchical reinforcement learning.
sections.heading: A MODEL ARCHITECTURE AND TRAINING DETAILS
sections.text: We use the latent dynamics and reward models from PlaNet (Hafner et al., 2019) with default hyperparameters. We set the image size to 64x64 and action repeat to 1 for both pointmass and Meta-World models. For every N = 1 episode collected, we train for It = 15 iterations. The pointmass models are trained with episode length 30 and βKL = 0.1, while the Meta-World models are trained with episode length 150 and βKL = 1. All models are trained on a single high-end GPU.
sections.heading: B PLANNING DETAILS
sections.text: CEM. we optimize for 100 iterations. In each iteration, 10000 action sequences are sampled and the distribution is refit to the 100 best samples. For this baseline, we have manually tuned the batch size and number of iterations, and we report the best results.
  GD. we optimize for 500 iterations using the Adam optimizer (Kingma & Ba, 2015) (which
  is a modified version of momentum gradient descent) with learning rate 0.05. We use dual descent to penalize infeasible action predictions. The Lagrange multipliers are updated every 5 optimization steps. For this baseline, we have manually tuned the learning rate and tried several first-order optimizers, and we report the best results.
  LatCo. we optimize for 200 iterations using the Levenberg-Marquardt optimizer with damping 10−3. The damping parameter controls the trust region, with smaller or zero damping speeding up convergence, but potentially leading to numerical instability or divergence. The Lagrange multipliers are updated every step using the rule from Section 5, with = 10−5 and η = 0.01. The threshold directly controls the magnitude of the final dynamics and action violations. In general, we found this parameter to be important for good performance, as a large threshold may cause infeasible plans, while low threshold would make the initial relaxation of the dynamics constraint less effective. However, we observed that a single threshold of 10−5 works for all of our environments. A larger η makes the optimization more aggressive but less stable, and a smaller η diminishes the effect of multiplier updates. We initilalize λdyn0 = 1, λ act 0 = 1.
  These planning hyperparameters remain fixed across the experiments as we observe that reward optimization converges in all cases. Planning a 30-step trajectory takes 12, 14, and 14 seconds for CEM, GD, and LatCo respectively on a single RTX 2080Ti graphics card. We trained the online models for 24 hours each.
sections.heading: C ADDITIONAL EXPERIMENTAL RESULTS
sections.text: Offline training. In Section 6.3, we have evaluated our algorithm in the online training regime, where each agent needs to collect the entire dataset from scratch by executing the policy in the environment. To further test whether collocation is beneficial for planning in a controlled setup, we have performed offline training experiments, where a single model trained with offline data is tested with different planning algorithms. This experiment further shows the applicability of our method to offline setting without additional data collection. We collect the offline data by training an oracle agent that observes dense reward, avoiding the exploration issue. We precollect 200K steps of this data for all tasks. For the Sawyer Push task, we collect an additional 200K environments steps of online interaction data according to Algorithm 1 as none of the methods were able to solve the task with offline data only. We use the same hyperparameters as in Section 6.3 otherwise.
  Table 2 shows the results of this experiment. We see that collocation outperforms shooting methods in this controlled setup, showing that it
  is better able to exploit a dynamics model trained offline.
  We visualize the optimization procedure on the obstacle task in Fig 6.
  We visualize the reward predictor output on the Pushing task in Fig 8.
  We visualize the additional curves of the optimization in Fig 7.
  We further compare the predicted and achieved reward by all methods in the offline setup in Table 3. The predicted reward is often higher than the achieved reward, indicating some degree of model exploitation. However, we see that the predicted reward reflects the general trends in the achieved reward and trajectories with high predicted reward also achieve high reward for our method.