Model-Based Reinforcement Learning via Latent-Space Collocation
Authors: Oleh Rybkin
Chuning Zhu
Anusha Nagabandi
Kostas Daniilidis
Igor Mordatch
Sergey Levine
Venue:  (2021)
Keywords: visual model-based reinforcement learning
visual planning
long-horizon planning
collocation

Abstract: The ability to construct and execute long-term plans enables intelligent agents to solve complex multi-step tasks and prevents myopic behavior only seeking the short-term reward. Recent work has achieved significant progress on building agents that can predict and plan from raw visual observations. However, existing visual planning methods still require a densely shaped reward that provides the algorithm with a short-term signal that is always easy to optimize. These algorithms fail when the shaped reward is not available as they use simplistic planning methods such as sampling-based random shooting and are unable to plan for a distant goal. Instead, to achieve long-horizon visual control, we propose to use collocation-based planning, a powerful optimal control technique that plans forward a sequence of states while constraining the transitions to be physical. We propose a planning algorithm that adapts collocation to visual planning by leveraging probabilistic latent variable models. A model-based reinforcement learning agent equipped with our planning algorithm significantly outperforms prior model-based agents on challenging visual control tasks with sparse rewards and long-term goals. 

Under review as a conference paper at ICLR 2021
MODEL-BASED REINFORCEMENT LEARNING VIA LATENT-SPACE COLLOCATION
Anonymous authors Paper under double-blind review
ABSTRACT
The ability to plan into the future while utilizing only raw high-dimensional obser- vations, 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, study how techniques trajectory optimization can enable more effective long-horizon reasoning. We draw on the idea of collocation-based planning and adapt the image-based setting by lever- aging 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.
INTRODUCTION
In order for autonomous agents to perform complex tasks in open- worldworld settings, they must be able to process sigh-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 (RI.) 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 the modeling side from better uncertainty estimates and incorporation representation spaces (Hafner et al. 2019) they leave much room for improvement the planning and optimization side. Most of the current best-performing deep learning approaches for vision- based planning only gradient-free action sampling the underlying optimizer, and typically applied to settings where a dense and well-shaped reward signal is available. In this work, we argue that more powerful planners necessary for longer-horizon reasoning.
temporal information (Ebert et al. 2017). to the explicit learning of latent
With this goal of long-horizon planning, we aim to extend the myopic planning behavior of existing reason
visual planning methods. Whether it's to avoid local minima due short-sightedness, further into the future in order to solve multi-step or sparse-reward tasks, this ability to perform long-hon horizon planning critical. Many of the current of-theof visual planning approaches use gradient-free sampling-base optimization methods such as shooting (Ebert et al. 2018 Nagabandi
wi lig sonee
Optimization
Figure 1: Collocation-bast planning. Each image shows full plan at that optimization step. Collocation jointly optimizes dynamics satisfaction constraints as well as rewards. This ability to violate dynamics allows for the rapid discovery of high-reward regions (where the object is next the goal), while the subsequent refinement of the planned trajectory focuses on feasibly achieving it.

---

Under review as a conference paper at ICLR 2021
et al., 2020); 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 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 learn about high-reward regions before figuring out how to get to those regions. In contrast to shooting approaches, or even to gradient-based 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 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- high-dimensional as well as partially observed nature of images, in addition to the fact that only thin manifold in the pixel space constitutes valid images. Instead, we draw from the representation learning literature and leverage latent dynamics models, which learn 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 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, longer-horizze reasoning in model- based RL. In our experiments, we analyze various aspects of our algorithm, and we demonstrate approaches on visual control tasks that require longer- planning with sparse rewards.
paper is the first to scale collocation to visual observations, thus enabling
approach significantly outperforming prior model-based
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 Naga bandi et al. 2020), while showing significant data efficiency improvements over model-free agents. Further, these neural network models can be scaled to high- high-dimensiooal 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-fre freee shooting, we propose 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-Quadratty 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., 22018] Savinov et al. 20188] Eysenbach et al. 2019 Liu et al. 2020) tree search (Schrittwieser et al. 2019 Parascandolo al. 2020), 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
2

---

