Learning Robust Representations with Long-Term Information for Generalization in Visual Reinforcement Learning

• Q6. The recent method MIIR [1] also considers learning robust representations using the information bottleneck framework and conducts experiments on the Distracting DMControl benchmark. Could the author discuss the differences between ROUSER and MIIR, such as their theoretical motivations and empirical performance comparisons?

  • A6. Thanks for bringing this to our attention. We summarize the major differences between ROUSER and MIIR as follows.
    1. Theoretical Analysis. ROUSER provides a theoretical guarantee, establishing an upper bound between the true action-value function and the action-value function on top of learned representations (see Theorem 4.4 in our main text). In contrast, MIIR only derives its optimization objectives without offering any guarantees.
    2. *Applicability. We integrate ROUSER with both policy gradient and value-based VRL algorithms to estimate the action values, making it suitable for continuous and discrete control tasks. MIIR, however, has only been implemented and evaluated on continuous control tasks.
    3. Empirical Performance. For comparison with MIIR, we conducted additional experiments on six DMC-GB tasks under video_easy and video_hard settings, as MIIR did not release its source code and only reported results for these two settings. We report the average results and standard deviations over three seeds in Tables C2 and C3 below. These results demonstrate that ROUSER outperforms MIIR in 9 out of 12 experiments.
       • We have included these results of Tables C2 and C3 in Table 8 (Lines 1467–1480) of Appendix E.4 of our revision. For your convenience, we have highlighted these updates in blue.

  [1] Wang S, Wu Z, Wang J, Hu X, Lin Y, Lv K. How to Learn Domain-Invariant Representations for Visual Reinforcement Learning: An Information-Theoretical Perspective.

Table C2. Results and standard deviations on DMC-GB under the video_easy setting. Results of our approach are over three random seeds.

Table C3. Results and standard deviations on DMC-GB under the video_hard setting. Results of our approach are over three random seeds.

We appreciate reviewer's insightful comments. We respond to each comment as follows and accordingly revise our manuscript, with all updates highlighted in blue for your convenience. We sincerely hope that our responses could properly address your concerns. If so, we would deeply appreciate it if you could raise your score. If not, please let us know your further concerns, and we will continue responding to the comments and enhancing our submission.

• Q1. Lack of training curves on Distracting DMControl benchmark for Figures 1 and 2: Including training curves would help readers understand the sample efficiency and stability of the proposed method.
  • A1. Thank you for pointing this out. We have included all training curves on Distracting DMControl benchmark in Figure 8 of our revision (see Lines 1188-1223).
• Q2. In Figures 1 and 2, the baselines include both standard methods for image-based RL (e.g., DrQv2, CURL, TACO) and methods designed to be robust against distractors (e.g., SRM, CRESP, RAP). It would be better to reorganize the results, such as grouping methods by type or using different visual indicators for a clearer comparison.
  • A2. Thanks for your valuable suggestions.
    • In Lines 432-453 of our revision, we updated Figure 2 (a merge of Figures 1 and 2 from our previous manuscript) by grouping methods based on their type. This improves clarity, making the comparisons across methods more intuitive.
      • We used warm colors to show VRL methods designed to improve robustness against distractions (i.e., SRM, CRESP, and RAP) on the left of each subplot. These methods were compared with ROUSER-SRM, which outperforms them in 11 out of 12 experiments.
      • We used cooler colors to present standard VRL methods (i.e., DrQv2, CURL, and TACO) on the right of each subplot. We compare these methods with ROUSER-DrQv2, which outperforms DrQv2 in all experiments and surpasses the other methods in 7 out of 12 experiments.
• Q3. Typo in line 430: Nore -> Note
  • A3. Thank you! We have corrected this typo in Line 431 of our revision.
• Q4. In Figure 4, axis labels should be included directly on the plots rather than in the caption.
  • A4. Thanks for bringing this to our attention. We have updated each subplot in Figure 4 to include x- and y-axis labels directly on the plots (see Lines 486-495), instead of referencing them in the caption.
