Paper: Stabilizing Contrastive RL: Techniques for Robotic Goal Reaching from Offline Data (ICLR-2024)

Reviewer comment (Novelty):
The major motivation of the paper is not clear. The primary objective seems to be stabilizing contrastive RL, but the paper does not explain why there is a need to stabilize "contrastive RL" in the first place or the inherent instability of contrastive RL.

Author response:
Our core focus is on (a) building a better goal-conditioned RL algorithm that is applicable to robotic tasks, (b) making initial attempts to understand the difference between our method and prior approaches, (c) and empirically analyzing why our method works better than baselines. We have revised the introduction to clarify these motivations. 

Prior methods sharing similar goals discussed various design decisions that would benefit the performance of an RL algorithm in both online [1] and offline setting [2][3]. “Stability” turned out to be an important step toward these goals: we applied layer normalization to balance feature magnitude of different tasks (Sec. 4.2), and use data augmentation to prevent overfitting of the GCBC regularization (Appendix Fig. 18). Appendix Fig. 13 shows that our method is more stable (less oscillatory) than the original contrastive RL algorithm. 

In addition to proposing design decisions that improve performance of an existing algorithm, we also included an visualization of representations learned by different methods (Sec. 4.4) and a quantitative analysis of learned Q values in the discussion of the arm matching problem (Sec. 4.5), providing empirical explanations for why stable contrastive RL learned to solve the tasks while baselines failed to do so. These studies potentially motivate future work on learning structural representations for robotic tasks.

Reviewer comment (Presentation):
The structure of the experiment section is messy. The section is segmented into seven sub-sections without clearly explaining how these experiments are related, leaving readers questioning their significance.

Author response:
To clarify the experiments, we have added a list of research questions we are trying to answer using our experiments at the beginning of the experiment section and also have updated headers of subsections of our experiments to refer to those questions. We have also revised this section to explain the relationships between the experiments.

Reviewer comment (Experiments):
Several experiments in the paper do not directly align with the core focus on stabilizing contrastive RL. Questions arise such as: Why introduce a comparison of pre-trained representation? What is the relevance of latent interpolation? How does the arm matching problem fit into the story? What motivates the test on generalization across unseen camera poses and object color?

Author response:
> Why introduce a comparison of pre-trained representation?

