Fast and Robust: Task Sampling with Posterior and Diversity Synergies for Adaptive Decision-Makers in Randomized Environments

We sincerely appreciate Reviewer beLW's efforts and recognition of our work. Below, we improve the manuscript based on beLW's feedback.

---

1. Terminology clarity about secret MDP

Thank you for your valuable feedback. We used the term "secret MDP" to highlight that we are the first to model robust active task sampling as an MDP and solve it with developed i-MABs. We will further explain this point in the revised manuscript and consider using a more intuitive term, such as "task-selection MDP" to avoid potential confusion.

2. Comparison betwen PDTS and MPTS

Thank you for the suggestion. Currently, we primarily summarize the orthogonal contributions of MPTS and PDTS in Table 1. For example, MPTS develops a VAE-like risk predictive model to achieve nearly $\text{CVaR}_\alpha$ optimization but suffers from the concentration issue. Our PDTS (i) introduces a versatile theoretical tool i-MAB to achieve a more robust solution, i.e., nearly worst-case optimization for meta RL and DR cases, (ii) resolve the concentration issue, and (iii) offer easier implementation with stochastic optimism. These three contribution points are new and increase scalability of RATS. Encouragingly, PDTS with i-MABs provides stable and tractable scheme for worst-case optimization in adaptive decision-making, broadening its applicability.

In the revised manuscript, we will involve more RATS background and further highlight the significance of PDTS for more general readers.

---

Once again, thank you for your valuable review and thoughtful recognition of our work. Your feedback is greatly appreciated and has significantly improved our manuscript. We hope our responses sufficiently answer your questions.

We sincerely appreciate Reviewer v1xi's efforts and positive feedback. Below, we provide our responses.

---

1. Clarification on the simplification of configurations

Apologies for any confusion. The simplification of configurations lies in two aspects:

• As stated in Lines 261–270 of Section 3.2, MPTS requires careful tuning of the candidate task batch size $\hat{\mathcal{B}}$. In contrast, PDTS eliminates this requirement and remains scalable even under extreme worst-case optimization (e.g.,$\hat{\mathcal{B}} = 64 \times \mathcal{B}$) by introducing a diversity-regularized acquisition function to mitigate the concentration issue.
• As stated in Lines 281–284 of Section 3.3, we use posterior sampling to utilize stochastic optimism while avoiding the calibration of exploration and exploitation weights in subset search, which is required by UCB-based methods.

We will highlight these in the revised manuscript.

2. Discussion on assumptions failure

Thank you for your thoughtful question. The assumptions in our analysis are consistent with those in prior works[1,2], as stated in Appendix A. If they do not hold, the effectiveness of the risk predictive model may degrade with extremely lower Pearson correlation coefficients, potentially impacting robust optimization. Fortunately, our experiments demonstrate that PDTS achieves higher PCC values, strong robust optimization performance, empirically supporting the validity of these assumptions. We will incorporate this discussion into the revised manuscript for clarity.

---

Once again, thank you for your valuable review. Your feedback is greatly appreciated and has helped improve our manuscript. We hope our responses adequately address your concerns, and we would be grateful if you could consider raising the score.

---

References:
[1] Greenberg I, Mannor S, Chechik G, et al. Train hard, fight easy: Robust meta reinforcement learning[J]. Advances in Neural Information Processing Systems, 2023, 36: 68276-68299.
[2] Wang, Q., Lv, Y., Xie, Z., & Huang, J. (2023). A simple yet effective strategy to robustify the meta learning paradigm. Advances in Neural Information Processing Systems, 36, 12897-12928.

We sincerely appreciate Reviewer b9NR's efforts and constructive feedback. Below, we provide our responses.

---

1. Additional computational overhead

Thanks for precious comment. We quantitatively analyzed the extra computational overhead in Fig.5. Even with a $64\times$ candidate batch (i.e., $\text{CVaR}_{1-1/64}$ approximating the worst case), the additional computational overhead remains negligible due to the efficiency of the risk predictive model—its cost is significantly lower than that of agent-environment interactions and policy optimization in MetaRL or DR. We will emphasize this point in the revised manuscript.

2. Ablation study on the risk predictive model

Thank you for the valuable question—it has been very helpful to our analysis. We conducted an ablation study to analyze the impact of the risk predictive model. We designed two variants, PDTS-Deep and PDTS-Shallow, by increasing and decreasing the number of encoder-decoder layers, respectively. Additionally, we replaced the encoder-decoder structure with an MLP to create PDTS-MLP. Results on Walker2dVel are summarized in the table below:

From these results, we observe:

• All PDTS variants achieve better task robust adaptation, confirming the effectiveness of PDTS and its generality across different risk predictive models.
• The encoder-decoder architecture generally outperforms MLP-based models, supporting the rationale behind this design.
• Deeper networks may introduce a performance-robustness trade-off in the current setting, which we plan to further investigate in more complex scenarios.
• Weaker risk prediction models degrade overall performance, due to poorer difficult MDP identification.

3. Scaling to high dimensions

Thank you! This is a very insightful and important question. Since RATS is still in its early stages, we are actively exploring its scalability to high-dimensional task spaces. We agree that heuristic or projection strategies could be viable solutions. One potential approach is to leverage lightweight general-purpose embedding models, such as WordLLaMA[1], to compress high-dimensional task identifiers from language or vision modalities[2]. We appreciate your insight and will continue to investigate this direction further to broaden the application scope of PDTS.

4. Sensitivity of hyperparameter $\gamma$ and tuning practices

We conducted an ablation study on $\gamma$ in Figure 13 and found that PDTS is relatively stable to $\gamma$ within a certain range. Below, we summarize two key tuning recommendations