• Q5. In Section 5.3, the authors analyze the effectiveness of ROUSER in capturing long-term information with a fixed length of 300. It would be helpful to see how the effectiveness changes with varying lengths (e.g., 50, 100, 500).
  • A5. Thanks for the kind and insightful comment.
    • We have conducted additional experiments with varying reward sequence lengths, including 50, 100, 300, and 500. We reported the results in Table C1 below (also see Table 3 with its description in Lines 1287-1304 in Appendix E.1 of our revision). These results, averaged over three seeds on cheetah_run task with unseen color distractions, demonstrate that the representations learned by ROUSER are able to accurately predict reward sequences of different lengths, with performance improving as the reward sequences become longer.

Table C1. MSE results and standard deviations with different reward sequence lengths over three seeds, including 50, 100, 300, and 500.

We appreciate the reviewer's insightful and constructive comments and suggestions, which help us significantly improve the clarity and quality of our manuscript. We have carefully addressed these concerns and accordingly revised the manuscript, with all updates highlighted in blue for your convenience. We sincerely hope that our responses could properly address your concerns. If so, we would deeply appreciate it if you could raise your score. If not, please let us know your further concerns, and we will continue actively responding to the comments and enhancing our submission.

• Q1. The paper is not clearly written, it took me some time to go through and understand it. I suggest improving the flow a bit. For example:

  (1) something that could help is better structure. In the beginning of section 4, I would say something about the big picture of the paper, as "We propose robust reward representations, then use td learning to learn robust value representations. Throughout, we use IB-based objective make representations robust."

  (2) At the end of section 4, it would also help to list the objectives, or add the figure or pseudocode from the appendix. I understand you don't have much space, maybe decreasing the size of figures 1 and 2, or aggregating results over tasks could be a way to go here.

  • A1. Thanks for your valuable feedback and suggestions regarding the structure of our paper.
    • For (1), to improve clarity, we have revised the beginning of Section 4 to provide a clear "big picture" overview of the paper. We highlight the update in blue in Lines 181-184 our revision, and for your convenience, we present it as follows.
      • "In this paper, we propose to learn robust action-value representations that capture long-term information for decision-making. Then, we introduce robust and compressed reward representations and use TD learning paradigm to learn robust action-value representations, guided by an IB-based objective to ensure robustness throughout the learning process."
    • For (2), we have moved the figure of ROUSER's architecture from the appendix to the end of Section 4 (see Lines 378-396) in our revision. Moreover, we have included the proposed objectives at the end of Section 4 (see Lines 404-406) to enhance readability. We highlight all these changes in blue.
      • Due to the page limitation, we have merged the figures related to the DCS results from our main experiments into Figure 2 (Lines 432-454) and resized Figures 2, 3 (Lines 470-479), and 5 (Lines 499-509) in our revision. We thank you for the valuable suggestions to help streamline the content.

• Q2.

  (1) Line 506: What information are you predicting in this frozen evaluation? Is it an average of 300 rewards in the future?

  (2) What exactly is the MSE between?

  (3) Line 502: all methods' prediction: what does that mean?

  (4) Line 504: "use a 2 layer mlp with a learning rate of 1e4". I understand what you're saying, but the sentence is awkward.

  • A2. We apologize for any confusion caused by our unclear descriptions. We response to each comment as follows.
    • For (1): No, we predict vectors with a dimension of 300, which are the future reward sequences over a length of 300.
    • For (2): We compute the MSE between the MLP outputs and samples of future reward sequences.
      • In this experiment, we first use a 2-layer trainable MLP to map the fixed representations into outputs, which are vectors with a dimension of 300. These fixed representations are extracted from the center layer of the critic in each method, using state-action pairs as inputs. We then compute the MSE between the MLP outputs and the samples of future reward sequences. A lower MSE indicates that the representations are better at capturing long-term information.
    • For (3) and (4): Thank you for pointing this out. We have revised the whole paragraph as follows (also see Lines 499-510 in our revision).
      • "We assess how ROUSER can capture long-term robust information. Specifically, we first collect 200K samples under unseen color distractions using a DrQv2's policy. Next, we compute the representations learned by each method from state-action pairs in the collected data. Such representations are extracted from the center layer of the critic in each method. Then, we fix these representations and input them into a 2-layer trainable MLP to predict the future reward sequences over a length of 300. During the prediction process, we use the Adam optimizer with a learning rate of 1e-4. As shown in Figure 4(a), we plot curves of the MSE loss on 1K evaluation samples across three seeds. The lower MSE losses of ROUSER indicate that its learned representations capture long-term information more effectively."

