Paper: Joint Representations for Reinforcement Learning with Multiple Sensors (ICLR-2024)

Reviewer comment (Novelty):
Firstly, if the joint representations (adding proprio and images) improves performance over learning an image-only or proprio-only representation, I do not find this surprising. It makes sense that adding more informations improves the performance.

Author response:
The key insight and contribution of our work is that the effectiveness of additional information depends on the way that it is used. Both the “Concat” and “Joint” agents receive the same information, but the joint representation matches concatenation on simple tasks, and outperforms it as task complexity increases. Similarly, we show that the combination of commonly used individual reconstruction and contrastive losses is crucial for performance. This combination is novel, and we believe the insights that we present to be helpful for many of our peers. Our combined methods outperform concatenation in model-free RL and improve upon SOTA model-based RL methods for tasks with natural video backgrounds. Additionally, they can improve performance in novel tasks with complex observations, such as the occlusion and locomotion tasks, where single modality methods fail.

Reviewer comment (Novelty):
Secondly, adding different losses for each modality, contrastive for images and reconstruction for proprioception, as the central contribution of this work is weak. What can however make this paper a stronger contribution is pursuing other sensor modalities (depth images, surface normals, segmentations, etc) and then exploring various appropriate losses there. As it stands, the paper is only applying a typical reconstruction-based loss to the proprio and contrastive to the image, which are the current norms in the field - no exciting surprise!

Author response:
The key insight and contribution of our work is that the effectiveness of additional information depends on the way that it is used. Both the “Concat” and “Joint” agents receive the same information, but the joint representation matches concatenation on simple tasks, and outperforms it as task complexity increases. Similarly, we show that the combination of commonly used individual reconstruction and contrastive losses is crucial for performance. This combination is novel, and we believe the insights that we present to be helpful for many of our peers. Our combined methods outperform concatenation in model-free RL and improve upon SOTA model-based RL methods for tasks with natural video backgrounds. Additionally, they can improve performance in novel tasks with complex observations, such as the occlusion and locomotion tasks, where single modality methods fail. 

We thank the reviewer for suggesting the use of other sensor modalities. We agree that experiments with different modalities would strengthen our contribution and currently run experiments using Depth Images and Proprioception for the OpenCabientDrawer Task. We will add these results to the final version of our manuscript.

Reviewer comment (Experiments):
Even though its nice the huge amount of analysis performed in this work, the concluding story is very hard to digest, specially since there are many different acronyms to their proposed method such as `Joint(R+R), Joint(CPC+R), Joint(CV+R)`. In addition, some baselines are missing for some tasks (e.g. `Figure 5`) and it makes drawing a final conclusion hard.

Author response:
We thank the reviewer for pointing this out. In the revision, we ensured that all acronyms were properly introduced in the “Representation Learning Methods” paragraph in Section 4 before presenting and discussing our results. We changed the reward curves to bar plots in the main part of the paper which allows a more comprehensive overview of all our results and added DrQv2, with and without proprioception, as a baseline for the OpenCabientDrawer Task (Figure 5). We believe that these changes make our findings and conclusion easier to digest, and would appreciate additional feedback if issues with them remain.

Reviewer comment (Novelty):
I also disagree with the following statement made in the paper `While many self-supervised representation ... neglect other available information, such as robot proprioception`. Adding proprioceptive state observations along with images is not novel in the field (especially in robotics).

Author response:
We agree that adding proprioceptive states to image observations is common in reinforcement learning with images, especially in robotics. However, existing approaches usually learn no explicit representations or concatenate an image representation with proprioception. Our work shows that using the proprioception to learn the representation in the first place, which is much less common, is usually beneficial. We clarified this distinction in the abstract.

Reviewer comment (Experiments):
Is there a reason many of the model-based baselines are missing for the `OpenCabinetDrawer` task?

