Contrastive Representations Make Planning Easy

Dear Reviewer,

We thank the reviewer for the detailed feedback. It seems like the main suggestion is to add more experiments, which we have done through new experiments on higher-dimensional tasks (up to 46-dimensional) and by demonstrating that the inferred representations are useful for control. Together with the revisions (red text in new PDF) and clarifications discussed below, does this fully address the reviewer's concerns about the paper? We look forward to continuing the discussion!

demonstrate that their method works on a real problem … robot turning a knob to open a door

As suggested, we have run additional experiments to validate our theory, including on a 39-dim robot door opening task:

• Fig 3 (left): We show how the analysis in Sec. 4 can be used to solve a maze. After learning representations, we used the closed form expression to infer the representations of intermediate states. This result shows how the representations effectively warp the state space so that linear motion in representations (i.e., $A \psi$) corresponds to fairly complex motion in the state space (i.e., navigating through a maze).
• Fig 3 (right): Using this same maze, we show that the inferred waypoints can be used for control (details below).
• Fig 4: We apply the representations to a 39-dimensional robotic door opening task. Compared to a PCA baseline, linearly interpolated waypoints are a better fit for the true data waypoints.
• Fig 6: We apply the representations to a 46-dimensional robotic hammering task. Compared to a PCA baseline, linearly interpolated waypoints are a better fit for the true data waypoints.

The term "planning" has a technical meaning in artificial intelligence. That meaning is very different from temporal inter- or extrapolation, which is what it refers to in the paper.

Thanks for pointing this out! We agree that there is no single clear meaning of this word and it is often misused. We have revised the introduction to explain that we will use ``planning’’ to refer to an inference problem, rather than an optimal control problem. We provide citations there to some prior work that takes a similar perspective (Botvinick 2012, Attias 2003), though we acknowledge that this may be a nonstandard usage within the AI community.

"representation", which is used as a map from observations to a latent space, and also for an element of that space

Thanks for raising this great point! It indeed seems like a potential source of confusion. To incorporate this feedback, we have significantly revised Sec. 4.1 to disambiguate “encoder” from “representation.” We have also reviewed every use of the word “representation” in the rest of the main text, making some minor changes in attempts to clarify this. Are there remaining parts of the paper where this is unclear? We would be happy to make additional changes to improve clarity and precision.

Abbreviations: CV, NLP

We have written these terms out in full the first time they are used.

Figure 1 is pretty meaningless.

As suggested, we have removed Figure 1.

The first paragraph of section 2 ends in the middle of a sentence.

We have completed the sentence in the revision.

I am surprised to see that the title of section 4 coincides with the paper title.

Our intention here was to make it easy for readers to identify where to find the main results. We also felt like it was an accurate description of the section. If the current title is unclear or imprecise, we would be happy to revise the section title.

At the point figure 2 is presented (section 1), the symbols in the figure are not introduced yet.

We have revised the caption to briefly explain the symbols

Why on earth should one gray out a proof?

We have removed this. Our original intention was to indicate that readers could skip over this section if they were not interested in the proof, but we realize in hindsight it can cause confusion.

Dear Reviewer,

Thank you for the detailed feedback. It seems like the main suggestion is to add more experiments, which we have done through new experiments on higher-dimensional tasks (up to 46-dimensional) and by demonstrating that the inferred representations are useful for control. Together with the revisions (red text in new PDF) and clarifications discussed below, does this fully address the reviewer's concerns about the paper? We look forward to continuing the discussion!

The inclusion of empirical evidence would significantly strengthen the case for acceptance.

As suggested, we have added 4 new figures with empirical results to validate our theory:

• Fig 3 (left): We show how the analysis in Sec. 4 can be used to solve a maze. After learning representations, we used the closed form expression to infer the representations of intermediate states. This result shows how the representations effectively warp the state space so that linear motion in representations (i.e., $A \psi$) corresponds to fairly complex motion in the state space (i.e., navigating through a maze).
• Fig 3 (right): Using this same maze, we show that the inferred waypoints can be used for control (details below).
• Fig 4: We apply the representations to a 39-dimensional robotic door opening task. Compared to a PCA baseline, linearly interpolated waypoints are a better fit for the true data waypoints.
• Fig 6: We apply the representations to a 46-dimensional robotic hammering task. Compared to a PCA baseline, linearly interpolated waypoints are a better fit for the true data waypoints.

Evidence that planning by waypoint inference enables long-horizon control:

To show waypoint inference is useful for control, we have added an experiment in a 2d maze environment (Fig 3). We define a simple proportional controller for navigating between nearby states; as expected (blue line in Fig 3), this proportional controller fails to reach distant goals. Using the proposed inference over representations, we infer waypoint representations, retrieve waypoints states (nearest neighbor), and iteratively have the proportional controller navigate to the waypoint states. The resulting success rates (orange line in Fig 3) are up to 4.5x higher (19% -> 87%).

