Paper: What Matters to You? Towards Visual Representation Alignment for Robot Learning (ICLR-2024)

Reviewer comment (Experiments):
The task complexity may not adequately highlight the significance of human feedback in robot learning.

Author response:
We would like to thank the reviewer for the suggestion. Per your suggestion, we implemented a more complex robot manipulation task (described in detail in the common response and in the Appendix of the updated manuscript). Specifically, we modified the robot manipulation environment described in subsection 5.2 by 1) increasing the manipulation difficulty of the task (more objects of various colors and sizes, robot needs to grasp the object while pushing), 2) introducing a more complex preference (grouping by object size), and 3) adding visual distractors (object color). Please find the detailed description in Appendix Section A.8. We find that in this more complex manipulation and preference-learning task, our method significantly outperforms all the baselines  (including the new baselines as introduced in our common response).

Reviewer comment (Experiments):
While the approach captures human-like preferences, there's no concrete evidence to validate its representation of genuine human preferences in real-world tasks. Enhancing the complexity of preference-based feedback might better align pre-trained visual models with user preferences.

Author response:
In this paper, we focused on learning representations and visual rewards from simulated human feedback, which indeed could be different from real human preference feedback. We agree that learning from real end-user feedback is extremely interesting. However, we focused on simulated human feedback in this investigation for several reasons. First, using simulated human feedback allows us to carefully ablate the size of the preference dataset. Second, it gives us privileged access to the ground truth task reward for direct comparison, which ultimately enables us to more rigorously evaluate our idea. Third, evaluating our method with real user feedback requires the Institutional Review Boards (IRB) approval, which is time consuming and requires detailed evaluation results on simulated humans. Based on your suggestion, we are currently preparing our IRB application given the promising results demonstrated in the paper and we are excited to evaluate our method with real user feedback through IRB approved user studies in our future work.

Reviewer comment (Reproducibility):
The author should elaborate in the main paper on their method of utilizing privileged access to state information for dataset generation.

Author response:
The procedure of generating the preference dataset is described in the second paragraph of Section 4. In summary, we utilize privileged state information of the environment and a designed reward function to rank trajectories for constructing the preference dataset. The detailed information about the privileged state and reward function for the X-magical environment is described in second paragraph of Section 5.1, and the detailed information about the privileged state and reward function for the manipulation environment is described in third paragraph of Section 5.2.

Reviewer comment (Experiments):
While the efficacy of this approach is demonstrated with the X-Magical and ISACCGym robots, it would be beneficial to see its applicability in other robotic embodiment or real-world scenarios.

Author response:
We would like to thank the reviewer for the helpful suggestion. We agree that deploying this approach in more robotic embodiments and real-world scenarios is very interesting. The main goal of this paper is to demonstrate that with an aligned visual representation, reward learning via optimal transport feature matching can generate successful robot behaviors with high sample efficiency, and show strong zero-shot embodiment generalization. Since our work is the first step towards visual representation alignment for robot learning, we focused on robotic experiments in simulation to thoroughly investigate our approach and control for confounds. We highlight in the conclusion that demonstrating RAPL with real human videos, preferences, and robotic hardware experiments is important future work.

Reviewer comment (Experiments):
Zhang et al (2020) is described to be the most similar to the proposed methodology. Why is this not compared to?

Author response:
A1.1 Indeed, this paper is related to our work as it studies representation alignment through reward signals. Specifically, this paper uses reward signals to learn representations which disentangle high versus low reward images and demonstrate that the learned representations outperformed the representations learned through self-supervised learning. However, a core assumption of the approach is a priori access to a task’s correctly specified reward function. Obtaining a correct reward function is precisely the learning goal of our setting: in other words, the learning agent does not know the human’s reward but needs to learn it through our proposed RAPL approach. Thus, Zhang et al (2020) is not a comparable baseline for our approach. We have modified our wording in Section 1 to make it more clear.

Reviewer comment (Experiments):
The main baseline appears to be TCC (Zakka et al., 2022; Kumar et al., 2023). Which of these (Zakka et al., 2022 vs Kumar et al., 2023) is used as the baseline? Furthermore, both Zakka et al., 2022 and Kumar et al., 2023 report results on {long-stick, medium-stick, short-stick, gripper} tasks -- why are these tasks not used to report results, for a completely fair comparison? This would be the most revealing comparison to the baselines.