Author response:
We omitted experiments with model-based approaches on the Locomotion Suite and OpenCabinetDrawer task because of the superior performance of their model-free counterparts in the previous experiments. We added DRQv2 with and without proprioception as a baseline for OpenCabinetDrawer in the revision.

Reviewer comment (Novelty):
The extension of RSSMs to model both image observations and proprioception is a straightforward one, which is the primary contribution of this work.

Author response:
Extending the RSSM to use multiple modalities is only a small part of our contribution. We think providing a unified framework to combine different loss functions for different modalities is more important and such a combination has not been considered before. Our experiments demonstrate that this combination is the main factor behind the benefits of our approach. This finding is particularly relevant given that practitioners typically employ concatenation or single modality methods, presumably because simply adding proprioception without careful consideration often fails to yield significant benefits. We thus agree with the reviewer that these insights are beneficial for practitioners in the field.

Reviewer comment (Evaluation):
It seems like the correct loss for each modality varies quite a bit across domains.

Author response:
We believe that the differences between reconstruction-based and contrastive image losses over different tasks and domains are expected and well explainable by intuition and theory: As reconstruction enforces the model to capture all information in the observation, it yields good representations if all information is relevant. If that is not the case, reconstruction focuses too much on irrelevant aspects that can be ignored by contrastive approaches.  We note that exploring the optimal choice of contrastive loss, though intriguing, is not the primary goal of our current research. That said, we believe that this is an interesting direction for future research for example by combining both paradigms, as also suggested by Reviewer ZjLm.

Reviewer comment (Experiments):
Why do you think there is such a gap between Joint and Concat, where Concat performs the same as ProprioSAC on the cabinet tasks? It seems like Concat should be able to produce any representation that Joint can. Does Concat eventually converge to the same performance as Joint in Fig. 5 if we let it train for longer?

Author response:
While both Joint and Concat receive the same information and thus should indeed be able to produce the same representations, we believe that including the proprioception in the representation learning guides the representation learning toward more relevant aspects of the image. The qualitative results in Figure 6 aim to provide some evidence and intuition for this. These results indicate that including the proprioception in representation learning helps the model to both ignore irrelevant aspects of the image and better capture the relevant aspects. We were unfortunately so far unable to run the OpenCabinetDrawer experiments for longer and thus cannot decisively answer the question regarding Fig 5.  Even if the concat baseline would eventually catch up, there are still large sample complexity benefits. Furthermore, we saw in several other experiments (Occlusions and Locomotion) how the concat baselines converge to policies that achieve less reward than the Joint Representation approach.

Reviewer comment (Experiments):
Do the DenoisedMDP and DreamerPro comparisons utilize observations from both modalities?

Author response:
Except for DrQv2 (I+P), the baselines use image modalities, which they were originally designed for. We use the code bases provided by the respective authors, which do not allow the use of multiple modalities. The baselines are meant to establish the difficulty of the new tasks and show that using joint representations can provide SOTA results. While the baselines can likely be improved with joint representations, such an improvement would further highlight the benefits of joint representation and thus would be in the spirit of our main message.

Reviewer comment (Experiments):
In Fig. 6 (right), are the images reconstructed from a separately trained decoder as a way to probe the representations?

Author response:
Yes. The decoder is trained entirely post-hoc to probe the representations and generate the plots. No information flows back from it to the original model. We clarified this in the revision.

Reviewer comment (Reproducibility):
How are the losses for different observation modalities weighted against each other?

Author response:
Except for the weighted KL for the CPC-Objective in Equation 4, there are no explicit weighting factors in the objectives. For the OpenCabinetDrawer task, we implicitly weigh the proprioception reconstruction by assuming a lower variance for the respective decoder. See Appendix B.2., for details.

Reviewer comment (Presentation):
While the overall logic is clear and smooth, some specific notations and figures are confusing. From the figure plotted in the main paper or appendix, it's very hard to draw any conclusions about which representation learning method is the best. Instead of presenting the curves, drawing some bar charts about the final performance average over different settings and environments can be more straightforward.