Under review as a conference paper at ICLR 2021
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 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 can benefit from improved latent variable models in the future.
Collocation- planning. Collocation is a powerful optimal control technique for trajectory optimization (Hargraves & Paris 1987 Witkin & Kass 1988) that optimizes sequence of states for the sum of expected reward, while eventually enforeing the constrained that the optimized trajectory conform to dynamics model (also see Kelly (217)) for recent tutorial). Prior work in optimal control has explored many versions of collocation for complex motion planning tasks, including Hamiltonian optimization (Ratliff' et al. 22009), contact-invariant optimization (Mordatch et al. 2012 sequential convex programming (Schulman al. 2014), as well as stochastic trajectory optimization (Kalakrishnan et al. 201T]. We note that in this paper 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 humanoid robots, contact-heavy tasks, and tasks with complex constraints. Our work most similar to that of Schulman et al.] (2014). however, all this prior work assumed availability of ground truth model of the environment. Some recent works have attempted using collocation with learned neural network dynamics models (Bansal al. 2016 Du al. 2019). but it only considered simple -wwdimensional 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 86+1 = [( az) that predicts the next state given the previous state and action, reward function r(s)), and the current state $1, the problem of trajectory optimization is to select actions that maximize the total reward:
max Σ r((:) [r(f(si,a)) + (f((f()))))))))) +r(((f(((L)a)) +...] (1)
01/T
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 [ as constrained optimization problem, optimizing sequences of actions and states, while ensuring that the constraint imposed by the dynamics model f is satisfied:
max
Σ E(()))
s.t.
8141
s( at).
$2:2-91:T
This constrained optimization problem, under some regularity conditions, can solved with a primal dual approach, which can be formulated as the following saddle point problem:
min max $2.7 R11TT
Al| s(r,, az)]
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 exploit the
3

---

Under review as a conference paper at ICLR 2021
Rewandss
p
LatCo
InLX T3
T3s
F
Resard Model
+(*)
r(2)
AA (,) = 2g
jA ag)
2g
p Mag ar)
p
Eatent Modil
1
2,
Z
Z1
a
Zy
Qg
Zy
ag
Z
Excoder
[(z|o)
Optiniand variable
Shooting
max
maX T3
ILLAX T
a2
a2
ag
Ofjictioe
Image
Modei
21
pa 0)
Ze
Hu (2,0)
Zz
μa 2,
21
Figure 2: Latent Collocation (LatCo). Left: Our latent state-space model, with an encoder g and latent state-space dynamies model p(2+41/2t; Qt) ~ [N(p(2, (t), (2,, az)). reward model r(zx) predicts the reward from the latent state. The model is trained with a variational lower bound to reconstruct the observations (not shown). Right: comparison of LatCo and shooting methods. LatCo optimizes sequence of latent states and actions 22.7, d1:T to maximize rewards r(x)) as well as satisfy dynamics 2++1 = p(2,, qz). This joint optimization allows to relax the dynamics constraint early on, which helps escape local minima. In contrast, shooting methods require recursive application of the dynamics and backpropagation through time, which is often difficult to optimize.
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- state space regions (potentially violating the dynamics), and then gradually modifying the trajectory to be more dynamically consistent, as illustrated in Figure This is in contrast to mootting-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).
3.2 LATENT VARIABLE MODELS
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 compact latent space. Specifically, this work learns a latent space 2x, a latent dynamics model po(2z+) (z+ | az). and decoding models that can be used decode the observation (o (zx) and the reward r(z)) 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 71:7 and images 01:7- While maximizing the exact likelihood is often intractable, we can construct and optimize a variational lower bound on it using a variational distribution (= 10441, ae, 2x) (Chung et al. 2015 Fraccaro al. 2016):
In pola(22: F1. 01, (1:T) (11,)))) 2 LIBBOooLT, (irnn,r, Soolsurroour
Σ [[npe(0+112))1) KL (go lou+1. ae, 22) po ax)))) (4)
LATENT COLLOCATION (LATCO)
Our aim is to design collocation method that can be used to produce near-optimal trajectories from raw image observations. A naive approach would learn a dynamics model over images, and directly optimize image sequence using Eq (2) However, such method would be impractical for several
4

---