Author response:
A1.2 The TCC representation baseline in the X-magical environment is from Zakka et al., 2022. Regarding embodiment combinations, we agree that it would be compelling to see all combinations of training and deployment embodiments. In our manuscript, we leveraged all four embodiments, but due to time and space constraints, we only investigated a subset of training and deployment embodiment pairs. Specifically, we chose embodiment pairs based on the hypotheses we were investigating in the paper.  For example, in the cross-embodiment experiment in the X-magical groping task (Section 6), we chose the short-stick as the training embodiment and chose the medium stick as the deployment embodiment to demonstrate that our approach can learn task-relevant representations:  because the short stick agent is so small, it has a harder time keeping the objects grouped together; in-domain results from Section 5 show a success rate of ≈ 60% (see Figure 4). In theory, with a well-specified reward, the task success rate should increase when the medium stick agent does the task, since it is better suited to push objects together. Interestingly, when the short stick visual representation is transferred zero-shot to the medium stick, we see precisely this: RAPL’s task success rate improves by 33% under cross-embodiment transfer (succ. rate ≈ 80%). This finding indicates that RAPL can learn task-relevant features that can guide correct task execution even on a new embodiment.

Reviewer comment (Experiments):
A comparison to other pre-training work is made by comparing to MVP-OT, which is revealing. MVP is trained on Ego4D, so there is a massive distribution shift, perhaps explaining the poor performance. For a fairer comparison, could Masked Visual Pre-training be performed on data from the environments directly (e.g. using the trajectories)?

Author response:
A1.3 We would like to thank the reviewer for the baseline suggestion, and we implemented their idea.  In order to preserve the original representations of MVP while adapting to the task domain, we utilized the low rank adaptation technique (LoRA; Hu et al., 2021) to fine-tune the MVP encoder (VIT as the backbone) on the images from the manipulation environment.. We then utilized the fine-tuned MVP representation to generate the OT-based visual reward. The new results are presented in Section A.7 of the appendix. While fine-tuning does slightly improve the MVP performance, this representation still struggles to achieve > 5% success rate, and our proposed method still significantly outperforms this new baseline (Fine-Tuned-MVP-OT) by approximately 70%. We hypothesize that this is due to the signal used during representation learning: RAPL specifically learns representations that disentangle preferred versus disliked robot behavior videos, while MVP learns to in-paint masked image patches (which does not necessarily provide insight into task-relevant visual features).

Reviewer comment (Experiments):
While the visual backbone models are initialized with ImageNet pre-trained weights, it may be revealing to include this (without any further training) as another baseline.

Author response:
A1.4 We would like to thank the reviewer for the baseline suggestion. We added this as an additional baseline in the manipulation setting. The new results are attached in Section A.7 of the appendix, and we find that our proposed method still significantly outperforms this new baseline (ImageNet-OT) by approximately 70%.

Reviewer comment (Experiments):
Data is simulated. A simulated human model is used, rather than learning from real end-user feedback.

Author response:
A2. We agree that learning from real end-user feedback is extremely interesting. However, we focused on simulated human feedback in this investigation for several reasons. First, using simulated human feedback allows us to carefully ablate the size of the preference dataset. Second, it gives us privileged access to the ground truth task reward for direct comparison, which ultimately enables us to more rigorously evaluate our idea. Third, evaluating our method with real user feedback requires the Institutional Review Boards (IRB) approval, which is time consuming and requires detailed evaluation results on simulated humans. Based on your suggestion, we are currently preparing our IRB application given the promising results demonstrated in the paper and we are excited to evaluate our method with real user feedback through IRB approved user studies in our future work.

Reviewer comment (Evaluation):
Figure 3 qualitatively shows the reward correlation -- quantification of such correlation using e.g. Spearman's correlation coefficient, compared to baselines, would be revealing.

Author response:
A3. We would like to thank the reviewer for this helpful suggestion. For each robot manipulation task, we computed the average Spearman’s correlation coefficient across 100 video trajectories, we found that our learned visual reward shows a stronger correlation to the GT reward compared to baselines. We added the detailed results in Section A.9

|              | Franka Group (Sec. 5.2) | Kuka Group (Sec. A.6) | Franka Clutter (Sec. A.8) |
|--------------|-------------------------|-----------------------|---------------------------|
| **RAPL**     | 0.59                    | 0.47                  | 0.61                      |
| **RLHF**     | 0.38                    | 0.31                  | 0.26                      |
| **MVP-OT**   | -0.1                    | 0.02                  | 0.08                      |
| **Fine-Tuned-MVP-OT** | 0.19            | 0.02                  | 0.11                      |
| **ImageNet-OT** | -0.09                | 0.12                  | -0.02                     |
| **R3M-OT**   | 0.03                    | -0.14                 | -0.17                     |

Table 1: Spearman’s rank correlation coefficient between the GT reward and each learned reward.

Reviewer comment (Evaluation):
How is the proposed approach better than GT for "Avoiding"?