Author response:
We thank the reviewer for their time and effort, and particularly for acknowledging our “solid and comprehensive” experiments. We appreciate them pointing out the typos and ambiguities in our initial submission and address them in our revision. We acknowledge the presentation issues (c.f. Weaknesses 1.) in the original manuscript. Following the reviewer’s suggestion, we introduce bar charts for all experiments in the revision and avoid discussing multiple task suits in a single figure. To summarize our experiments, we conduct model-free and model-based experiments for the standard DMC tasks as well as those with modified images (video background and occlusion). We focus on model-free approaches for the Locomotion Suite and OpenCabinetDrawer as they showed better performance than their model-based counterparts in the DMC tasks with all types of images. We included DRQ-v2  as a baseline for OpenCabinetDrawer in the revision to also provide the compressions to SOTA related work for this task.

Reviewer comment (Presentation):
The figures in the main paper (Figure 2, 3, 4) are interleaving with model-free/model-based and occlusion/locomotion, which makes it hard to understand what's been delivered.

Author response:
We thank the reviewer for their time and effort, and particularly for acknowledging our “solid and comprehensive” experiments. We appreciate them pointing out the typos and ambiguities in our initial submission and address them in our revision. We acknowledge the presentation issues (c.f. Weaknesses 1.) in the original manuscript. Following the reviewer’s suggestion, we introduce bar charts for all experiments in the revision and avoid discussing multiple task suits in a single figure. To summarize our experiments, we conduct model-free and model-based experiments for the standard DMC tasks as well as those with modified images (video background and occlusion). We focus on model-free approaches for the Locomotion Suite and OpenCabinetDrawer as they showed better performance than their model-based counterparts in the DMC tasks with all types of images. We included DRQ-v2  as a baseline for OpenCabinetDrawer in the revision to also provide the compressions to SOTA related work for this task.

Reviewer comment (Writing):
There is no explanation of e.g., Joint(CV+R), Joint(CPC+R) which make the confusion that the "+R" is for reward reconstruction.

Author response:
Yes. For all model-free experiments, we follow the findings of Tomar et al (2023), which show reward reconstruction for representation learning helps, even if the policy learning can work ground truth rewards.

Reviewer comment (Experiments):
The results are not clear or convincing enough to draw a strong conclusion as in the discussion section. The environmental and experimental design is not delivered clearly. Please address Questions 1. a for clarification.

Author response:
> The environmental and experimental design is not delivered clearly. Please address Questions 1. a for clarification. (a) Is the "standard image" /"video background" / "occlusion" suite all modified from environments in Table 1?

All three suites contain the 7 environments listed in Table 1 and differ only in the image observations. We clarified this in the revision.

> Are the curves averaged in each suite? How many seeds do you run for each task? The curves in each figure have low variance, which doesn't look like an average as different tasks require very different sample complexities in each suite. If you normalized that, how the normalization was done?

The results are aggregated over each suite. We ran 5 seeds per task for the locomotion and DMC tasks, so each curve/bar represents a total of 30 runs for the locomotion suit and 35 runs for the other three DMC suits. We use interquartile means and stratified bootstrapped confidence intervals for the error bars. This procedure follows the widely adopted (in RL) suggestions from Argawal et al (2021) and is specifically designed to aggregate results over an experiment suite in a principled manner. We compute the results using the code published by Argawal et al and provide per-task results for completeness in Appendix C. For OpenCabientDrawer we use 20 seeds to establish significance.

Reviewer comment (Evaluation):
The claim "In the more difficult settings, i.e., Occlusions, Locomotion (Fig. 4), and OpenCabinetDrawer (Fig. 5), using a joint representation gives the largest benefits", "In the Locomotion experiments, the CPC approaches (Fig. 4) have a significant edge over reconstruction", is only true for model-based occlusion (Fig. 4), and for Fig. 5 the gain is unclear (see question 1. b).

