Stabilizing Contrastive RL: Techniques for Robotic Goal Reaching from Offline Data

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.

The entire experimental framework is based on the method "Contrastive learning as goal-conditioned reinforcement learning". This singular focus raises concerns about the generalizability of the conclusions.

One of the aims of this paper is to find design decisions that improve an existing goal-conditioned RL algorithm based on contrastive learning. While we did indeed use this method as a starting point, we evaluate our resulting method by comparing it to 8 prior methods (Sec. 4.3), including those based on different backbone algorithms. Note that we have also run experiments to study the effect of applying some of components (e.g. auxiliary perception loss and sub-goal planning) from prior methods to our approach (Appendix F.11 & Appendix F.12).

Limited Task Support for Major Conclusions

While Fig. 1 in the main text does the ablation on the $`drawer`$ environment, Fig. 14 in the appendix repeats these ablation experiments on the $`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.

In Sec 3.2, it mentions "using larger batches increases not only the number of positive examples (linear in batch size) but also the number of negative examples (quadratic in batch size)". Could you further explain this point?

This point is derived from the original implementation of contrastive RL [6] and we have added a citation in this sentence to clarify. As mentioned in Sec. 3.1, the critic function is parameterized by the inner product between the state-action representation $\phi(s, a)$ and future state representation $\psi(s_{t+})$: $f(s, a, s_{t+}) = \phi(s, a)^{\top} \psi(s_{t+})$. In practice, we first sampled a batch of state-action representations $\\{\phi(s^{(i)}, a^{(i)})\\}^N_{i = 1}$ and the corresponding future state representations $\\{\psi(s_{t+}^{(i)})\\}^N_{i = 1}$, and then compute their $*outer products*$ to construct a critic matrix $F \in \mathbb{R}^{N \times N}$ with $F[i, j] = \phi(s^{(i)}, a^{(i)})^{\top} \psi(s_{t+}^{(j)})$. The diagonal entries in the critic matrix $F$ are positive pairs while the off-diagonal entries in the critic matrix $F$ are negative pairs. Thus, increasing the batch size $N$ results in a linear increase in positive examples and a quadratic increase in negative examples.

[1] Fujimoto, S., Hoof, H. and Meger, D., 2018, July. Addressing function approximation error in actor-critic methods. In International conference on machine learning (pp. 1587-1596). PMLR.

[2] Ball, P.J., Smith, L., Kostrikov, I. and Levine, S., 2023. Efficient online reinforcement learning with offline data. arXiv preprint arXiv:2302.02948.

[3] Kumar, A., Agarwal, R., Geng, X., Tucker, G. and Levine, S., 2022, September. Offline Q-learning on Diverse Multi-Task Data Both Scales And Generalizes. In The Eleventh International Conference on Learning Representations.

[4] Nair, S., Rajeswaran, A., Kumar, V., Finn, C. and Gupta, A., 2023, March. R3M: A Universal Visual Representation for Robot Manipulation. In Conference on Robot Learning (pp. 892-909). PMLR.

[5] Laskin, M., Lee, K., Stooke, A., Pinto, L., Abbeel, P. and Srinivas, A., 2020. Reinforcement learning with augmented data. Advances in neural information processing systems, 33, pp.19884-19895.

[6] Eysenbach, B., Zhang, T., Levine, S. and Salakhutdinov, R.R., 2022. Contrastive learning as goal-conditioned reinforcement learning. Advances in Neural Information Processing Systems, 35, pp.35603-35620.

We thank the reviewers for the detailed responses and for the suggestions for improving the paper. It seems like the main question is about the motivation and structure of experiments. We agree that this paper is an empirical report – it documents a year's worth of experiments aimed at building a better offline algorithm for goal-conditioned RL problems. We believe these findings may be of interest to some of the ICLR community that works on robotics, scaling offline RL, and related topics. We agree that the empirical flavor of these results may not be interesting to all facets of the ICLR community. Together with the discussion below and the revision in the paper (orange texts), does this fully address the reviewer’s questions?

Motivation is not clear

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.

The experiment section, segmented into seven sub-sections, lacks clarity.

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.

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.

Author response

Thanks for clarifying that the rebuttal addresses the concerns about organization and objectives of the experiments, but not the concerns about why contrastive RL needs stabilization and why drawer is a representative task to use in the ablations.

We have further revised the paper to address these two points. We describe the changes below:

evidence about why drawer is a representative task to run ablations.

The ablation experiments are done on two tasks: $`drawer`$ and $`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 ($`drawer`$, which is still challenging for baselines) and one of the most complex tasks ($`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.

Why is there a need to stabilize "contrastive RL" in the first place?

Perhaps “contrastive RL v2” or “improved contrastive RL” would be a more clear name for the method. We called our method “stable contrastive RL” because, when we looked at the learning curves (Fig. 13), they were smoother and had much lower variance across random seeds. We would be happy to update the method / title if the reviewer thinks that would address this potential for confusion.

We thank the reviewer for the responses and suggestions for improving the paper. We have revised the paper based on these suggestions (orange texts) and describe in more detail below.

… some analysis of learned representations of stable contrastive RL in relation to performance combined with visual situations would help to understand those performance comparison analysis.

We agree that analysis of learned representations of stable contrastive RL in addition to simply comparing performance would help understand the difference between our method and baselines. As an initial step in this direction, we have visualized the representations learned by our approach and 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. We welcome additional suggestions for analyzing the performances.

But most of the paper are descriptions for simulated situations. This may blur your aims and contributions.

Thanks for the suggestion. We have added this into the limitation section.

We thank the reviewer for the detailed responses and helpful suggestions for improving the paper. The reviewer raised a number of good questions about the paper, which we have attempted to address below and which we have already incorporated into the paper (orange texts). We believe that these revisions to the paper make it stronger. Together with the discussions below, does this fully address the reviewer’s concerns?

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.

Overall, the paper does not do a good job of explaining why some of the seemingly simple tasks are difficult … 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 $`push block, open drawer`$ includes first pushing the orange block away and then opening the drawer. Similarly, to solve the task `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.

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

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.

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.

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

We thank the reviewers for the responses and suggestions for improving the paper. We have revised the paper based on these questions (orange texts), which we describe in more detail below.

The conclusion that a deeper CNN performs worse than a shallow one, likely because of overfitting, indicates that maybe the benchmark tasks used do not align with the paper's goal of leveraging a vast amount of unlabeled data, like current approaches in computer vision and NLP.

We have revised the introduction to emphasize that our aim is to get goal-conditioned RL methods working on the real-world robotics datasets that we have today. While these datasets are big from a robotics perspective, they are admittedly much smaller than the datasets in CV / NLP (potentially explaining why we still see overfitting).

There are some missing citations … McCandlish et al and Bjorck et al …

We have added these citation to the “Layer normalization” paragraph of Sec. 3.2.

A full grid-search of all the combinations of design decisions would be great to have, but the reviewer understands that this may not be feasible.

We agree this is computationally infeasible, but would be happy to brainstorm and run more targeted experiments to answer directed questions.