We appreciate the insightful and constructive comments. We respond to each comment as follows and accordingly revise our manuscript, with all updates highlighted in blue for your convenience. We sincerely hope that our responses could properly address your concerns. If so, we would deeply appreciate it if you could raise your score. If not, please let us know your further concerns, and we will continue actively responding to the comments and enhancing our submission.

• Q1. The idea of incorporating Q-values to mitigate visual distractions is a commonly adopted design choice in VRL, e.g. [1,2].

  • A1. Thank you for your valuable comments. We present the major contributions of ROUSER as follows.
    1. ROUSER provides a theoretical guarantee, establishing an upper bound between the true action-value function and the action-value function on top of learned representations (see Theorem 4.4 in our main text). In contrast, [1,2] only derive their optimization objectives without offering any guarantees.
    2. We integrate ROUSER with both policy gradient and value-based VRL algorithms to estimate the action values, making it suitable for continuous and discrete control tasks. In contrast, [1,2] have only been implemented and evaluated on continuous control tasks.
    3. To the best of our knowledge, ROUSER is the first to extend action-value representations to VRL against visual distractions, providing a new perspective on task decomposition and the development of VRL methods.
    4. We have conducted additional experiments on six DMC-GB tasks across all four settings to comprehensively demonstrate that ROUSER outperforms the generalization baselines (including [1,2]) in 17 out of 24 settings. Please refer to A2 below for details.

  [1] Look where you look! Saliency-guided Q-networks for generalization in visual Reinforcement Learning. NeurIPS 2022.

  [2] How to Learn Domain-Invariant Representations for Visual Reinforcement Learning - An Information-Theoretical Perspective. IJCAI 2024.

• Q2. (1) The DMC-GB is a commonly adopted visual generalization benchmark with color_hard and video_hard setting [1,2]. I believe the evaluation results on DMC-GB would be more persuasive. (2) For baselines, I believe comparing with generalization baselines is necessary. For example: TPC [3], MIIR [4], TIA [5], and PIEG [6].

  • A2. Thank you for the valuable feedback.

    • For (1), to comprehensively evaluate our approach, we have conducted additional experiments on six DMC-GB tasks across all four settings (i.e., color_easy, color_hard, video_easy, and video_hard). Please refer to Tables B1, B2, B3, and B4 below.

    • For (2), we have incorporated results from six VRL generalization baselines (including DrG [7], MIIR, PIE-G, SVEA [8], SGQN [9], and SECANT [10]) into our evaluations (also see Tables B1, B2, B3, and B4).

      • Note that we replaced TPC with DrG, a model-based VRL generalization method similar to TPC, which uses the current state space model to enhance the robustness of latent representations against visual distractions. DrG demonstrates better generalization performance on DMC-GB than TPC.
      • Due to computational resource constraints, the TIA experiments are still ongoing. We are committed to reporting TIA's results on DMC-GB as soon as they are completed.

      All tables show that ROUSER outperforms the aforementioned baselines in 17 out of 24 settings, demonstrating ROUSER's robustness and effectiveness in DMC-GB tasks.

  [1] Look where you look! Saliency-guided Q-networks for generalization in visual Reinforcement Learning. NeurIPS 2022.

  [2] Stabilizing Deep Q-Learning with ConvNets and Vision Transformers under Data Augmentation. NeurIPS 2021.

  [3] Temporal Predictive Coding For Model-Based Planning In Latent Space. ICML 2021.

  [4] How to Learn Domain-Invariant Representations for Visual Reinforcement Learning - An Information-Theoretical Perspective. IJCAI 2024.

  [5] Learning Task Informed Abstractions. ICML 2021.

  [6] Pre-Trained Image Encoder for Generalizable Visual Reinforcement Learning. ICML 2022.

  [7] Dream to Generalize: Zero-Shot Model-Based Reinforcement Learning for Unseen Visual Distractions. AAAI 2023.

  [8] Stabilizing Deep Q-Learning with ConvNets and Vision Transformers under Data Augmentation. NeurIPS 2021.

  [9] Look where you look! Saliency-guided Q-networks for generalization in visual Reinforcement Learning. NeurIPS 2022.

  [10] Secant: Self-expert cloning for zero-shot generalization of visual policies. ICML 2021.