Author response:
We apologize for the confusing presentation in our original submission. The updated bar plots of the revision show the benefit of the joint representation for all mentioned cases. The new bar plots also clarify that the CPC approaches have an edge over the reconstruction approaches on the locomotion tasks. For example, Joint(CPC+R) achieves a reward of almost 800 while Joint(R+R) achieves only slightly more than 600. Further, we made our statements more precise and updated the references to the new figures in the revision. We ensured all our claims are backed up by experimental findings and added missing details for the experimental description (number of seeds, aggregation) and baselines (DrQv2 for ReplicaDrawer).

Reviewer comment (Experiments):
The statement "In the Locomotion experiments, the CPC approaches (Fig. 4) have a significant edge over reconstruction" is problematic because first, the gain is not significant, and second, Locomotion's model-based results are missing.

Author response:
We ran only model-free experiments for locomotion, as model-free methods outperformed model-based ones on the DMC. We thank the reviewer for pointing out the mistake in the caption of Figure 12 and fixed it in the revision.

Reviewer comment (Experiments):
Given that many details are missing, the conclusion seems too strong and not rigorous enough. Understandably, there may not exist a unified framework that works best for model-based and model-free RL. Some possible improvements include separately discussing (i) model-free and model-based, (ii) locomotion and manipulation, and (iii) standard image and background changes, to make a less strong but more rigorous conclusion. Also, DMC's results are very saturated, so it might be more convincing to include more diverse domains (e.g., more Maniskill/RLBench/FrankaKitchen results).

Author response:
Understandably, there may not exist a unified framework that works best for model-based and model-free RL. Given that many details are missing, the conclusion seems too strong and not rigorous enough. Some possible improvements e.g. separately discuss (i) model-free and model-based, (2) locomotion and manipulation, and (3) standard image and background changes, to make a less strong but more rigorous conclusion.

We apologize for the confusing presentation in our original submission. The updated bar plots of the revision show the benefit of the joint representation for all mentioned cases. The new bar plots also clarify that the CPC approaches have an edge over the reconstruction approaches on the locomotion tasks. For example, Joint(CPC+R) achieves a reward of almost 800 while Joint(R+R) achieves only slightly more than 600. Further, we made our statements more precise and updated the references to the new figures in the revision. We ensured all our claims are backed up by experimental findings and added missing details for the experimental description (number of seeds, aggregation) and baselines (DrQv2 for ReplicaDrawer).  

> Also DMC's results are very saturated, it might be more convincing to include more diverse domains (e.g. more Maniskill/RLBench/FrankaKitchen results).

The occlusion and locomotion tasks are novel tasks building on DMC that we designed to study hard information fusion problems in representation learning for RL. While natural video backgrounds have been previously used for a similar purpose, both related work and our experimental results indicate that there is still room for improvement. We are currently running additional Maniskill experiments, investigating the combination of Depth Images and proprioception. We will add those to the final version of our paper.

Reviewer comment (Theory):
The formulation in equation (4) is not mathematically rigorous. If CPC is applied to estimate the MI between the current representation and the next observation, the KL part should be factorized differently but not naively apply equation (1).

Author response:
Adding the weighted KL term to the CPC objective is based on theoretical (Shu et al, 2021)) and empirical (Nguyen et al (2021),  Srivastava et al (2021)) findings which highlight the importance of enforcing consistency between prior and posterior beliefs when training state space representations with predictive losses.  We conducted preliminary experiments further validating these findings and decided to adopt the method of Srivastava et al (2021).

Reviewer comment (Writing):
Typo in caption: Figure 9 should be model-based results.

Author response:
We thank the reviewer for their time and effort, and particularly for acknowledging our “solid and comprehensive” experiments. We appreciate them pointing out the typos and ambiguities in our initial submission and address them in our revision.