Under review as a conference paper at ICLR 2021
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 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 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 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 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. (left).
Latent space collocation. We adapt the collocation procedure described in Section 3.1 into the probabilistic latent state setting. Instead optimizing sequence observations, we optimize sequence of latent variables 22:T and actions a1., Since the latent state space model predicts the reward directly from the latent state ot using an approximation r(z)), we never need to decode images during the optimization procedure, which makes it memory-efficient. We adapt the collocation algorithm to use probabilistic dynamiis by enforcing that the next latent state the mean of the distribution predicted by the model: 2+11 Ep. ( [24+1], where the expectation is given simply as the mean of the Gaussian 2t+1 = (==t,aa)) Other approaches
possible, such as maximizing the likelihood the latent state, but these require introduction of additional balance hyperparameterss while constrained optimization approach automatically tunes this balance. We visualize the collocation procedure in Fig 2igg2an and provide the detailed algorithm as Algorithm
Algorithm 1 Latent Collocation (LatCo)
1: Start with any available data DD
2: while not converged do
3: for each environment step / = 1... .....2 Tux with step Teache do
4: Infer latent state: 2₁ ~ q(=z(0t)
5:
Define the Lagrangian:
[204110- HH.- A) Σ [((z)) - 21 - a(2) -e) X² maxx0,, lael am)
6:
for each optimization step & 1... K do
7:
Update plan:
2++1+H OcL+H 20+1++ Ot&+ VL
Eq D Eq
8:
Update dual variables: A := UPDATE C, A
9:
Execute Tua++Toee in environment: Otst+Tuae" Titt+Tate Penv Add episode to replay buffer: D = DU (OI:T. 01/Ta: (1:Tum for training iteration i 1.. do
10:
12: Sample minibatch from replay buffer: (OLTT 01.T, (1:T): ~ D
Train dynamics model: # == V an.r, F1:1 F1:111
Eq
5 OPTIMIZATION FOR LATENT COLLOCATION
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 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-Marquardd optimization. We use the Aevenberg-Marquaar optimizer for the states and actions, which pre-conditions the gradient direction with the matrix (J"J)-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:
V (J') + -1T
(5)
5

---

Under review as a conference paper at ICLR 2021
Obstacle
Sawyer Reach
Sawyer Push
Figure 3: LatCo executed trajectories on the three considered tasks. all these sparse reward settings requiring temporally extended reasoning, LatCo is able to optimize effective plans and execute them.
The Levenberg-Marquardd 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 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 naive gradient descent wall clock time in our experiments.
The Levenberg
algorithm optimizes the sum of squares E, p defined terms of residuals p. Any bounded objective can be expressed in terms of residuals by combination shifting and square root operations. For the dynamics objective, we use the žt+1 4(zz, ar) differences as residuals directly, with one residual per state dimension. For the action objective, we use max(0, |ae]- (eel-am). For the reward objective, we form residuals by passing the negative reward through softplus operation: p = (1 + +e="
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 zero. We therefore design update rule that scales with the relative difference between the current dynamics constraint || fu(z,, at [l2 and the target constraint e, such that the update is large when the constraint violation is large and vice versa. We found the following rule to work well in experiments:
|| - ((21,,)))) + A
dyn += 0.1 log
where 77 0.01 is used to ensure numerical stability. We found that using 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.
6 EXPERIMENTS
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-an visual planning algorithms? (4) What is the effect of LatCo allowing dynamics violations during optimization?
6.1 EXPERIMENTAL SETUP
We evaluate our method on three visual tasks (see Fig. By: Obstacle, a navigation environment with obstacle, Sawyer Reach, robotic reaching task, Sawyer Push, robotic pushing task. The first

---

Under review as a conference paper at ICLR 2021
of these tasks is custom environment, while the last two from the Meta-World benchmark Yu jet 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 image- amage-based 2D navigation task where the agent must reach a goal on the other side of a wall. The reward defined the negative distance to the goal. The agent plans over planning horizon H 30 steps, executes every action Tenche = 30, Tot = 30, and results are averaged over 100 runs. In the Sawyer Reach task, a 7-DoF Sawyer arm needs reach the goal. For this task, H = 30, Teache = 30, Ti 150. In the Sawyer Push task, a 7-DoF Sawyer am needs to push an object (puck) on the table to the goal. For this task, H 40, Tecche 20, Two 120. In both Sawyer tasks, we use sparse reward of 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 initialized to its default position, and results are averaged over 20 runs. We train all methods online according to Algorithm The hyperparameters used in the experiments are detailed in Appendix A
6.2 DOES LATCO PRODUCE EFFECTIVE PLANS ON REWARD VISUAL CONTROL TASKS?
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 in- deed able perform collocation in a learned latent space in order to successfully produce plans that solve these tasks, with effective longer- horizon reasoning. In these results, we see LatCo (top) avoiding local minima that would arise from myopic plan- ners, (middle) extending to higher- dimensional 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.
LatCo Success (plan)
LatCo Success (execution)
LatCo Failure (plan)
LatCo Fallure (execution)
PlaNet (plan)
6.3 How DOES
PlaNet (execution)
LATCO COMPARE TO ALTERNATIVE SHOOTING-BASE BASED APPROACHES?
To specifically ascertain the bene- fits of collocation, must deter- mine whether the benefits of LatCo stem from gradient-based optimiza- tion, from optimizing over states, or both. Therefore, we include prior method based on zeroth-order CEM optimization, PlaNet (Hafner et al. 2019), which is representative of broader class of shooting-bars 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. which denote GD actions. To provide fair comparison that isolates the effects of different planning methods, we use the same dynamics model architecture for all agents.
Figure 4: Planned and executed trajectories on the Sawyer Push task. LatCo produces feasible and effective plan, and is able to execute it. Shooting struggles to solve this task as dense or shaped rewards are not provided.
From Table I 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 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 correctly optimize the sparse signal. Specifically,
7

---

Under review as a conference paper at ICLR 2021
Table 1: Comparison of shooting and collocation with online data. Shooting struggles with long horizon tasks and sparse rewards, while the powerful trajectory optimization with LatCo is able to find a good trajectory.
Sawyer Reaching
Sawyer Pushing
Shaped reward Visual observations
Return
Success
Return
Success
Shooting CEM (PlaNet)
0.0 +0.0
0%
0.0 +:0.0
0%
Shooting GD LatCo (Ours)
38.0 1:5.1
42%
4.3 41.5
10%
75.8 +200
100%
10,9 =2.1
29%
Table 2: Comparison of shooting and collocation with offline data.
Obstacle
Sawyer Reaching
Sawyer Pushing
Shaped reward
Visual observations
Return
Success
Return
Success
Return
Success
Shooting CEM (PlaNet)
-52.4 +0.6
25%
1.4 +0.8
15%
0.0 =0.0
0%
Shooting GD
-54.5
+0.5
34%
4.8 13.8
10%
0,0 10.0
0%
LatCo (Ours)
-47.3
+1.1
54%
90.9 +5.8
100%
18.7 +7.0
40% %
there is no partial reward for reaching down the object, so the planner has 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 considerable margin, achieving 29% success rate. It is worth noting that even 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.
6.4 THE EFFECT OF ALLOWING DYNAMICS VIOLATIONS
In this section, we analyze the ability of our method to temporarily violate dynamics order to effectively plan for long-term reward. We
1
Dynamics violation
100
Predicted rewards
show the dynamics violation costs and the pre- dicted rewards associated with the planned tra- jectories over the course of optimization quan- titatively Fig. and qualitatively in Figs. Since the dynamics constraint is only gradually enforced with the increase of the Lagrange mul- tipliers, the first few steps of optimization allow for dynamics violations order to focus on dis- covering the high-reward region (corresponding to steps 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.
10a
50
10-1
10-3
-50
50
100
150
200
Optimization step
Figure 5: Dynamics violation and reward predic- tions of planned trajectories over the course of 200 optimization steps on the obstacle task. Earlier on, the planner explores high-reward regions. converg- ing to the goal state. Eventually, as the dynamics constraint is enforeed, the plan becomes feasible, while maintaining a high predicted reward.
CONCLUSION
We presented LatCo, an algorithm for latent collocation that improves performance visual-moded 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.
8

---

Under review as a conference paper at ICLR 2021
Future work will examine improvements to the collocation procedure, such as Hamiltonian Rathif
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, believe the ability to perform reinforcement learning state space as opposed to action space opens up many avenues for algorithmic development, such in imitation learning hierarchical reinforcement learning.
REFERENCES
Masataro Asai and Alex Fukunaga. Classical planning in deep latent space: Bridging the subsymbolic- symbolic boundary, AAAI, 2018.
Somil Bansal, Anayo K Akametalu, Frank Jiang, Forrest Laine, and Claire J Toomlin. Learning quadrotor dynamics using neural network for flight control. In 2016 IEEE 55th Conference on Decision and Control (CDC), pp. 4653-4660. IEEE, 2016.
Lars Buesing, Theophane Weber, Sébastien Racanière, M. Ali Eslami, Danilo Jimenez Rezende, David P. Reichert, Fabio Viola, Frederic Besse, Karol Gregor, Demis Hassabis, and Daan Wierstra. Learning and querying fast generative models for reinforcement learning. arXiv:/802.03006, 2018.
Yize Chen, Yuanyuan Shi, and Baosen Zhang. Optimal control via neural networks: A convex approach. ICLR, 2019.
Kurtland Chua, Roberto Calandra, Rowan McAllister, and Sergey Levine. Deep reinforcement learning in handful of trials using probabilistic dynamics models. In Advances in Neural Information Processing Systems, pp. 4754-4765, 2018.
Junyoung Chung, Kyle Kastner, Laurent Dinh, Kratarth Goel, Aaron Courville, and Yoshua Bengio. A recurrent latent variable model for sequential data. 2015.
E. Denton and R. Fergus. Stochastic video generation with learned prior. 2018.
Yilun Du, Toru Lin, and Igor Mordatch. Model based planning with energy based models. arXiv preprint arXiv:1909.0665 06878, 2019.
Frederik Ebert, Chelsea Finn, Alex X Lee, and Sergey Levine. Self-supervised visual planning with temporal skip connections. 2017.
Frederik Ebert, Chelsea Finn, Sudeep Dasari, Annie Xie, Alex Lee, and Sergey Levine. Vi- sual foresight: Model-based deep reinforcement learning for vision-based robotic control. arXiv: 1rXiv:1812.00555 2018.
Ben Eysenbach, Russ R Salakhutdinov, and Sergey Levine. Search on the replay buffer: Bridging planning and reinforcement learning. In Advances in Neural Information Processing Systems, pp. 15257, 2019.
Marco Fraccaro, Søren Kaae Sønderby, UIrich Paquet, and Ole Winther. Sequential neural models with stochastic layers. In Advances in neural information processing systems, pp. 2199-2207, 2016.
Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. 2019.
Danijar Hafner, Timothy Lillicrap. Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. ICLR, 2020.
Charles R Hargraves and Stephen W Paris. Direct trajectory optimization using nonlinear program- ming and collocation. Journal of guidance, control, and dynamics, 10(4):338-32,, 1987.
Mrinal Kalakrishnan, Sachin Chitta, Evangelos Theodorou, Peter Pastor, and Stefan Schaal. Stomp: Stochastic trajectory optimization for motion planning. In 201 IEEE international conference on robotics and automation, pp. 4569-4577 IEEF, 2011.

---

Under review as a conference paper at ICLR 2021
Matthew Kelly. An introduction trajectory optimization: How to do your own direct collocation.
SIAM Review, 59())8449-))))) 904, 2017.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. 2015.
Thanard Kurutach, Aviv Tamar, Ge Yang, Stuart J Russell, and Pieter Abbeel. Learning plannable representations with causal infogan. In Advances Neural Information Processing Systems, pp- 8733-8744, 2018.
Kara Liu, Thanard Kurutach, Christine Tung, Pieter Abbeel, and Aviv Tamar. Hallucinative topologi- cal memory for zero-shot visual planning. arXiv preprint arXiv:2002. 2020.
Igor Mordatch, Emanuel Todorov, and Zoran Popović. Discovery of complex behaviors through contact-invariant optimization. ACM Transactions on Graphics (TOG), (1)::11-8, 2012.
Anusha Nagabandi, Kurt Konolige, Sergey Levine, and Vikash Kumar. Deep dynamics models for learning dexterous manipulation. In Conference on Robot Learning, pp. 1101-1112, 2020.
Suraj Nair and Chelsea Finn. Hierarchical foresight: Self-supervised learning of long-horizon tasks via visual subgoal generation. arXiv preprint
2019.
Soroush Nasiriany, Vitchyr Pong, Steven Lin, and Sergey Levine. Planning with goal-conditioned policies. In Advances in Neural Information Processing Systems, pp. 14843-14854, 2019.
Giambattista Parascandolo, Lars Buesing, Josh Merel, Leonard Hasenclever, John Aslanides, Jes- sica B Hamrick, Nicolas Heess, Alexander Neitz, and Theophane Weber. aniie-and-conquer monte carlo tree search for goal-directed planning. arXiv preprint arXiv:2004. /410, 2020.
Karl Pertsch, Oleh Rybkin, Frederik Ebert, Chelsea Finn, Dinesh Jayaraman, and Sergey Levine Long-horizon visual planning with goal-conditioned hierarchical predictors, 2020.
Nathan Ratliff, Matt Zucker, J Andrew Bagnell, and Siddhartha Srinivasa. Chomp: Gradient
optimization techniques for efficient motion planning. In 2009 IEEE International Conference on Robotics and Automation, pp. 489-494. IEBE, 2009.
Nikolay Savinov, Alexey Dosovitskiy, and Vladlen Koltun. Semi-parametrie topological memory for navigation. arXiv preprint arXiv:/803.000
2018
Buser Say, Ga Wu, Yu Qing Zhou, and Scott Sanner. Nonlinear hybrid planning with deep net learned transition models and mixed-integer linear programming. In IJCAI, pp. 750-756, 2017.
Julian Schrittwieser, loannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, et al. Mastering atari, go, chess and shogi by planning with a learned model. arXiv preprint arXiv: 2019.
John Schulman, Yan Duan, Jonathan Ho, Alex Lee, Ibrahim Awwal, Henry Bradlow. Jia Pan, Sachin Patil, Ken Goldberg, and Pieter Abbeel. Motion planning with sequential convex optimization and convex collision checking. The International Journal of Robotics Research, 3319): 2014.
Manuel Watter, Jost Springenberg, Joschka Boedecker, and Martin Riedmiller. Embed to control: A locally linear latent dynamics model for control from images. In Advances in neural information processing systems, pp. 2746-2754, 2015.
Andrew Witkin and Michael Kass. Spacetime constraints. ACM Siggraph Computer Graphics, 22(4): 159-168, 1988.
Tianhe Yu, Deirdre Quillen, Zhanpeng He, Ryan Julian, Karol Hausman, Chelsea Finn, and Sergey Levine. Meta-world: A benchmark and evaluation for multi-task and meta reinforcement learning. In Conference on Robot Learning, pp. 1094-1100, 2020.
Marvin Zhang, Sharad Vikram, Laura Smith, Pieter Abbeel, Matthew Johnson, and Sergey Levine. Solar: Deep structured representations for model-based reinforcement learning. In International Conference on Machine Learning, pp. 7444-7453. PMLR, 2019.
10