• In most cases, setting $\gamma$ to 1 or a nearby value secures superior enough performance.
• For scenarios with an extremely low-dimensional task identifier, increasing $\gamma$ appropriately may improve performance.

We will incorporate these recommendations into the revised manuscript and specify the $\gamma$ values used in practice in the released code to aid practitioners.

5. Would you expect the best $\gamma$ to scale with $\hat{\mathcal{B}}$ or with dimension $d$?

We suggest using a larger candidate batch, $\hat{\mathcal{B}}$, to better capture the worst-case scenario and improve performance, as shown in Figure 13. Therefore, we set $\hat{\mathcal{B}} = 64 \times \mathcal{B}$ in all main experiments without delicate tuning and believe it's unnecessary to co-tune it with $\gamma$. We hypothesize that $\gamma$ scales positively with $\hat{\mathcal{B}}$, as larger batches exacerbate concentration issues.

Additionally, we expect that, in comparable scenarios, a smaller $d$ leads to a larger $\gamma$, as concentration issues become more likely. For example, we use $\gamma = 1$ for the task with 2D task identifier (Walker2dMassVel) and $\gamma = 5$ for the 1D task (Walker2dVel).

---

Once again, thank you for your valuable review and thoughtful recognition of our work. Your feedback is greatly appreciated and has significantly improved our manuscript. We hope our responses sufficiently address your concerns.

---

References:
[1] Miller, D. L. (2024). WordLlama: Recycled token embeddings from large language models. https://github.com/dleemiller/wordllama
[2] Kim M J, Pertsch K, Karamcheti S, et al. Openvla: An open-source vision-language-action model[J]. arXiv preprint arXiv:2406.09246, 2024.

We sincerely thank Reviewer xw4Y's efforts and thoughtful feedback. Below, we provide our responses.

---

1. Discussion of curriculum learning methods

Thank you for your valuable suggestion. We'll cite these important works and discuss them in the revised manuscript, adding contents below to Line 667:

Curriculum learning is also a crucial topic related to adaptive decision-making. Dennis, Michael, et al. [1] develop unsupervised environment design (UED) as a novel paradigm for environment distribution generation and achieve SOTA zero-shot transfer. Jiang, Minqi, et al. [2] cast prioritized level replay to enhance UED and formulate dual curriculum design for improving OOD and zero-shot performance. In [3], heavy-tailed distributions are incorporated into the automated curriculum, which leads to robustness improvement. In contrast, our work emphasizes robust task adaptation. Integrating the idea of surrogate evaluation from PDTS into curriculum design could be an interesting direction for future research.

2. Writing optimization

Thank you for your constructive suggestion. We will further simplify the math notations in Sections 2 and 3 to improve comprehension. For instance, instead of introducing the new notation $\mathbf{R}$, we will use the more intuitive summation form of the cumulated return to enhance clarity.

3. Trade-off between worst-case performance and the computational costs

Thanks for the insightful comment. This is a crucial point that we also investigated in Lines 432–438, as well as Figures 5 and 13. Figure 5 demonstrates that even with a $64\times$ candidate batch (i.e., $\text{CVaR}_{1-1/64}$ approximating the worst case), the additional computational overhead remains negligible due to the efficiency of the risk predictive model—its cost is significantly lower than that of agent-environment interactions and policy optimization in MetaRL or DR. Figure 13 shows that increasing the candidate batch size improves nearly worst-case performance. Therefore, we adopt the $64\times$ pseudo batch consistent in all main experiments without extra adjustments. We will emphasize this point in the revised manuscript.

4. Limitations discussion

Thank you for your insightful feedback. We briefly discussed the limitations of the proposed method in the Conclusion (Line 435-438) and Appendix E (Line 1347-1348). Specifically, our approach (1) relies on the risk predictive model, for roughly scoring task difficulties, and (2) depends on identifier information and the inherent smoothness of the adaptation risk function, which might not hold in restricted scenarios. We will provide a more detailed discussion in the revised manuscript.

5. Meta-RL agent have access to task identifiers

Apologies for the confusion. We will clarify this in the revised manuscript: The task identifiers are indeed visible in Meta-RL, as is consistent with prior works such as ROML [4] and MPTS [5]. Detailed information on the task identifiers can be found in Table 3.

---

Once again, thank you for your valuable review. Your feedback is greatly appreciated and has helped improve our manuscript a lot. We hope our responses adequately address your concerns, and we would be grateful if you would reconsider the evaluation and update the score accordingly.

---

References:
[1] Dennis, M., Jaques, N., Vinitsky, E., Bayen, A., Russell, S., Critch, A., & Levine, S. (2020). Emergent complexity and zero-shot transfer via unsupervised environment design. Advances in neural information processing systems, 33, 13049-13061.
[2] Jiang, M., Dennis, M., Parker-Holder, J., Foerster, J., Grefenstette, E., & Rocktäschel, T. (2021). Replay-guided adversarial environment design. Advances in Neural Information Processing Systems, 34, 1884-1897.
[3] Koprulu, C., Simão, T. D., Jansen, N., & Topcu, U. (2023, July). Risk-aware curriculum generation for heavy-tailed task distributions. In Uncertainty in Artificial Intelligence (pp. 1132-1142). PMLR.
[4] Greenberg I, Mannor S, Chechik G, et al. Train hard, fight easy: Robust meta reinforcement learning[J]. Advances in Neural Information Processing Systems, 2023, 36: 68276-68299.
[5] Wang Q C, Xiao Z, Mao Y, et al. Beyond Any-Shot Adaptation: Predicting Optimization Outcome for Robustness Gains without Extra Pay[J]. arXiv preprint arXiv:2501.11039, 2025.