• We conducted comprehensive experiments across all four DMC-GB settings: color_easy, color_hard, video_easy, and video_hard.
  • Most of the baseline results we report are taken directly from their respective papers. For cases where results are not provided, we ran experiments using the available source code and hyperparameters under three random seeds.
    • We ran PIE-G, DrG, and SGQN on all DMC-GB tasks across color_easy and color_hard settings.
    • We ran DrG on finger_spin and walker_stand tasks across video_easy and video_hard settings.
    • We ran SVEA on all DMC-GB tasks under the color_easy setting and on cheetah_run task under the color_hard setting.
  • Note that since SECANT and MIIR did not release their source code, we only included the results reported in their respective papers.
• We have included all these updates in Appendix E.4 of our revision, including Tables B1 and B2 (see Table 7 in Lines 1422–1436), Tables B3 and B4 (see Table 8 in Lines 1467–1480), and their descriptions (see Lines 1394-1403). For your convenience, we have highlighted these updates in blue.

Table B1. Results and standard deviations on DMC-GB under the color_easy setting. Results of our approach are over three random seeds.

Table B2. Results and standard deviations on DMC-GB under the color_hard setting. Results of our approach are over three random seeds.

Table B3. Results and standard deviations on DMC-GB under the video_easy setting. Results of our approach are over three random seeds.

Table B4. Results and standard deviations on DMC-GB under the video_hard setting. Results of our approach are over three random seeds.

Dear Reviewer 238T,

We deeply appreciate the time and effort you've invested in reviewing our work.

We have completed the additional experiments of TIA on six DMC-GB tasks across all four settings and updated all DMC-GB results in Tables D1 and D2 below. In addition, we have updated Appendix E.4 in our revision, where Table D1 is presented as Table 7 in Lines 1422–1436, and Table D2 is presented as Table 8 in Lines 1459–1472.

• Since the TIA agent based on its original source code always produced NaN outputs, we used the reproduced version of TIA provided at https://github.com/zchuning/repo.

The results remain unchanged, demonstrating that ROUSER outperforms the aforementioned baselines in 17 out of 24 DMC-GB experiments.

We hope our responses have thoroughly addressed your concerns, and we remain fully prepared to address any additional questions you may have to further refine our manuscript. If you find the revisions satisfactory, we would be grateful if you could raise your score.

Thank you once again for your invaluable guidance and support.

Best,

Authors

We appreciate the reviewer's insightful and constructive comments and suggestions. We respond to each comment as follows and accordingly revise our manuscript, with all updates highlighted in blue for your convenience. We sincerely hope that our responses could properly address your concerns. If so, we would deeply appreciate it if you could raise your score. If not, please let us know your further concerns, and we will continue actively responding to the comments and enhancing our submission.