Prior work [3][4] has used pre-trained off-the shelf visual encoder for an algorithm. We introduce this comparison to study whether pre-training the representation with an auxiliary objective boosts performance (it doesn't). We have added a sentence in the “Network architecture” paragraph of Sec. 4.2 to clarify this.

> What is the relevance of latent interpolation?

The aim of the latent interpolation experiment (Fig. 6) is to provide some intuition into the learned representations. We believe these visualizations are relevant to the discussion because they demonstrate that our method captures causal relationships in its representation space while baselines failed; this offers one hypothesis for why stable contrastive RL outperforms prior methods in Sec. 4.3. We have added a sentence at the end of Sec. 4.3 to clarify.

> How does the arm matching problem fit into the story?

The arm matching problem is a common failure mode of prior methods; we wanted to analyze this case to understand if and how our method was avoiding this failure case, achieving a higher success rate. The discussions in Sec. 4.5 and Appendix F.6 provide a quantitative analysis of the (normalized) Q values learned by our method and baselines, showing that stable contrastive RL achieves lower error in credit assignments. We believe these experiments offer another hypothesis for why stable contrastive RL outperforms prior methods in Sec. 4.3. We also revised the final paragraph of Sec. 4.3 to clarify this.

> What motivates the test on generalization across unseen camera poses and object color?

We have revised the introduction to clarify that we focus not only on studying design decisions that produce a better goal-conditioned RL algorithm based on contrastive RL, but also on investigating the properties of the policy learned by our method. Prior methods have used contrastive learning as auxiliary loss besides RL objective to improve the robustness of RL algorithm with image inputs [4][5]. This motivates us to study the robustness of the learned policy by testing the generalization of stable contrastive RL on unseen camera poses and object colors.

Reviewer comment (Presentation):
The presentation in section 4.2 further complicates comprehension. While five design decisions are discussed in section 3.2, their individual results are not distinctly showcased in section 4.2, leading to confusion.

Author response:
While five design decisions are discussed in section 3.2, their individual results are not distinctly showcased in section 4.2, leading to confusion.

We grouped those 5 design decisions (Sec. 3.2) into two main classes due to limited pages in Sec. 4.2. We have added labels (D1, D2, etc.) in Sec. 3.2 and Sec. 4.2 for each of the design decisions to clarify.

Reviewer comment (Experiments):
Major conclusions are primarily derived from the `drawer` task experiments (Sec 4.2). Although additional tasks are covered in the appendix, the range of tasks considered is limited.

Author response:
Limited Task Support for Major Conclusions

While Fig. 1 in the main text does the ablation on the $\texttt{drawer}$ environment, Fig. 14 in the appendix repeats these ablation experiments on the $\texttt{push block, open drawer}$ environment, drawing similar conclusions. We have revised the first paragraph of Sec. 4.2 to clarify. Because of compute constraints, it's not feasible to run the ablation on every environment. Nonetheless, we'd be happy to run additional narrow experiments to answer any targeted questions the reviewer may have.

Reviewer comment (Experiments):
The paper does not explain why the ablation study was not conducted on all the benchmarks mentioned in Sec 4.1.

Author response:
> evidence about why `drawer` is a representative task to run ablations.

The ablation experiments are done on two tasks: $\texttt{drawer}$ and $\texttt{push block, open drawer}$, both using image-based observations. Here's why we selected these two tasks:
* While using state-based tasks would have been computationally less expensive, we opted for the image-based versions of the tasks because they more closely mirror the real-world use case (Fig. 5), where the robot has image-based observations.
* Prior methods struggle to solve both these tasks, including (original) contrastive RL (Fig. 2 (Right)). Thus, there was ample room for improvement.
* Out of the tasks in Figure 2, we choose the easiest task ($\texttt{drawer}$, which is still challenging for baselines) and one of the most complex tasks ($\texttt{push block, drawer open}$).

Of course, the most representative task would have been to ablate these decisions on the physical robot, but that would have been extraordinarily expensive. We believe that the two representative tasks used are a good faith effort to evaluate our design decisions. We have added this discussion to Appendix F.4. 

If it would help to address the concern, we would be happy to run an additional ablation experiment on another one of the tasks.

Reviewer comment (Experiments):
The paper demonstrates that contrastive RL is a brittle objective, since small changes in the network architecture and initialization lead to huge changes in agent performance. The paper does not study if the objective could be changed to make it more robust and less prone to overfitting.

Author response:
The paper does not study if the objective could be changed to make it more robust and less prone to overfitting.

We agree that our focus is on one contrastive RL objective. However, our experiments do compare to prior methods that use different goal-conditioned RL objectives, e.g., advantage weighted policy regression (WGCSL [6]). In particular, VIP [7] and R3M [8] also use the idea of contrastive learning to obtain representations from large real-world robotic datasets, potentially more robust and less prone to overfitting. We compared to GCBC variants trained on top of those learned representations in Sec. 4.3 and found that stable contrastive RL still outperformed them on 4 / 5 tasks.

Reviewer comment (Evaluation):
It is unclear why none of the methods can learn, e.g., the “push can” policy in Figure 5. Overall, the paper does not do a good job of explaining why some of the seemingly simple tasks are difficult and why SOTA robot learning methods cannot solve them.

Author response:
… why SOTA robot learning methods cannot solve them.

As an initial step in explaining the failure of baselines, we have visualized the representations learned by our approach and other baselines in Sec. 4.4 and Appendix F.5, showing that stable contrastive RL captures causal relationships in its representation space while baselines (GCBC) failed to do so. Quantitatively, we measure the alignment of the interpolated representations with respect to the ground truth sequences in Appendix F.6, showing that stable contrastive RL achieves lower error than the alternative methods. Additionally, Appendix F.7 contains a comparison of the normalized Q learned by stable contrastive RL to those learned by baselines through an optimal trajectory, suggesting that stable contrastive RL assigns values to different observations correctly over baselines. These experiments provide empirical explanations for why stable contrastive RL outperforms prior methods. We have revised Sec. 4.3 to add references to these empirical analysis.

Reviewer comment (Writing):
The second and third paragraph of the introduction could be more specific. E.g. “various design decisions” is too ambiguous.

Author response:
The second and third paragraph of the introduction could be more specific. E.g. “various design decisions” is too ambiguous.

We have revised this to clarify.

Reviewer comment (Novelty):
“However, to the best of our knowledge these contrastive RL methods have not been applied on real-world robotic systems.” – https://arxiv.org/abs/2306.00958, https://arxiv.org/abs/2203.12601 use contrastive learning with offline RL and are deployed on real-world systems. This claim is only true if we use a very narrow definition of “contrastive RL methods”.

Author response:
“However, to the best of our knowledge these contrastive RL methods have not been applied on real-world robotic systems.” … This claim is only true if we use a very narrow definition of “contrastive RL methods”.

We have removed this sentence in the introduction.

Reviewer comment (Experiments):
Could standard imitation learning approaches that use point clouds (like PerAct https://arxiv.org/abs/2209.05451, RVT https://arxiv.org/abs/2306.14896 and Act3D https://arxiv.org/abs/2306.17817) be able to solve the tasks in Figure 5 with high success rate?

Author response:
Could standard imitation learning approaches that use point clouds (like PerAct https://arxiv.org/abs/2209.05451, RVT https://arxiv.org/abs/2306.14896 and Act3D https://arxiv.org/abs/2306.17817) be able to solve the tasks in Figure 5 with high success rate?

Comparing to imitation learning methods using point clouds as inputs may take some time, as the standard benchmarks we've used don't support point cloud observations. One way of interpreting the results in Fig 3 is that the gap between representation-based and image-based results suggests that there is room for more sophisticated representation learning methods (e.g., based on point clouds) to boost performance. We have added a discussion of point-cloud based methods (including [1][2][3]) to the “Representation Learning in RL” paragraph of the related work section.

Reviewer comment (Evaluation):
Why is it difficult to reach high success rates on the seemingly simple robot tasks in Figure 2 and 5?

Author response:
Why is it difficult to reach high success rates on the seemingly simple robot tasks in Figure 2 and 5?
 
We have discussed the challenges in solving these tasks in Sec. 4.1 and Appendix E. Our experiments use a suite of simulated goal-conditioned control tasks based on prior work [4] (Fig. 2 left). These tasks are challenging because we directly use the 48 x 48 x 3 RGB images as observations and goals and solving most of them requires multi-stage reasoning. For example, the optimal trajectory for $\texttt{push block, open drawer}$ includes first pushing the orange block away and then opening the drawer. Similarly, to solve the task $\texttt{pick \\& place (drawer)}$, the policy needs to pick and place the green block first and then close the drawer. Note that the orders of completing those stages (sub-tasks) cannot change, suggesting that an algorithm achieving high success rates should be able to figure out the causal relationship. For the real-world manipulation and locomotion benchmarks, these tasks are difficult because (a) the offline datasets used for training contain suboptimal trajectories missing the objects or target poses[5], (b) evaluation goals are not revealed during training, (c) and the policy learns to solve different evaluation tasks by learning on one large offline dataset.