GRAM: Generalization in Deep RL with a Robust Adaptation Module

Thank you for your time and effort in reviewing our paper. It seems that there is a misunderstanding about the assumptions made in our approach:

• GRAM makes no assumptions about which parameters will vary at test time, and automatically adapts at test time to both in-distribution (ID) scenarios seen during training and out-of-distribution (OOD) scenarios not seen during training using only a history of observations and actions. Our experiments analyze both of these settings.
• We utilize a common form of adversary policy for dynamics robustness that does not assume knowledge about the parameters that will vary at test time. We demonstrate the robustness of GRAM across a variety of OOD scenarios that are different from the structure of the adversary policy.

Please see below for detailed responses to your comments, which we hope clarify this misunderstanding and address your main concerns. We ask that you please consider updating your review scores to reflect these clarifications. Based on your suggestion, we have also updated the paper to include sensitivity analysis in Figure 13 related to the GRAM coefficient $\alpha$.

---

The assumption underlying this paper is notably strong. While assuming access to context variable values within the simulation environment is reasonable, assuming the knowledge of which parameters constitute the context is a strong assumption. In their experiments, contexts include friction, added base mass, motor strength, and joint angle bias - effectively acknowledging a priori which parameters will vary in test environments.

It seems that there is a misunderstanding about the assumptions made by our algorithm. As mentioned above, GRAM makes no assumptions about which parameters will vary at test time, and automatically adapts at test time to both ID scenarios seen during training and OOD scenarios not seen during training using only a history of observations and actions.

• We assume access to a set of ID contexts $\mathcal{C_\textnormal{ID}} \subset \mathcal{C}$ during training, and we refer to all other contexts as OOD contexts $\mathcal{C_\textnormal{OOD}}$. Our algorithm is specifically designed to achieve adaptive ID performance for $c \in \mathcal{C_\textnormal{ID}}$, and robust OOD performance for $c \in \mathcal{C_\textnormal{OOD}}$. Please refer to the Problem Statement in lines 115-128.
• We follow the standard teacher-student training procedure used in Contextual RL, where the user defines the range of ID contexts $\mathcal{C_\textnormal{ID}}$ by selecting environment parameters to vary during training (e.g., the parameters in Table 1 for our Base ID setting). These parameters are provided as input to the context encoder $f$ during training (all other possible parameters that impact dynamics would be constant during training, so do not need to be provided as an input to the network). This does not imply that these are the only parameters that may vary at test time.
• Our experiments include OOD scenarios that consider variations of parameters that are different from the ones varied during training. In Figure 5, we test on rough terrain despite only training on flat terrain. In the top row of Figure 4, we test on frozen joints even though the Base ID setting did not observe frozen joints of any kind during training.

Your point highlights the main limitation of Contextual RL methods that we address in this work. Contextual RL often only works well within the range of ID contexts seen during training (as we show in our experiments), which are typically selected based on domain knowledge about likely test-time conditions. GRAM addresses this main limitation of Contextual RL by automatically identifying and reacting to OOD scenarios in a robust fashion at test time, while still retaining the ID adaptation capabilities of Contextual RL.

Thank you for your supportive feedback and detailed suggestions that have helped to improve our paper. We have updated the paper to incorporate additional experimental results and implementation details based on your suggestions (changes highlighted in blue). Please see below for detailed responses to all of your comments and questions. We hope that we have addressed your main concerns, and we ask that you please consider updating your review scores to reflect the revised version of the paper.

---

The authors should consider citing and discussing recent works [1] and [2].

Thank you for pointing out these interesting recent works. We have added references to these in the Related Work section. Note that the methods proposed in these works could be used within our GRAM framework: the robustness curriculum in [1] could be applied in our robust OOD training, and the domain randomization curriculum in [2] could be applied rather than considering a fixed set of ID contexts throughout training.

References:

[1] Reddi et al. Robust Adversarial Reinforcement Learning via Bounded Rationality Curricula. In ICLR 2024.

[2] Tiboni et al. Domain Randomization via Entropy Maximization. In ICLR 2024.

---

The authors provide little information to the implementation of domain randomization (DR) and no comparison of GRAM vs. DR in fig. 4 and 5.

In response to your comment, we have included additional implementation details related to our baselines in Appendix A (lines 796-801). We have also added detailed experimental results comparing GRAM and DR in Figures 8 and 9 in Appendix B, which contain the same analyses as Figures 4 and 5 in the main paper.

---

I highly suggest the authors to compare GRAM against a DR baseline which also uses a history of previous state and actions, as this is the notorious way to implement DR. Furthermore, DR may and should also leverage privileged information at training time, by using the notorious asymmetric actor-critic paradigm.

Thank you for the suggestions.

• Please note that Contextual RL represents a baseline that is trained across a variety of contexts, uses a history of previous states and actions as input to the policy (via the history encoder $\phi$), and uses privileged context information in the critic during training. This is one of our main baselines (which allows for ID adaptation), and GRAM extends this approach to also incorporate OOD robustness.
• Our implementation of DR considers a policy that takes as input the current observation, which is a common baseline used in prior works that consider a teacher-student architecture (e.g., [3, 4]). We chose to compare against a version of DR that does not use privileged context information in the critic because we found that this choice led to slightly more robust performance. Based on your suggestion, we have also added as a baseline a version of DR that uses privileged context information in the critic. We have included this additional baseline in Figure 3, and provide detailed comparisons against GRAM in Figures 8 and 9 in Appendix B.

References:

[3] Kumar et al. RMA: Rapid motor adaptation for legged robots. In RSS 2021.

[4] Margolis et al.. Rapid locomotion via reinforcement learning. The International Journal of Robotics Research, 2024.

---

The analysis is carried out on locomotion environments only, and simulated environments only. The generalization problem under unobservable dynamics is likely exacerbated and more challenging for manipulation and contact-rich settings.

We chose to focus our experimental analysis on quadruped locomotion, which represents a complex task where Contextual RL has been applied for ID adaptation in prior works [3, 4]. We consider the realistic Isaac Lab simulator that includes sensor noise and has been shown to allow for zero-shot sim-to-real transfer [3, 4]. As mentioned in the Limitations and Future Work paragraph (lines 535-538), we focused on simulated experiments to conduct a comprehensive empirical study of GRAM across a variety of ID and OOD settings, and we plan to deploy our GRAM framework in real-world hardware experiments as part of future work.

Note that our GRAM framework can be applied to a variety of applications where dynamics generalization is important, and we agree that it would be interesting to consider other tasks such as manipulation in future work. We have added commentary on this in the Limitations and Future Work paragraph (lines 538-539).

Thank you for your positive feedback on our paper. Based on your suggestions, we have incorporated additional experiments in the revised version of the paper (changes highlighted in blue). Please see below for detailed responses to your comments. If we have addressed your main concerns, we ask that you consider updating your review scores to reflect the revised version of the paper.

---

The paper does not consider other easy to implement baselines, e.g. domain invariance prediction, adding noise to the context, or use a larger replay buffer for the context encoder training. I think to show that this specific design is useful, one should also consider other easy-to-implement baselines.

Thank you for these suggestions. We have added the following experiments to address your comment:

• We include an additional baseline in In Figure 3 that adds random noise to the latent feature variable in Contextual RL, which is related to smoothing methods for robustness [1]. We provide detailed results for this baseline in Figures 10 and 11 in Appendix B, and discuss this baseline in lines 961-968. We see that adding random noise to the latent feature improves the robustness of Contextual RL, but GRAM remains significantly more robust than this baseline in OOD scenarios.
• In Figure 12, we analyze the impact of using a larger replay buffer to train the history encoder $\phi$ in Contextual RL. Our main experiments follow the standard approach from [2], which uses on-policy data collected by the student policy. Using a parallel training scheme with 4096 environments, this results in a large and diverse on-policy batch of ~100k samples. We compare this to using larger replay buffers that store ~1M and ~5M samples, respectively (by reusing samples from the last 10 and 50 policy updates, respectively). We see in Figure 12 that reusing data does not improve the performance of Contextual RL, which still performs significantly worse than GRAM in OOD scenarios.

References:

[1] Kumar et al. Policy smoothing for provably robust reinforcement learning. In ICLR 2022.

[2] Kumar et al. RMA: Rapid motor adaptation for legged robots. In RSS 2021.

---

GRAM biases towards 0, which is the mean at the beginning of the training. It is unclear if this is the best choice. How about bias toward the mean at the end of the training of the context encoder?

Other values of $z_{\textnormal{rob}}$ are possible, but we believe that $z_{\textnormal{rob}} = \mathbf{0}$ is a simple and logical implementation choice for the reasons discussed in lines 242-248. Our experiments demonstrate that this implementation choice leads to strong ID and OOD performance, so we did not experiment with other values. Note that $z_{\textnormal{rob}}$ is provided as an input to the policy for robust OOD training (see the right side of Figure 2), so it is not possible to consider a $z_{\textnormal{rob}}$ that depends on values of latent vectors at the end of training.

Thank you for the detailed comments and positive feedback. Based on your suggestions, we have updated the paper to emphasize our focus on dynamics generalization and include additional experimental results to support our implementation choices (changes highlighted in blue). Please see below for detailed responses to each of your comments. We ask that you please consider updating your review scores to reflect our responses and the revised version of the paper.

---

[W1 / Q1] It is unclear how general is the proposed approach in dealing with different types of variations. What about variations in visual scenes? generalization to different tasks with the same robot?

• This paper focuses on the problem of dynamics generalization, where the context impacts the transition model (see lines 106-107). We have updated the language in the Abstract and Introduction to make this more clear at the beginning of the paper.
• Please note that variations in visual scenes would not impact our experiments, which consider a policy based on proprioceptive observations available from standard onboard sensors (see Table 2). However, it would be interesting to extend our approach to address visual variations as you have mentioned. We have added commentary about this in our Limitations and Future Work paragraph (see lines 534-535).
• Multi-task RL represents an orthogonal line of research, where task information is typically known and can be provided to the policy as an additional input (similar to the velocity command input in our experiments). Our framework could be applied to provide dynamics generalization in a multi-task setting.

---

[W2 / Q2] The in-distribution and out-of distribution generalization motivation is a bit confusing in the introduction. What are examples of scenarios where out-of-distribution generalization is possible without in-distribution generalization?

Robust RL is the main example of this scenario. Robust RL is designed to achieve robust out-of-distribution (OOD) performance, but does not perform in-distribution (ID) adaptation across different training environments. Instead, it outputs an action that maximizes worst-case performance, leading to overly conservative behavior in ID environments.

ID generalization refers to the ability to identify and adapt to different environments seen during training to achieve the best performance. OOD generalization refers to the ability to achieve robust performance in environments that were not seen during training. See lines 33-38 in the Introduction, and lines 124-128 in the Problem Statement. Our algorithm GRAM achieves both adaptive ID performance and robust OOD performance.