Paper: PcLast: Discovering Plannable Continuous Latent States (ICLR-2024)

Reviewer comment (Experiments):
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).

Author response:
> “... highly complicated, involving five moving components …”

Yes, it may be possible to eliminate some components under certain conditions. But, in principle, each component of our method plays its own critical role for real-world scenarios i.e. multi-step inverse dynamics helps in learning exogenous noise free representation, forward dynamics helps with planning during inference , temporal contrastive representation helps with retaining metric space, hierarchical planner reduces plan computation time and robustifies it, and low-level MPC planner helps in navigating to a goal state. Please note, hierarchical planner and low-level MPC are non-learning components of our approach. They are required as we are in a reward-free setting and don’t learn any policy. 
 
In the future, we can consider potential approaches of merging these components into a single module. Also, if there is no exogenous noise, one can probably ignore inverse-dynamics.

Reviewer comment (Experiments):
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.

Author response:
>  “… PCLast is at least better than all the basic representations… ”.

In ACRO, authors have shown that their method works better than IDM and FDM. This is the reason that we focus on simply comparing it with ACRO and show limitations of ACRO.

Reviewer comment (Experiments):
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).

Author response:
> “  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? “

In this paper, we evaluated on the hardest exogenous noise setting from the ACRO paper, which is called “HARD-EXO-Background Agent”, these are in Figure 20 of that paper’s appendix.  Our ACRO baseline is very similar to the numbers reported in that paper.

Reviewer comment (Experiments):
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.

Author response:
>   “... how planning helps on MuJoCo tasks… ”

When it comes to MuJoCo tasks, we simply focus on the method’s ability to learning a better representation than baselines. This empirically shows improvement in policy learning. This is also an example, where not all the components of our method are required to be used since the considered mujoco-tasks are not goal-conditioned.

Reviewer comment (Presentation):
The paper is using the ICLR 2023 format.

Author response:
>   “...The paper is using the ICLR 2023 format…”
We have fixed this.

Reviewer comment (Writing):
What does LNS stand for in Appendix A?

Author response:
>   What does LNS stand for in Appendix A?

LNS refers to “Local Neighbourhood Structure” which is enforced by the contrastive loss in Section 3.3. We have updated the text to be more clear about it.

Reviewer comment (Experiments):
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).

Author response:
> Is it possible to train a single state encoder using both the inverse dynamics and contrastive learning losses?
 
There are known counterexamples where mixing contrastive temporal and multi-step inverse leads to contamination of the latent state with exogenous noise.  Even a small weighting on contrastive temporal in a mixed loss create this contamination effect.

Reviewer comment (Writing):
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.

Author response:
With “x”, we simply imply observations from the environment; “z” refers to the learnt latent state from the multi-step inverse-dynamics model, and “s” refers to the true-latent state. Though we attempted to be consistent in our usage of these symbols, we will revise our document to fix any inconsistencies and re-emphasize them in our “Section 3.1 Notation and PRELIMINARIES”

Reviewer comment (Experiments):
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.

Author response:
In principle, our work is not restricted to the dimensionality of the true state. We learn our latent state representation from high-dimensional image inputs and keep the dimensionality of the latent state in the order of 128 or 256 for our experiments. Our work rather emphasizes learning latent states for any high-dimensional environment which requires goal conditioned planning. Our experiments on half-cheetah and walker-2d suggest our learning method seems to be helping.

Reviewer comment (Writing):
In Section 3, $x$ is often used for a variable in both $\phi(x)$ and $\psi(x)$. But, since $x$ is used to describe an observation in the section, it might be better denoted as $\psi(s)$. Similarly in Section 3.2, $z$ is introduced for a latent state but it could be simply $s$?

Author response:
We have made the suggested change in the later half of section 3.1. For section 3.2, we continue to use “z” to maintain continuity with the rest of the document.

Reviewer comment (Writing):
In Section 3.2, $f_\text{AC}$ is defined to take two latent states but the following sentence says it takes a concatenation of $\tau + 1$ latent states, which is confusing.

Author response:
Yes, it takes concatenation of the z_t, z_{t+k} and embedding of “k”. We have reworded this in section 3.2 to make it more clear. Thanks for pointing it out.

Reviewer comment (Presentation):
Figure 2 is useful to understand the holistic view of the proposed method. However, it is a bit overwhelming with many symbols. It could be easier to understand this figure if it includes a few image observations for $x$s and uses different shapes or colors for different symbols.

Author response:
Thanks for the suggestion, we will update the figure to have different symbols and colors for different modules.