• Q1. To provide a more comprehensive evaluation, you should run additional experiments on "walker_stand" to better assess your method's generalization performance in alignment with this benchmark.

  • A1. Thanks for the kind and insightful comment.
    • For a comprehensive evaluation, we have conducted additional experiments on six DMC-GB tasks (including the "walker_stand" task) across all four settings (color_easy, color_hard, video_easy, and video_hard). We report the average results and standard deviations over three seeds. These results demonstrate that ROUSER outperforms others in 17 out of 24 experiments. Please refer to Tables B1, B2, B3, and B4 below for details.
    • Moreover, it is worth noting that Distracting Control Suite (DistractingCS)---which we use in our paper to evaluate generalization capabilities of agents---introduces visual distractions (i.e., random colors and video backgrounds) similar to those of the DMC-GB [1], but with higher stochasticity (see DMC-GB).

  [1] Generalization in Reinforcement Learning by Soft Data Augmentation. ICRA 2021.

• Q2. To provide a more relevant comparison, you should include results for your method alongside established methods like SECANT and PIE-G, which were specifically developed to enhance generalization in VRL.

  • A2. Thanks for your suggestion.
    • We have included results from six VRL generalization methods (including PIE-G [1], SECANT [2], DrG [3], SVEA [4], SGQN [5], and MIIR [6]) in our evaluations. Tables B1, B2, B3, and B4 present these results on DMC-GB, demonstrating the superiority of ROUSER for generalization.
      • Note that SECANT did not provide DMC-GB results under color_easy and video_hard settings, nor did it release the source code. Thus, we only included the results in color_hard and video_easy settings reported in its paper.

  [1] Pre-trained image encoder for generalizable visual reinforcement learning. ICML 2022.

  [2] Secant: Self-expert cloning for zero-shot generalization of visual policies. ICML 2021.

  [3] Dream to Generalize: Zero-Shot Model-Based Reinforcement Learning for Unseen Visual Distractions. AAAI 2023.

  [4] Stabilizing Deep Q-Learning with ConvNets and Vision Transformers under Data Augmentation. NeurIPS 2021.

  [5] Look where you look! Saliency-guided Q-networks for generalization in visual Reinforcement Learning. NeurIPS 2022.

  [6] How to Learn Domain-Invariant Representations for Visual Reinforcement Learning - An Information-Theoretical Perspective. IJCAI 2024.

Table B3. Results and standard deviations on DMC-GB under the video_easy setting. Results of our approach are over three random seeds.

Table B4. Results and standard deviations on DMC-GB under the video_hard setting. Results of our approach are over three random seeds.

• We conducted comprehensive experiments across all four DMC-GB settings: color_easy, color_hard, video_easy, and video_hard.
  • Most of the baseline results we report are taken directly from their respective papers. For cases where results are not provided, we ran experiments using the available source code and hyperparameters under three random seeds.
    • We ran PIE-G, DrG, and SGQN on all DMC-GB tasks across color_easy and color_hard settings.
    • We ran DrG on finger_spin and walker_stand tasks across video_easy and video_hard settings.
    • We ran SVEA on all DMC-GB tasks under the color_easy setting and on cheetah_run task under the color_hard setting.
  • Note that since SECANT and MIIR did not release their source code, we only included the results reported in their respective papers.
• We have included all these updates in Appendix E.4 of our revision, including Tables B1 and B2 (see Table 7 in Lines 1422–1436), Tables B3 and B4 (see Table 8 in Lines 1467–1480), and their descriptions (see Lines 1394-1403). For your convenience, we have highlighted these updates in blue.

Table B1. Results and standard deviations on DMC-GB under the color_easy setting. Results of our approach are over three random seeds.

Table B2. Results and standard deviations on DMC-GB under the color_hard setting. Results of our approach are over three random seeds.

Dear Reviewer M8ko,

Thanks for your kind support and the further improvements you have suggested for our manuscript. We sincerely appreciate your insightful feedback, which has significantly helped us in ensuring a more comprehensive evaluation of our approach.

Thank you again for your valuable comments and guidance.

Best,

Authors