PcLast: Discovering Plannable Continuous Latent States

• The proposed method is highly complicated, involving five moving components: a forward dynamics model (FDM), a multi-step inverse dynamics model (IDM), a contrastive representation learner (CL), a hierarchical graph-based planner, and a low-level MPC planner, with several newly introduced hyperparameters for each module. While I see the motivation behind this approach (i.e., we want to first eliminate all the uncontrollable parts, and then impose a metric on another latent representation space), I'm not fully convinced that all of these components are necessary in practice (see the next bullet points for more specific concerns).
• In the Maze/Sawyer domains, the authors only compare PCLast with ACRO. Given the high complexity of the method, I think it is necessary to show that PCLast is at least better than all the "basic" representations that are used as building blocks for PCLast, i.e., IDM, CL, and FDM, to justify the added complexity.
• The performance of ACRO (the best baseline considered in this work) on the exogenous noise MuJoCo environments seems to be worse than the original performance reported in Islam et al. (2023) (e.g., in cheetah_run_expert). Could the authors explain this performance gap? Moreover, given the complexity of PCLast, I believe it requires more thorough empirical comparisons to demonstrate the effectiveness of PCLast (e.g., comparisons on the same set of environments (12 datasets x n distractor configurations) used in ACRO).
• Also, while the only "non-toy" benchmark considered in this work is MuJoCo in Section 4.5, it is unclear how planning helps on MuJoCo tasks (since hierarchical planners seem to be only used in the Maze environments). Given that the main benefit of PCLast is its ability to plan (as the title says), I believe it would be much more convincing if the authors could show this benefit in more practical environments.

[Minor]

• The paper is using the ICLR 2023 format.
• Typo (small $\\\|$) after "we formulate the cost as" in Appendix D.

I believe this is a borderline paper. While I appreciate the motivation behind PCLast, the empirical results are not strong/extensive enough to justify the need for all five components of PCLast.

There are a several experiments for which results are not presented which I think are necessary to clarify the benefits of this approach. In particular, the related work lists HOMER and DRIML as methods which attempt to encode reachability information in the latent space. However, these techniques are only included in the experiments on exogenous information, and not in the experiments on goal-conditioned RL. Similarly, I'd be interested in the results you can obtain using the hierarchical planner with alternative latent space learning techniques besides PCLaSt.

More generally, I'm a bit confused by the relation of PCLaSt to HOMER and DRIML. The related work section argues that these two approaches are different from PCLaSt because they don't deal with exogenous noise. However, in the technical development of the paper, it seems that the denoising effects are due primarily to ACRO, whereas the contribution of PCLaSt is primarily in enforcing state-space geometry.

On a more minor notes, while the presentation is generally clear there is one omission that I found confusing at first. In Section 3.3, the process for generating samples for the contrastive loss omits a description of how $s_t$ and $s_{t+k}$ are generated based on $y$.

1. The first main contribution is a new contrastive loss function. Although I like the justification of the loss function through the model of Brownian motion (Appendix B), there are several previously proposed loss functions that seem to solve the same problem. Specifically, the two time-contrastive loss functions I list below pull embeddings of frames nearby in time together and push embeddings of frames far away in time or from completely different videos apart. In the related work, the authors state that time-contrastive learning “is not invariant to exogenous information”, but neither is the proposed loss function. Invariance to exogenous information is achieved by first training the latent space using a prior method, ACRO.

The loss function proposed in this paper: $- \log \sigma ( \beta - \alpha D(s_t, s_{t+k}) ) - \log (1 - \sigma ( \beta - \alpha D(s_t, s_r) ) )$

Wang et al., 2015: $max(0, D(s_t, s_{t+k}) - D(s_t, s_r) + m)$

Nair et al., 2022: $- \log \frac{\exp S(s_t, s_{t+k}) }{\exp S(s_t, s_{t+k}) + \exp S(s_t, s_{t+k+l}) + \exp S (s_t, s_r)}$

Where $D$ is a distance function and $S$ is a similarity function between a pair of states.

2. The second main contribution is the entire system of first learning an encoder $\phi$ using an inverse dynamics loss function, followed by a contastively-learned encoder $\psi$, followed by a forward model, followed by K-means clustering, followed by a hierarchical planner. I would like to see more justification of this approach. For example, the combination of $\psi(\phi(x))$, where $\phi$, is trained first and frozen, is not fully justified. Could we train a single encoder with multiple losses? Further, the choice of the hierarchical clustering and planning methods could be compared to other hierarchical systems (e.g. Kurutach et al., 2018).

3. (Relatively minor) From the introduction to the methods section, the paper does not flow very well. I think the reason is that the paper lacks focus on its key contribution. If the key contribution was a new and well-justified contrastive learning function, the paper could be written differently. If the key contribution was a system of learning a hierarchical latent space, the paper could also be written differently. This might be just a matter of personal taste, so feel free to disregard.

Minor:

• Should “τ ≫ t ≫ t0” be “τ >= t >= t0”?
• As I stated in my previous review, the proof of Proposition 1 contains the term $P_k(z′|z′, y = 1)$, which I think is a typo since $z’$ appears twice?

References:

Wang et al., 2015: https://arxiv.org/abs/1505.00687

Nair et al., 2022: https://arxiv.org/abs/2203.12601

Kurutach et al., 2018: https://arxiv.org/abs/1807.09341

Post rebuttal: Thank you for answering my questions! Unfortunately, I do not think this paper can be accepted without more comparisons to prior time-contrastive learning losses as well as prior hierarchical learning methods. I highlighted these in my review. Nevertheless, I believe the work addresses an important problem and is eventually going to be a strong submission.

• Throughout the paper, there are several symbols used to denote different levels of latent states. However, each of the symbols $x$, $z$, and $s$ sometimes means different levels of abstraction. It might be easier to follow if each symbol is used to represent a single entity and a summary of these symbols is illustrated as a figure or list. If I didn't understand correctly, the paper writing could be improved to make it straightforward.

• The proposed approach makes sense in 2D/3D navigation domains where reachability can be represented with a very low dimensional (2D/3D) latent space. However, it is unclear whether this also applies to more complicated environments with high-dimensional C-space, e.g., robotic manipulation with many objects. If the proposed method is not scalable but still useful for navigation tasks, the paper needs to explain which family of problems can be handled by the proposed approach and which family of problems cannot.

• The proposed planner assumes that latent states in the same cluster are reachable from each other, which may not be true since the latent state representations are approximated and learned from random exploration trajectories. It might not be a problem in the experiments in this paper since the experiments are done in simple 2D navigation and 3D reaching tasks. However, this may fail when a higher-level plan cannot be achieved because an agent cannot reach a waypoint (the center of a cluster in the plan) from the current state. It is required to discuss this issue and how to resolve it.

• The results in Sawyer-Reach do not demonstrate the advantage of the proposed approach as explained in Section 4.3. Extensive experiments with more challenging environments that can show the benefit of the proposed approach would strengthen the paper.

• The experiments rely on exploratory data to pre-train the representations. This may not be scalable to more complex environments with many narrow passages in C-space. One alternative could be learning the representations using play data or offline data. This can be an interesting experiment to see.