---

Under review as a conference paper at ICLR 2021
10
Dynamics violetion
10'
coeficieni
100
Predicted
10-
10a
1/h
10
10
10.
100 Optimization stop
Optimization
Optimization
Figure 7: Additional optimization curves. The dynamics coefficient (magnitude of Lagrange multipli-
ers) increases exponentially as the dynamics constraint is enforced, and eventually converges
A MODEL. ARCHITECTURE AND TRAINING DETAILS
We use the latent dynamics and reward mod- els from PlaNet (Hafner et al. 2019) with de- fault hyperparameters. We set the image size to 64x64 and action repeat mass and Meta- World models. For every N = episode collected, we train for It 15 iterations. The pointmass models are trained with episode length 30 and SKi.- = 0.1, while the Meta- World models trained with episode length 150 and 3KL- = 1. All models are trained on single high-end GPU.
Shooting
Collocation
for both point-
Stort
Optiintattin
Figure 6: Visualization of collocation-based plan- ning. Left: shooting-based planning searches for actions that maximize rewards of the transitions as predicted by a model. struggles on tasks that re- quire long-term reasoning, as its myopic approach prone to local minima such as being trapped behind the obstacle here. Right: in contrast, col- location jointly optimizes dynamics satisfaction constraints as well as rewards. This ability to vi- olate dynamics allows for the rapid discovery of high-reward regions (here, the goal state), while the subsequent refinement of the planned trajectory focuses feasibly achieving it.
B PLANNING DETAILS
CEM. we optimize for 100 iterations. In each iteration, 10000 action sequences are sampled and the distribution is refit to the 100 best sam- ples. For this baseline, we have manually tuned the batch size and number of iterations, and we report the best results.
GD. optimize for 500 iterations using the
Adam optimizer (Kingma & Ba (2015) (which a modified version of momentum gradient descent) with learning rate 0.05. We 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.
optimizer with damping 10-3, The damping parameter controls the trust region, with smaller 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 e 10 and r 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 single threshold of 10-5 works for all of our environments. A larger y makes the optimization more aggressive but less stable, and a smaller 7 diminishes the effect of multiplier updates. We initilalize 2 1, Aast 1.
LatCo. we optimize for 200 iterations using the Levenberg-
These planning hyperparameters remain fixed across the experiments as we observe that reward optimization converges in all cases. Planning 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.
11

---

Under review as a conference paper at ICLR 2021
Table 3: Comparison of predicted and achieved reward.
Obstacle
Sawyer Reaching
Sawyer Pushing
Predicted
Real
Predicted
Real
Predicted
Real
Shooting CEM (PlaNet)
-8.3 #0.1
-54.7 +2.3
#2.3
0.0 +0.0
10.0
0.0 #0.0
Shooting GD
-00.8 +2.3
-49.3 +2.0
5.9 +0.7
0.0 +0.0
3.0 4..0
0.0:0..0
LatCo (Ours)
-20.7 +0.9
-46.9 +2.7
141.4 +2.8
90.8 +4.77
29.3 +5.9
18.2 +6.5
C ADDITIONAL EXPERIMENTAL RESULTS
we have evaluated our algorithm the online training regime, where each agent needs to collect the en- tire dataset from scratch by executing the policy
Offline training. In Section 6.3
1.2
1.0
object y-cooddinate
0.2
0.8
in the environment. To further test whether col- location beneficial for planning a controlled setup, have performed offline training experi- ments, where a single model trained with offline data is tested with different planning algorithms. This experiment further shows the applicability of method to offline setting without addi- tional data collection. We collect the offline data by training an oracle agent that observes dense reward, avoiding the exploration issue. We pre- collect 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 | 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
0.6
Goo!
0.0
0.4
Surt
0.2
-0.
0.0
-0.2
-0.4
0.4
0.6
0.8
object x-coordinate
Figure 8: Visualization of the reward predictor for the Sawyer Pushing task. The output of the reward predictor is shown for each object position on the 2D table. We see that the reward predictor correctly predicts value of 1 at the goal, and low values otherwise. In addition, there is certain amount of smoothing induced by the reward predictor, which creates gradient from the start the goal posi- tion. This explains why gradient- based planning is applicable even in this sparse reward task.
Table shows the results of this experiment. We see that collocation outperforms shooting methods in this controlled setup, showing that better able to exploit dynamics model trained offline.
We visualize the optimization procedure the obstacle task in Fig
We visualize the reward predictor output on the Pushing task Fig
We visualize the additional curves of the optimization in Fig
We further compare the predicted and achieved reward by all methods in the offline setup in Table
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 method.
12


Reviews:
[{"review_id": "PaPsLydRIdG", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary: 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.\n\nStrengths:\n\ni) The motivation, organization and the overall writing of the paper are clear.\n\nii) The tested experimental domains are good representatives of the realistic planning setting identified in the paper.\n\nWeaknesses:\n\ni) Discussion of literature on planning in latent spaces [1,2,3,4,5] is left out and should be included. Namely, [1,2] performs (classical) planning from images, and [3,4,5] perform planning with learned neural models. Here, space can be saved by removing Figure 4 since all of its subfigures look identical given their (visual) quality.\n\nii) Have you tried solving Eq. 2. directly similar to [4]? It seems more appropriate baseline compared to c) (i.e., as labeled above).\n\niii) How do you reason about the length of the horizon T? For example [1,2] use heuristic search.\n\niv) There does not seem to be any presentation of hyperparameter selection/optimization, runtime results or quality of solutions. Table 1 is too high-level to provide any meaningful insight into understanding how each method compares. Similarly, Figure 5 is very hard to read and not clear what each axis represents. Overall, I would say this is the weakest part of the paper.\n\nReferences:\n\n[1] Classical Planning in Deep Latent Space: Bridging the Subsymbolic-Symbolic Boundary, Asai and Fukunaga AAAI-18.\n\n[2] Learning Neural-Symbolic Descriptive Planning Models via Cube-Space Priors: The Voyage Home (to STRIPS), Asai and Muise IJCAI-20.\n\n[3] Nonlinear Hybrid Planning with Deep Net Learned Transition Models and Mixed-Integer Linear Programming, Say et al., IJCAI-17.\n\n[4] Scalable Planning with Deep Neural Network Learned Transition Models, Wu et al. JAIR.\n\n[5] Optimal Control Via Neural Networks: A Convex Approach, Chen et al., ICLR 2019.\n\n** Post Rebuttal **\n\nTo best of my understanding, the authors have addressed all my questions and suggestions with the appropriate revision of their paper. Specifically, the necessary discussion of hyperparameter selection is added and presentation of the runtime&solution quality results (i.e., raised in point iv)) have been improved with the inclusion of important details, additional discussion of related work is added (i.e., raised in point i)) and questions are addressed (i.e., raised in point ii) and iii)). As such, I have updated my rating accordingly.", "extra_scores": {}, "extra_text": {}}, {"review_id": "PCH864P5UCH", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "## Paper summary\n\nThis 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.\n\nI like the idea, but it is a relatively small extension to existing work, so I am inclined to rate this paper as marginally below the acceptance threshold. I would be willing to revise my score if the paper was revised to \n- better clarify the algorithm to align with the methods used in experiments\n- better justify the reasons why ilQR trajectory optimisation with locally linear dynamics models was not used as a baseline  (or even better, include this as a baseline)\n\n### Pros\n\n- The paper is well written and clearly laid out.\n- Solving a collocation problem in the latent space is a sensible approach, and a much better idea than using CEM planning or shooting. \n\n### Cons\n- It's a reasonably straightforward application of collocation in a learned latent space. While I have not seen this done previously, it is a relatively obvious improvement.\n- The paper motivates the need for collocation in the context of *long horizon tasks*, where shooting performs poorly. However, none of the tasks (pushing and reaching in free space) considered in this work are long horizon tasks, or particularly challenging. \n\n### General recommendations for improvement and queries \n\n-  I'd recommend replacing the term *long horizon tasks* with something more suitable, along the lines of what is actually demonstrated in the experimental results, eg. *vision-based motion planning*.\n- Page 2 - Latent Planning.  The paper mentions work on structured latent dynamical systems (Watter  et al. 15), but disregards these \"*However, these approaches relied on locally-linear predictive models, which may be difficult to design.*\" No design is required for latent dynamical systems with local linear latent dynamics (eg. Watter  et al. 15, [Fraccaro et al. 17](https://arxiv.org/pdf/1710.05741.pdf)) - all transition matrices and parameters are learned, using a slightly different ELBO. The benefit of this approach is that it allows for standard trajectory optimisation approaches like iLQR to be applied directly. I would like to see a comparison against trajectory optimisation using a dynamical system with learned locally linear models, which arguably allows for simpler planning and control.\n- Along the lines above, there is a recent body of work looking at imposing more structure in the latent dynamical system to simplify and improve downstream control (eg. embedding for proportionality - [Jaques et al.](https://arxiv.org/abs/2006.01959), koopman embeddings for open loop control with QP [Li et al.](https://openreview.net/forum?id=H1ldzA4tPr) In contrast, this work seems to advocate the opposite approach - ignoring the latent dynamical system learned, and focusing on better methods to solve a more challenging optimisation problem. I  believe that more discussion on the contrasts between these ideas would be a useful addition to this paper.\n- Algorithm 1. The algorithm and training approaches lack clarity and cause some confusion, which needs to be improved. The algorithm seems to indicate that dynamics model learning and planning happen jointly,  which doesn't really make sense - we shouldn't need to re-learn a dynamics model at planning time. Unless the intention was to imply that this is an online learning approach? I assume that this is not the case, as experimental methods seem to indicate that dynamics and reward models are pre-trained, separately from trajectory optimisation using collocation. Please clarify, and ensure that the methodology lines up with what was demonstrated in the experiments section.\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "GiK0Ken4PBK", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "## summary\nThe paper proposes to transpose colloction methods to solve planning problems in a learned latent state space.\nThis can then be used as a replacement for shooting methods in model-based RL, particularly suitable\nfor image-based tasks, where planning in the observation space is impractical.\n\n## pros\n- Basic shooting methods are a primitive planning technique; we should be able to do much better.\n  Using collocation methods in learned latent state spaces makes sense. This paper is one of the first\n  to provide a working realization of this.\n\n## cons\n- The problem is only difficult because of the attempt to learn the task directly from visual inputs.\n  From a practical robotics and planning perspective, the task problems are very dated, e.g., from 30 years ago.\n  In this sense, the tasks are \"straw man\" problems that are uninspiring.\n- Shooting methods provide exploration that the gradient-driven collocation methods do not allow for.\n  The tradeoffs are not as simple as portrayed.\n\n## recommendations\nI currently lean marginally in favor of acceptance, purely on the grounds that transposing collocation\nmethods to latent spaces does havae future potential.  However, the given examples are uninteresting.\n\n## questions\n- How would the results compare to simply using the latent state to estimate a traditional compact state descriptor\n  and then using that with a classical motion planner? For the given example tasks, that seems very feasible.\n- Can planning methods like CHOMP also be realized in the latent space?\n  What are the general constraints or restrictions, if any, on transposing the many known planning methods into \n  the latent space?\n- What is the impact of choosing a time horizon T that is too short or too long?\n- What is stochastic about the dynamics, if anything, for the chosen experimental tasks? \n- What is the action space for the given tasks? What is a-max for the tasks?\n\n\n## feedback\nThe output is a trajectory, not a policy. To make it actionable would require using the optimized\ntrajectories to learn a policy or to use MPC.  This aspect is missing from the paper. Similarly,\nthe exploration issue is avoided (cf sec 6.1).  Thus, overall, the paper is not really solving an\nRL problem.  The title could more directly address the contribution, i.e., motion planning via\nlatent-space collocation.\n\n\"To this, collocation methods\" (sic)\n\nFigure 2: the text refers to a decoder, but this is missing in the figure. \nThe dynamics model is left unlabeled.\n\nIt is worthwhile briefly discussing the broader space of collocation methods, and where your method\nfits within that taxonomy.\n\nSection 5, Constrained optimization: \"balance between the strength of the dynamics constraint.\"\nmissing: \"and the objective\" ?\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "BKUpMJe2Fz", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "#### Summary:\nIn 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.\n\n\n#### Pros:\n- The paper is clearly written and experiments demonstrated improved performance over CEM and gradient descent optimization of actions.\n\n#### Weaknesses:\n- The experiments are limited to sparse-reward tasks, it may be interesting to compare the performance of LatCo and CEM on DeepMind control suite tasks (same as PlaNet), also to see how LatCo performs on dense-reward tasks.\n- It is unclear why collocation should find goals better than CEM or gradient descent for sparse rewards. If the reward function network learns this sparse reward, there is no meaningful gradient towards the goal for an optimization based method.  CEM seems to have a better chance to find the goal due to randomization of actions. If not reward shaping has been used, why is the learned reward by the PlaNet network useful for collocation?\n- Conclusions claims that the approach would be \"removing the need for reward shaping\", however the task is simplified by the oracle agent for training data collection which uses reward shaping. The manual interaction is shifted from reward shaping to training data augmentation. Please clarify.\n\n#### Recommendation: \nThe main concern about the paper is that optimization-based collocation might not be appropriate for the sparse reward case for a method that learns to predict reward for states. Hence experimental results are questionable. The rebuttal should carefully address this issue.\nThe idea is evaluated in a sufficient range of experiments,  although further experiments on standardized benchmarks (DeepMind control suite) would significantly improve the paper. The points raised in weaknesses above should be addressed.\n\n\n#### Questions for rebuttal:\n- See \"weaknesses\".\n- Why not use gradient descent to update the Lagrange multipliers?\n- What is the role of $\\epsilon$ in the Lagrangian in algorithm 1 / l5?\n- How do the terms in the Lagrangian relate to the residual terms? Especially, why does the quadratic action objective in the Lagrangian relate to the residual $\\max(0, |a_t| - a_\\mathrm {max})$?\n- In 6.3, you write \"To provide a fair comparison that isolates the effects of different planning methods, we use the same dynamics model architecture for all agents\". Is it only the same architecture, or the same dynamics model (at least for the models trained only on the oracle data)?\n- What is the task in Sec. 6.4 to generate the plots in Fig. 5?\n- Why do the returns get negative if the reward is sparse and positive ?\n\n#### Further comments:\n- Rename $\\lambda_t$ in eq. 6 to $\\lambda_t^\\mathrm{dyn}$,  to match l5 of algorithm 1\n- What is the value of $\\lambda_t^\\mathrm{act}$?\n- \"For the reward objective, we found it convenient to map the reward to  the negative part of the real line with the softplus operation\"  sounds confusing to me, I associate negative numbers with the  negative part of the real line. Maybe phrase it like  \"For the reward objective, we form residuals by squashing the negated reward through a softplus function\".\n- Algorithm 1: $T_\\mathrm{rep}$ is not defined\n- Algorithm 1 / l13: The ELBO is maximized -> gradient *ascent* (with some learning rate) $\\theta := \\theta + \\alpha \\nabla ...$\n- \\emph{} seems to give underlined instead of italic characters (see the references section), this is probably not intended\n- Please plot lagrange multiplier values in Fig 5\n\n#### Post-rebuttal comments\n- The paper should further elaborate on the smooth reward predictions and how online learning in the sparse reward setting can be possible with LatCo. It seems the method requires a specific initialization/implementation of the reward predictor, for instance, to overestimate rewards so that the method has to explore the areas where reward is overestimated and pull down the predicted reward. The paper should explain how this was implemented. This kind of exploration would be prone to the curse of dimensionality if the state representation of the environment is high-dimensional. The authors should discuss this limitation thoroughly. This might also explain why the tasks in the experiments are limited to 2-dimensional states.\n- I wonder about the discretization of the colors in Fig 8. Higher quantization of color should be provided so gradients of the reward landscape can be assessed.\n- The paper still does not detail the update rule for \\lambda_act\n\nOverall, the author response has addressed some of my technical concerns, but the main challenges are only addressed partially. The paper is still borderline and might need another thorough round of improvement and resubmission to another venue. ", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This work applies collocation, a well known trajectory optimization technique, to the problem of planning in learned visual latent spaces. Evaluations show that collocation-based optimization outperforms shooting via CEM (PlaNet) and  shooting via gradient descent.