"Our analysis will also look…": information resides in the Appendix

We have added another half paragraph here to explain the intuition behind the contents in the Appendix. We welcome additional suggestions for details that should be moved to the main text.

Is the primary assertion in Section 3, which states "learned representations follow isotropic Gaussian distributions," a novel contribution?

This assertion is a very minor extension of prior work, which is why we've placed it in the preliminaries section. We've added a note to explain this at the end of Section 3.

For instance, the reference to "supported by analysis based on Wang and Isola" is somewhat ambiguous. Does this imply that such analysis has been previously conducted in their work?

We have revised this sentence to clarify the relationship with prior work. Prior work provides some theoretical intuition for why Assumption 2 should hold. For clarity, the Appendix explains why this intuition from prior work also holds for Gaussian distributions, rather than von Mises Fisher distributions.

Section 3: it would be advantageous to formalize the statements in this section within theorem/lemma/proposition frameworks, similar to the structure used in Lemma 4.1.

We have reformatted this section to typeset the Assumptions in the same manner as the Lemmas in the paper. Our motivation for not phrasing Assumption 1 and Assumption 2 as "assumptions," rather than theorems, is that we want to be transparent about the assumptions behind our method. While it is true that both these assumptions are guaranteed to hold in certain settings, real world settings may violate these assumptions (e.g., because of sampling noise, function approximation error). We welcome other suggestions for making this section as precise as possible.

In (4), there is no $c / (c+1)$.

We have fixed this typo.

Is the current use of the term "planning" widely accepted within the community?

We have added a sentence to the introduction to explain that we will use planning to refer to an inference problem, rather than an optimal control problem. We provide citations there to prior work that takes a similar perspective.

Dear Reviewer,

We thank the reviewer for the detailed feedback. It seems like the reviewer's main concern is about the significance of the work. We have revised the paper to clarify the significance (start of Sec. 4), and provide more details below. As suggested, we have also added a new experiment to demonstrate how the sampling states (based on our representations) can be used for planning (Fig 4). Do these revisions address the reviewer's concerns about significance? We look forward to continuing the discussion.

how effective the sampled states can be used in planning, which is not studied in this paper.

We ran an additional experiment to study how the learned representations might be used for control (new Fig. 3). These results show that the inferred representations can boost the performance of a simple proportional controller by up to 4.5x (for reaching the most distant goals).

there is no new contribution in terms of proof techniques.

We agree that our paper does not contribute new analysis of Gaussian distributions; rather, it combines (known) results about Gaussian distributions with (known) results about contrastive learning (that it learns a probability ratio) to produce a result that is not known (to the best of our knowledge): that representations learned by temporal contrastive learning are distributed according to a Gaussian Markov chain. We have revised the start of Sec. 4 to mention these points.

Our empirical results demonstrate some potential use cases of this result, including analyzing 46-dimensional time series data (new Fig 6) and solving control problems (new Fig 3).

calculation of the conditional probability given certain states’ representations is not a difficult task

We agree that calculating conditional probabilities of Gaussian Markov chains is easy. We contend that calculating conditional probabilities of arbitrary time series data is very challenging (see, e.g., Fig 4). The contribution of this paper is to show how to learn representations of time series data to reduce this challenging problem to an easy problem. We have revised the start of Sec. 4 to mention these points.

For example, consider the problem of inferring a sequence of states (i.e., a plan) between an initial state and a final state. Such planning problems typically require solving a combinatorial optimization problem (e.g., Williams '17, Fang '23). Even among prior learning-based methods that acquire compact representations, it can be unclear how to efficiently perform planning over the space of representations. In contrast, our work shows how to learn representations s.t. the planning problem becomes easy, reducing it to a simple matrix inversion (Sec. 4.4), avoiding the need for combinatorial optimization.

Comparison of PILCO and Gaussian processes for data-efficient learning in robotics and control

We have added a discussion of these related works to the revised paper. These methods are similar to ours in that they aim to build a probabilistic model of time series data. Unlike GPs, our method avoids having to invert a matrix whose size depends on the number of data points. A second difference is the aims of the papers: while those prior papers aim to learn reward-maximizing policies, our aim is to infer likelihood-maximizing sequences of observation representations. It's worth noting that the representation method we use (contrastive learning) has already seen widespread adoption in audio [e.g., Oord '18], computer vision [e.g., Sermanet '18], NLP [e.g., Mikolov '13], and reinforcement learning [e.g., Laskin '20], so providing a theoretically-grounded way of using these representations for inference may open new avenues for research.