Author response:
A4. Our hypothesis is that the manually-defined GT reward (without extensive reward weight tuning) is not “shaped” as well as the preference-based reward to enable fast RL policy optimization. We would like to stress that the manually-defined GT reward *does* capture all the right features and their relative importance, but from a RL policy learning perspective, we hypothesize that human preferences actually help shape the visual reward in such a way that makes RL optimization easier compared to using the manually-defined reward. However, we would like to note that both the GT reward policy and our visual reward policy eventually converge to a similar final success rate.

Reviewer comment (Evaluation):
There seem to be error bars on many of the plots -- how many trials is this over?

Author response:
A5. We would like to thank the reviewer for reminding us about this missing information. The success rate plots show the mean and standard deviation of five trials with different random seeds. We added this information to the main paper.

Reviewer comment (Experiments):
Experimental setup is promising but the evaluation is done on simpler tasks. With just 150 triplet dataset used for both the tasks to learn preferences the tasks being studied brings up a question about how well this approach works on more complex tasks? For example, the block movement task in x-magical environment only requires agent to avoid the blue box in the environment. This can essentially be treated as a simple obstacle to avoid in the environment which is simple to learn in general. Some example tasks authors could use for evaluation are: block stacking task with constraints where a user prefers stacking blocks based on size i.e. largest block is at lowest followed by smaller blocks irrespective of the color of boxes. This setup tests generalization ability of the policy to different colors and ability to reason about object size based on preferences. My point being preference-based learning needs to be evaluated with tasks that have distractors that are difficult to pick up. It’d be nice if authors can add more experiments on different manipulation tasks with such properties.

Author response:
A1&3. We agree that this is a very interesting suggestion! Per your suggestion, we implemented a more complex robot manipulation task (described in detail in the common response and in the Appendix of the updated manuscript). Specifically, we modified the robot manipulation environment described in subsection 5.2 by 1) increasing the manipulation difficulty of the task (more objects of various colors and sizes, robot needs to grasp the object while pushing), 2) introducing a more complex preference (grouping by object size), and 3) adding visual distractors (object color). Please find the detailed description in Appendix Section A.8. We find that in this more complex manipulation and preference-learning task, our method significantly outperforms all the baselines (including the new baselines as introduced in our common response).

Reviewer comment (Experiments):
One baseline I’d like to see in addition to RAPL is how well does a simple contrastive pretrained visual representation as reward with same optimal transport logic performs in comparison to RAPL. It is unclear how beneficial RAPL’s triplet sampling is for the current set of tasks being considered.

Author response:
A2. We would like to thank the reviewer for the baseline suggestion. We additionally compare our method with OT-based visual reward using R3M representation (Nair et al., 2022), which is a off-the-shelf ResNet-18 encoder pre-trained on the Ego4D data set via a learning objective that combines time contrastive learning and video-language alignment. The results are reported in Appendix Section A.7. This new baseline achieves near 0 success rate. Similar poor results are observed also in the (Haldar et al., 2023b) where an OT-based visual reward with a pre-trained R3M representation model gives near 0 success rate for manipulation tasks.

Reviewer comment (Experiments):
For the cross-embodiment experiments the simple tasks of block movement doesn’t necessarily highlight how well the proposed approach is performing. I’d be interested in seeing similar results in simple object pick and place task similar to one asked in W1.

Author response:
A1&3. We agree that this is a very interesting suggestion! Per your suggestion, we implemented a more complex robot manipulation task (described in detail in the common response and in the Appendix of the updated manuscript). Specifically, we modified the robot manipulation environment described in subsection 5.2 by 1) increasing the manipulation difficulty of the task (more objects of various colors and sizes, robot needs to grasp the object while pushing), 2) introducing a more complex preference (grouping by object size), and 3) adding visual distractors (object color). Please find the detailed description in Appendix Section A.8. We find that in this more complex manipulation and preference-learning task, our method significantly outperforms all the baselines (including the new baselines as introduced in our common response).

Reviewer comment (Experiments):
Human preferences are multi-modal in nature i.e. different people have different preferences. The current setup doesn’t consider multi-modal preferences and show experiments comparing different methods in such a setup which seems like a big flaw in current experimental setup. I’d like authors to consider adding experiments that’d demonstrate effectiveness of RAPL under such setting.

Author response:
A4. Indeed, when learning from multiple people, each user could have a different preference than other users. However, in our paper, we focus on a personalized setting in which the robot aims to align its representation and visual reward with a single end-user, i.e., the user’s preference is uni-modal. We are very excited about the proposed extension where preference feedback is crowdsourced from different people (i.e., multi-modal preference), and we included this in the conclusion and future work.