Pros:
- I agree with the reviewers that this idea makes sense, and will very likely be built on in future work
- the authors have very actively addressed most comments of all reviewers that engaged in discussion

cons:
- I agree with the reviewers that this is a very simple and straightforward application of collocation methods to the visual latent space domain. Furthermore, the chosen tasks are fairly simplistic, meta-world has a variety of tasks, most of which are more complex than the reaching and pushing task that were chosen for this manuscript.  
- Even with all the updates, the evaluation is still very shallow. I agree with the reviewers that obtaining results for both settings: a) visual MPC with pre-trained (or even ground truth) dynamics model b) in the model-based RL setting, for which the model is being learned, is important. While the authors have added some of these experiments, a detailed discussion of how the results change from a) to b) is missing.  Furthermore, when using collocation in this MBRL setting, how should dynamics constraints be enforced (should they even be enforced when the model is still really bad?). How does the comparison between collocation and shooting fare when you use dense/shaped rewards for the sawyer tasks? Many questions come to mind, some of which that have been raised by the reviewers, and my main point is that simple idea + in-depth analysis of some of these questions would have created a stronger contribution.
- Alternatively, real system experiments would have increased the significance of this work. 
- I don't see any direct references of gradient-based visual latent-space planning (shooting), but related work on this does exist. 

In my opinion, a simple straightforward idea is no reason to reject a paper. However, currently, the reader does not learn when collocation should be considered over other trajectory optimization methods, when attempting to plan in a learned visual latent space. And what some of the main remaining challenges are. Because of this I lean towards recommending reject, and would encourage the authors to deepen their analysis of collocation in visual latent space.
Decision: Reject