Venue: ICLR 2021 Conference

Reviews:

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


#### Pros:
- The paper is clearly written and experiments demonstrated improved performance over CEM and gradient descent optimization of actions.

#### Weaknesses:
- 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.
- 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?
- 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.

#### Recommendation: 
The 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.
The 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.


#### Questions for rebuttal:
- See "weaknesses".
- Why not use gradient descent to update the Lagrange multipliers?
- What is the role of $\epsilon$ in the Lagrangian in algorithm 1 / l5?
- 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})$?
- 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)?
- What is the task in Sec. 6.4 to generate the plots in Fig. 5?
- Why do the returns get negative if the reward is sparse and positive ?

#### Further comments:
- Rename $\lambda_t$ in eq. 6 to $\lambda_t^\mathrm{dyn}$,  to match l5 of algorithm 1
- What is the value of $\lambda_t^\mathrm{act}$?
- "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".
- Algorithm 1: $T_\mathrm{rep}$ is not defined
- Algorithm 1 / l13: The ELBO is maximized -> gradient *ascent* (with some learning rate) $\theta := \theta + \alpha \nabla ...$
- \emph{} seems to give underlined instead of italic characters (see the references section), this is probably not intended
- Please plot lagrange multiplier values in Fig 5

#### Post-rebuttal comments
- 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.
- 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.
- The paper still does not detail the update rule for \lambda_act

Overall, 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. 
review_title: Review: Model-Based Reinforcement Learning via Latent-Space Collocation

final_score: rating: 6: Marginally above acceptance threshold
review_content: ## summary
The paper proposes to transpose colloction methods to solve planning problems in a learned latent state space.
This can then be used as a replacement for shooting methods in model-based RL, particularly suitable
for image-based tasks, where planning in the observation space is impractical.

## pros
- Basic shooting methods are a primitive planning technique; we should be able to do much better.
  Using collocation methods in learned latent state spaces makes sense. This paper is one of the first
  to provide a working realization of this.

## cons
- The problem is only difficult because of the attempt to learn the task directly from visual inputs.
  From a practical robotics and planning perspective, the task problems are very dated, e.g., from 30 years ago.
  In this sense, the tasks are "straw man" problems that are uninspiring.
- Shooting methods provide exploration that the gradient-driven collocation methods do not allow for.
  The tradeoffs are not as simple as portrayed.

## recommendations
I currently lean marginally in favor of acceptance, purely on the grounds that transposing collocation
methods to latent spaces does havae future potential.  However, the given examples are uninteresting.

## questions
- How would the results compare to simply using the latent state to estimate a traditional compact state descriptor
  and then using that with a classical motion planner? For the given example tasks, that seems very feasible.
- Can planning methods like CHOMP also be realized in the latent space?
  What are the general constraints or restrictions, if any, on transposing the many known planning methods into 
  the latent space?
- What is the impact of choosing a time horizon T that is too short or too long?
- What is stochastic about the dynamics, if anything, for the chosen experimental tasks? 
- What is the action space for the given tasks? What is a-max for the tasks?


## feedback
The output is a trajectory, not a policy. To make it actionable would require using the optimized
trajectories to learn a policy or to use MPC.  This aspect is missing from the paper. Similarly,
the exploration issue is avoided (cf sec 6.1).  Thus, overall, the paper is not really solving an
RL problem.  The title could more directly address the contribution, i.e., motion planning via
latent-space collocation.

"To this, collocation methods" (sic)

Figure 2: the text refers to a decoder, but this is missing in the figure. 
The dynamics model is left unlabeled.

It is worthwhile briefly discussing the broader space of collocation methods, and where your method
fits within that taxonomy.

Section 5, Constrained optimization: "balance between the strength of the dynamics constraint."
missing: "and the objective" ?

review_title: useful idea;  uninteresting examples

final_score: rating: 6: Marginally above acceptance threshold
review_content: ## Paper summary

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.

I 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 
- better clarify the algorithm to align with the methods used in experiments
- 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)

### Pros

- The paper is well written and clearly laid out.
- Solving a collocation problem in the latent space is a sensible approach, and a much better idea than using CEM planning or shooting. 

### Cons
- 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.
- 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. 

### General recommendations for improvement and queries 

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

review_title: Sensible idea, some concerns about method clarity 

final_score: rating: 7: Good paper, accept
review_content: 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.

Strengths:

i) The motivation, organization and the overall writing of the paper are clear.

ii) The tested experimental domains are good representatives of the realistic planning setting identified in the paper.

Weaknesses:

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

ii) Have you tried solving Eq. 2. directly similar to [4]? It seems more appropriate baseline compared to c) (i.e., as labeled above).

iii) How do you reason about the length of the horizon T? For example [1,2] use heuristic search.

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

References:

[1] Classical Planning in Deep Latent Space: Bridging the Subsymbolic-Symbolic Boundary, Asai and Fukunaga AAAI-18.

[2] Learning Neural-Symbolic Descriptive Planning Models via Cube-Space Priors: The Voyage Home (to STRIPS), Asai and Muise IJCAI-20.

[3] Nonlinear Hybrid Planning with Deep Net Learned Transition Models and Mixed-Integer Linear Programming, Say et al., IJCAI-17.

[4] Scalable Planning with Deep Neural Network Learned Transition Models, Wu et al. JAIR.

[5] Optimal Control Via Neural Networks: A Convex Approach, Chen et al., ICLR 2019.

** Post Rebuttal **

To 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.
review_title: Review of LatCo

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