Reviewer comment (Experiments):
In Figure 6, the RLHF results are quite poor. Have authors tried scaling training for RLHF for the same experiment to figure out how long does it take to reach RAPL/GT performance?

Author response:
A5.  We investigated both RLHF performance as a function of training time and as a function of preference dataset size. First, we compared the true success rate of the RLHF policy (under the true human’s measure of success) vs. the “perceived” policy performance under the RLHF reward. These results are visualized in Figure 18 in Appendix Section A.10. We see that after 350 epochs, the RLHF learned reward perceives the policy as achieving a high reward. However, the true success rate is still near zero. This indicates that the RL optimization is capable of improving reward over time, but it is optimizing a poor reward signal that does not correlate with the true measure of success. Second, we investigated if RLHF performance can be improved by scaling the preference dataset. In the last paragraph of Section 6 and Appendix Section A.5, we doubled the training dataset set for training the RLHF reward and tested the policy training with this reward in the Franka manipulation task (Figure 12 in Appendix Section A.5). We observe that policy performance is slightly improved (indicating that with more feedback data, RLHF could yield to an aligned policy), but our method still outperforms RLHF by 75% with 50% less training data. To build intuition for this, we visualized the attention map in Appendix Section A.4. We can observe that the RLHF embedding pays attention to irrelevant areas that can induce spurious correlations, which may also contribute to the poorer performance of the RLHF visual reward.

Reviewer comment (Experiments):
My major concern is evaluation on simple object movement tasks and no evaluation on tasks with multi-modal preferences. If authors add more experiments that'd address my concerns I’d be happy to update the rating.

Author response:
A1&3. We agree that this is a very interesting suggestion! Per your suggestion, we implemented a more complex robot manipulation task (described in detail in the common response and in the Appendix of the updated manuscript). Specifically, we modified the robot manipulation environment described in subsection 5.2 by 1) increasing the manipulation difficulty of the task (more objects of various colors and sizes, robot needs to grasp the object while pushing), 2) introducing a more complex preference (grouping by object size), and 3) adding visual distractors (object color). Please find the detailed description in Appendix Section A.8. We find that in this more complex manipulation and preference-learning task, our method significantly outperforms all the baselines (including the new baselines as introduced in our common response).

Reviewer comment (Evaluation):
No discussion on what is still hard to do or not reliable.

Author response:
We have expanded our conclusion section to reflect the limitations of our method. One limitation of the method is that we assume the user is rational, which motivates the usage of the Bradley-Terry model for interpreting the preference feedback. Although such an assumption is widely used in the preference learning community, humans can exhibit irrationality, systematic bias, and suboptimality. Future study can focus on the open problem of robustly aligning visual representations under such forms of feedback. Moreover, although our method is much more sample efficient compared to the vanilla RLHF, it still passively learns from the human feedback which makes it still inefficient for product deployment. Future work could investigate the challenging problem of actively querying humans for faster aligning visual representations.

Reviewer comment (Evaluation):
Evaluation based on difficulty level of the configuration, shape, number of objects needed to be cleared would help toward answering the previous question.

Author response:
We would like to thank the reviewer for your suggestion! Per your suggestion, we implemented a more complex robot manipulation task (described in detail in the common response and in the Appendix of the updated manuscript). Specifically, we modified the robot manipulation environment described in subsection 5.2 by 1) increasing the manipulation difficulty of the task (more objects of various colors and sizes, robot needs to grasp the object while pushing), 2) introducing a more complex preference (grouping by object size), and 3) adding visual distractors (object color). Please find the detailed description in Appendix Section A.8. We find that in this more complex manipulation and preference-learning task, our method significantly outperforms all the baselines  (including the new baselines as introduced in our common response).

Reviewer comment (Experiments):
Another is to compare with a second, very different, task.

Author response:
We would like to thank the reviewer for your suggestion! Per your suggestion, we implemented a more complex robot manipulation task (described in detail in the common response and in the Appendix of the updated manuscript). Specifically, we modified the robot manipulation environment described in subsection 5.2 by 1) increasing the manipulation difficulty of the task (more objects of various colors and sizes, robot needs to grasp the object while pushing), 2) introducing a more complex preference (grouping by object size), and 3) adding visual distractors (object color). Please find the detailed description in Appendix Section A.8. We find that in this more complex manipulation and preference-learning task, our method significantly outperforms all the baselines  (including the new baselines as introduced in our common response).

Reviewer comment (Reproducibility):
How does the visual attention map (figure 11, supplementary material) computed?

Author response:
We added a detailed explanation and graphical visualization of how we compute the attention map in Appendix Section A.4. In short, we visualize the attention map by leveraging linear permutation and kernel inflation.