Fast and Robust: Task Sampling with Posterior and Diversity Synergies for Adaptive Decision-Makers in Randomized Environments
Authors: Yun Qu
Cheems Wang
Yixiu Mao
Yiqin Lv
Xiangyang Ji
Venue: ICML 2025 poster (2025)
Keywords: Meta Learning
Robust Reinforcement Learning
Primary area: general_machine_learning->transfer_multitask_and_metalearning
TL;DR: We propose an easy-to-implement method, referred to as Posterior and Diversity Synergized Task Sampling (PDTS), to accommodate fast and robust sequential decision-making

Abstract: Task robust adaptation is a long-standing pursuit in sequential decision-making.
Some risk-averse strategies, e.g., the conditional value-at-risk principle, are incorporated in domain randomization or meta reinforcement learning to prioritize difficult tasks in optimization, which demand costly intensive evaluations.
The efficiency issue prompts the development of robust active task sampling to train adaptive policies, where risk-predictive models can surrogate policy evaluation. 
This work characterizes robust active task sampling as a secret Markov decision process, posits theoretical and practical insights, and constitutes robustness concepts in risk-averse scenarios.
Importantly, we propose an easy-to-implement method, referred to as Posterior and Diversity Synergized Task Sampling (PDTS), to accommodate fast and robust sequential decision-making.
Extensive experiments show that PDTS unlocks the potential of robust active task sampling, significantly improves the zero-shot and few-shot adaptation robustness in challenging tasks, and even accelerates the learning process under certain scenarios.

Fast and Robust:
Task Sampling with Posterior and Diversity Synergies for Adaptive Decision-Makers in Randomized Environments
Yun Qu Qi Wang 1
Yixiu Mao"/ 1 Yiqin Lv
Xiangyang Ji
Abstract
environment, c.g., placing distribution over Markov de- cision processes (MDPs), for policy search in zero-shot or few-shot manner. This facilitates the rise of domain ran- domization (DR) (Muratore et al., 2018) and Meta-RL (Finn et al., 2017) paradigms, which train adaptive policies task episodic learning. Simultaneously, adaptation robustness to worst-case scenarios catching increasing attention most real-world decision-making scenarios are inherently risk- sensitive, where failures in adaptation can cause catastrophic outcomes, c.g., damage to robots (Carpin et al., 2016) or accidents in autonomous driving (Rempe et al., 2022).
Task robust adaptation is a long-standing pur- suit in sequential decision-making. Some risk- averse strategics, c.g., the conditional value-at- risk principle, are incorporated in domain ran- domization
reinforcement learning pri- oritize difficult tasks in optimization, which de- mand costly intensive evaluations, The efficiency issue prompts the development of robust active task sampling to train adaptive policies, where risk-predictive models are used to surrogate pol- icy evaluation. This work characterizes the opti- mization pipeline of robust active task sampling Markov decision process, posits theoretical and practical insights, and constitutes robustness concepts in risk-averse scenarios. Importantly, we propose an easy-to-implement method, referred to as Posterior and Diversity Synergized Task Sam- pling (PDTS), to accommodate fast and robust sequential decision-making. Extensive experi- ments show that PDTS unlocks the potential of bust active task sampling, significantly improves the zero-shot and few-shot adaptation robustness in challenging tasks, and even accelerates the learning process under certain scenarios. Our project website is at https://thu-rlaan to/puu.ooe//DDSssooeeee....e.
Active Inference's Promise for Adaptive Robust Decision- Maker: When risk-averse principles are incorporated in DR and Meta-RL to enhance adaptation robustness (Wang et al., 2024c; Lv et al., 2024; Greenberg et al., 2024); prioritiz- ing challenging tasks to optimize demands intensive and expensive policy evaluation in massive environments over iteration. To overcome the efficiency bottleneck, Wang et al. (2025) constructs risk predictive model to actively infer MDP difficulties for worst subset selection in policy search, which we identify as a method of the robust active task sampling (RATS) paradigm in Fig. 1. As policy evalua- tion in arbitrary MDPs can be amortized by executing this risk predictive model, the resulting model predictive task sampling (MPTS) (Wang et al., 2025) enables expanding the scope of task subset selection, c.g., screening B from B MDPs, to learn adaptive policies without extra environment interaction cost. This reflects RATS's huge potential for efficient, robust decision- making when exhaustive policy evaluation is prohibitive in the vast MDP space.
1. Introduction
Challenges in Theoretical Analysis and Implementations: Despite RATS's promise in decision- decision-making, we can still perceive several issues from its latest SOTA method MPTS, (i) No versatile tool is developed for theoretical analysis, e.g., the robustness concept in optimization, which is an indispensable consideration in risk-averse cases. (ii) It requires dedicated pseudo batch size B and other con- figurations. As implied in Meta-RL results (Wang et al., 2025), appropriately mixing up the random and predictive samplers decisive in task subset selection; otherwise, degrades generalization and robustness (See See. 3.2). (iii) There lack comprehensive discussions about acquisition
Deep reinforcement learning (RL) has garnered remarkable progress in solving complicated sequential decision-making problems in the past few years (Sutton & Barto, 2018). How- ever, an existing challenge is effectively transferring the RL policy unseen but similar scenarios without learning from scratch. A commonly used strategy is to randomize the
Equul contribution 'Department of Automation, Tsinghua University, Beijing, China. Correspondence to: Xiangyang Ji Exyji@tsinhua.ccuu.nn
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

PDTS for Adaptive Decision-Makers in Randomized Environments
Decision-Maid Makinn in Randomized Environments
Robust Active Task Sampling
Candidate Tisk Identifiers
f1
fiel
p(7)
Batch
1. Random Sampler
Tasks {nu
Predicted Reaults
Candidate Tasks
Amortioed Evaluation with PS
Subsets Diversity
(fu)
Ria Values
py
(z.
T
Evaluation
Evaluation
Amortized
di
As...
Adaptive Coriisio-MMaker
Rink Predictive Model
T
T
Top
Subacts
dh...
Mernory Buffer
He = (ir 4
Subeet Selection in PDTS
max
Ar(T" 7 ((d((,
7
(a) General Framework of RATS in Risk-Averse Socision-Making
(b) PDTS as RATS Method
Figure 1. (a) General RATS in risk-averse decision-making. The pipeline involves amortized evaluation of task difficulties, robust subset selection, policy optimization in the MDP batch, and risk predictive models' update. [fire: updates; snow: evaluation] (b) PDTS as a RATS method. PDTS treats task subsets as bandit arms, evaluates values through posterior sampling, and solves a regularized problem.
principles in MPTS. The adopted upper confidence bound (UCB) principle must strike balance between worst-case and uncertainty, with hyper-parameters carefully adjusted in implementation.
from stochastic optimism in posterior sampling for decision-male makingg.
Empirically, the most thrilling finding is that PDTS exhibits adaptation robustness superior to existing SOTA baselines in typical DR and Meta-RL benchmarks without complicated configurations. Even in more realistic and challenging sce- narios, such as vision- based decision-making. PDTS retains remarkable performance over others.
Making Sense of RATS in Decision-Making: Regarding theoretical understanding, we first abstract the general RATS process as a task-selection MDP M, construct an infinitely many-armed bandit (i-MAB) (Carpentier & Valko, 2015) for task subset selection and demonstrate MPTS (Wang al., 2025) as a special solution the i-MAB. Aim at exploring more optimal subsets with the risk predictive model; we make a diagnosis of the concentration issue under a larger B and enhance the acquisition function with the diversity regularization (Wang et al., 2023b; Borodin et al., 2). To simplify the amortized evaluation and exploit the stochas- tie optimism i-MAB, we adopt the posterior sampling strategy (Russo& & Van Roy, 2014) to search for the optimal task subset with fewer configurations. Under these modifi- cations, we present the Posterior and Diversity Synergized Task Sampling (PDTS) as competitive RATS method.
2. Research Background
Notations. For conciseness and coherence, we retain most notations in (Wang et al., 2025) and leave the MDP dis- tribution P(T) details and RL preliminaries in Appendix A. Both DR and Meta-RL -LLL perform policy search in p(T). where MDPs as tasks are specified by physics identifiers
Rdd
The primary goal of DR (Tobin al., 2017; Muratore et al., 2018 8) and Meta-RL. (Beck et al., 2023) is seek a policy O θ that adapts well to new MDP with zero a few roll- outs. We denote the task rask-specific dataset by D, = DUD where DD are the K-rollouts for fast adaptation in MDP r and D are rollouts for after-adaptation policy evaluation. DR differs from Meta-RL and learns to adapt a zero-shot manner (Mehta et al., 2020; Chi al., 2024a), indicating DE 0. The risk function is : D O - IR, mapping the adaptation risk value, e.g., negative average returns of query rollouts.
Contributions and Fascinating Discoveries in PDTS. This work is built on empirical findings and the risk predictive model in MPTS, while the research focus is orthogonal (See Table 1). Surrounding risk-averse adaptive decision-male maainng. the primary contributions are:
1. Our constructed i-MAB provides versatile model to achieve RATS under various principles, including but not limited to MPTS and PDTS. The separate robust-
In task episodic learning of DR and Meta-RL, the optimiza- tion history is written A. {0z, (Tr,;, Dr. le.)) J with B the number of tasks to optimize in t-th iteration. MPTS (Wang et al., 2025) further prepares it as H, )
ness concepts can be refined accordingly.
2. The designed diversity regularized acquisition function fixes the concentration issue, allows for exploration in wider range of task sets (e.g., B- = 64B), and secures nearly worst-case MDP robustness.
{ (Tr.,, le,,)) _ to feed into the risk learner to predict adaptation risk on arbitrary task. Importantly, it introduces the pseudo task set T. with B > B and employ the acqui-
3. The resulting PDTS is easy-to-implenee to-implement and benefits
2

---

PDTS for Adaptive Decision-Makers in Randomized Environments
sition function A(-) to actively select the subset for next iteration. Throughout this work, we leave the exact opti- mization task batch size B fixed and same for all methods in both theoretical analysis and evaluation.
The adaptation capability in either few-shot (Wang al., 2022; Chi al., 2021) zero-shot is our primary focus to evaluate in RATS, Here, MPTS designs risk predictive model p(4|T, H1::: 0.) to surogate expensive evaluation, e.g., negative average rollout returns l of the policy 0g in a MDP r. Hence, we identify it as a method of RATS. In particular, the empirical evidence in (Wang et al., 2025) Fig. 5 validates its risk predictive model's feasibility of approximately scoring MDPs difficulties with high Pearson correlation coefficients between the model predictive ones and exact evaluation. Next, we overview its construction together with the acquisition function.
2.1. Task Robust Optimization Methods
Robustness (Carlini et al., 2019; Chi et al., 2024b) is en- tangled with risk minimization principles. And this part recaps those incorporated in DR and Meta-RL for robust decision- meeisi.n-making..
Expected/Empirical Risk Minimization (ERM). ERM originates from the statistical learning theory (Vapnik al., 1998). Such a principle minimizes the expectation of risk values under p(r):
Generative Modeling Adaptation Optimization. MPTS treats the optimization history as sequence generation and involves latent variables 2; to summarize batches of adapta- tion risk over iterations, which leads to:
min Bp(c) DI:
(I)
Ac
T
T
p(CoT zo.r (urr)) p(zo)
IPC 1e 0.) II p(zen1z()) t=
II
where p(T) is mostly a uniform distribution default.
(4)
with the evaluation risk batch 2T = ((TT,,, )
Group Distributionally Robust Risk Minimization (GDRM). Such a principle aims to boost the adaptation ro- bustness under certain proportional scenarios, e.g., a group of some under-sampled yet challenging tasks (Sagawa et al., 2019). In mathematics, we write the expression as:
With the Bayes rule and the streaming variational inference (Broderick et al., 2013; Nguyen et al., 2017) w.t.. Eq. (4), it obtains the approximate evidence lower bound of the risk learner to maximize in each batch:
max SELBO Φ) = (x( In ple (lo.l
min max po(Fr) oce SEB
(2)
BDxx [00(2) He) 9 (xil
(5)
with 9 to denote collection of groups over a task dataset. GDRM handles extremely worst subpopulation shifts (Koh et al., 2021) through a risk-reweighting mechanism.
with e the fixed conditioned prior from the last update and B E IR+ the penalty weight. Then the amortized evaluation of adaptation performance is approximated as p(E|T, Hirri 0z) s Bgg(s, (P ze; 0,)] through Monte Carlo estimates.
Distributionally Robust Risk Minimization (DRM). As a typical strategy DRM, CVaR, selects (1-a) proportional worst tasks after exact evaluation to optimize:
max (mb = In Pab Pep (H, H1:-1) NEW
(6a)
min CVaR,, : (p [( D; (3) 8E0
p(E 4,
Hy::: 0,)
MC,
(m(()), 0 (4)))
(6b)
J arg
Au
(6c)
J CT A
In Meta-RL, this corresponds to the hard MDP prioriti- zation (Greenberg et al., 2024; Lv al., 2024). With mingee lima-11 (O), degenerates to the worst- case optimization in (Collins al., 2020).
Optimization Pipeline. MPTS contains three critical steps in accordance with Eq. (6). ! In approximate posterior inference for Eq. (6)a, MPTS optimizes the risk predictive model through maximizing Eq. (5) with Hr; ② In amortized evaluation, samples J22 from p(r) and runs stochastic forward passes to estimate m(&,) and (44,) Vi e J in Eq. (6)b; 3 In subset selection, it picks up the subset T/t among Top-B acquisition scores from Ihf for next iteration optimization. By repeating 1-0) steps for -ruudds, MPTS derives the robust adaptive decision-maker.
2.2. RATS Preliminaries & MPTS Modules
Here, we specify RATS paradigm, which slightly differs from traditional active learning purposes (Cohn et al., 1996). RATS is mainly incorporated into risk-avers learning with traits: (i) active inference towards task difficulties with lim- ited cost, and (ii) acquisition rules to select the subset from the pseudo task set to optimize for robustness.
3

---

PDTS for Adaptive Decision-Makers in Randomized Environments
Acquisition Function in MPTS. The subset selection rule is built on the principle of optimism in the face of uncer-
Transition Dynamics. We describe the dynamical system p(0 (0,, 3++1 E P conditioned 0g and the action T with the transited state (after-adaptation) Oz+1;
tainty (OFU) (Auer, 2002b), and tasks with worse adap- tation performance and higher epistemic uncertainty are prioritized. This leads to the UCB principle, i.e., Au( (Auer, 2002a; Garivier & Moulines, 2008):
Reward Function. The step-wise reward quantifies adap- tation robustness improvement after state transitions. With CVaR,, risk-averse measure, we define it as RiO,, J (41) == CVaRa(O)) (O+1).
(T(T)) Σ rom((e)) + no(f), with T. E T2t1.
Given M in Fig. 2, the maximum iteration step T also corresponds to the total interaction rounds. The agent ac- tively selects J from the time-varying T/ and optimize the machine learner to increase adaptation robustness in T -'ruuuns.
where the mean m(()) Bas(zi(Mt)) p (4]T: zi) and the standard deviation (44) = v (o(2)111)) ( zi) of task-specific adaptation risk and are estimated through multiple stochastic forward passes, i.e., 2e ~ (2 and l (Cg T., z1). (1o: } are trade-off parameters.
Accordingly, the sequential actions are the outcome of a series of deterministic functions as the policy set ToT = [ TT-1 Here, the policy maps Ho. to the next subset from J for optimization. i.e., ne : Ho:: + + These ingredients can depict M in a probabilistic graph as Fig. 2.
3. Theoretical Investigations and Practical Enhancements
Formally, we can express the agent's ultimate goal as maxi- mizing the cumulative reward in a sequential manner,
This section first studies RATS with a task-selection MDP, proposes i-MABs as a versatile tool for RATS methods and establishes its connections with MPTS. To promote RATS's use in risk-averse decision-making, we analyze acquisition rules' influence on robustness concepts and present PDTS as an masy-to-implemen yet powerful scheme.
Tor arg max
Σ R(O., (H1). HoT--
(8)
Note that solving Eq. (8) is equivalent to reaching the opti- mal state 07. with lowest CVaR,(O), after repeating Tround optimization steps.
3.1. Enable Robust Active Task Sampling with i-MABs
We also denote the policy subset of an arbitrary intermediate decision-making sequence by Hgy == { with its optimal solution marked in * The associated cumulated return is Di R(O,, J conditioned on the starting state
When RATS meets decision-making, it involves scoring MDPs' difficulty from amortized evaluation, selecting a subset, and performing adaptive policy search in either a zero-shot or few-shot manner. The following introduces a theoretical tool to analyze these steps in RATS.
Hg.
Remark 3.1 (Bellman Optimality). For the studied T- horizon M, we write its Bellman optimality as:
T
8
((O.TT11) = max
Σ R(O., THi) Σ R(O., Jh
Hoo
00
01
Or
(9)
revealing that the optimal solution to the sub- uub-problem also reserves its global optimality. Hence, we can break the problem-solving into optimal subset selection in each round.
R
R
R
Figure 2. Task Robust Episodic Learning Task-Selection MDP.
From i-MABs to MPTS's Robustness Concept. Finding plausible strategies to solve M non-trivial as policy search is considered in discrete space with the varying action set A,. To this end, we further simplify M into infinite MAB (Mahajan & Teneketzis, 2008) to enable an online search of the optimal subset and interpret the optimization pipeline of MPTS as a special case.
Task Robust Episodic Learning as a MDP. The primary insight lies in a fnite-horizon Markov decision process (Put- erman, 2014), denoted by M =< S, A, P, R >. Here, we specifyy the essential components of M as:
State Space. M treats the feasible machine learner's 's pa- rameter, such as Meta-RL policies, as the reachable state, i.e., S = {0 E e):
Distinguished from the vanilla multi-armer bandit, M involves the state 0, and A441 induced by resampled J The arm corresponds to a feasible subset T E A--1- Hence, we can associate the reward distribution p( R|O, Iht with the chosen action TB As a result, one
Action Space. The action space constitutes collection of task subsets with cardinality constraints A, (J/² C
Tht with |T21 B) in RATS or CVaR,, methods;
4

---

PDTS for Adaptive Decision-Makers in Randomized Environments
Table 1. Comparison between MPTS and PDTS in Contributions.
Research Lens
Robustness Concept
Subset Selection
Acquisition Rule
MPTS (Wang et al., 2025)
Generative Modeling
Nearly CVaR
Max-Sum (Top-B)
UCB
PDTS (Ours)
MDP & i-MABs
Nearly Worst-Case
Max-Sum Diversity
Posterior Sampling
optimistic strategy for i-MABs is execute greedy search J arg max EA111 E [R]0,,T% in each round.
PDTS (Ours, 64x)
MPTS (BK, no, no_mixed)
MPT5 mixed)
NPT5 (8x. mised)
Sampling Frequancy
Walker
Walkerrddii
The regret in the i-MAB measures the performance differ- ence between the cumulated rewards of the optimal policy - and the actually executed policy I1o:T-1 over T- rounds:
Fruqaaiay
12.5
10.0
Selected
Chanoo Butur
-50
Aversge Roturn
1o
75
-75
20
5.0
2.5
-125
20
0.0
-150
200 Iteration
403
205
400
(Sentifior
Iteration
Regret(T, (or))
Σ t=
((11) - R(O,, [++1
Figure 3. MPTS's Performance Collapse with Greater B. We report the performance collapses of MPTS on Walker2dVel in the case B- = 8B. The task sampling frequency reveals the presence of the concentration issue.
(10)
We can further simplify the definition expression as Regret(T, [1: = CVaR, (0r) - CVaR,, which quantifies the performance gap between the optimal state and the practical state under (1 a) -ttai robustness.
interactions and policy optimization in Meta-RL DR.
Proposition 3.2 (MPTS a UCB-guided Solution to i-MABs). Executing MPTS pipeline in Eq. (6) is equiv- alent to approximately solving M with the i-MAB under the UCB principle.
Unfortunately, MPTS might encounter performance col- lapse with greater B, as indicated in Fig. 3 when B = 8B without a remedy. This empirical result can be attributed to the selected subset's concentration a narrow range, which motivates the proposal of heuristic tricks in Meta-RL to adopt mixed sampling strategy to ensure sufficient visita- tions to the whole task space (Wang al., 2025; Greenberg et al., 2024). Although these heuristics mitigate the issue, they do not provide a complete solution, as performance degradation persists with increasing B.
Consequently, our i-MAB is a theoretical model for induc- ing RATS methods, and MPTS can be viewed as a special case in Proposition 3.2.
3.2. Acquisition Functions Matter in Improving Coverage & Boosting Robustness
Several maciiion--makiggaee scenarios, e.g., robotics, are risk- averse, which makes worst-case optimization more advan- tageous (Greenberg al., 2024). However, (i) the search scope to worst cases is restricted by the batch size, and (ii) minimax optimization requires carefully designed relaxation in the field (Collins et al., 2020; Sagawa et al., 2019).
Theoretical Diagnosis of Sealed Exploration Potential in MPTS. In practice, we propose to circumvent complicated heuristies and retain the simplicity to enable sufficient explo- ration in i-MABs. To this end, analyze the concentration issue and report the theoretical analysis as Proposition 3.3.
Proposition 3.3 (Concentration Issue in Average Top-B Selection). Lei I(T): IRd - R be a unimodal and con- tinuous function, where d E N and t E RR., with maximum value f(T") at T" We uniformly sample a set of points B {T.) where To are i.id. with probability Pe falling within e-neighborhood T* as [(T')) S e. Following MPTS, we select the Top-B
Benefits of Enlarging B in Subset Selection. Note that the feasible subset is the arm in the i-MAB, enlarging B increases its number to Ar| = C. One promising trait of the risk predictive model in MPTS is to amortize the pol- icy evaluation in arbitrary MDP without exact interactions. Hence, greater B in implementation reserves at least two bonus: (i) it encourages exploration in the task space with more candidate subsets at no actual interaction cost; (ii) un- der high-risk prioritization rule, the optimization pipeline in RATS approximately executes CVaR,- - in each iteration,
- ( samples with the largest function values, i.e.,
J TTT- .f),
B,B E N+ B
h B-242 B+1
For any probability
0 such that p. <
the concentration
i.c., worst-case optimization with B + co. Notably, the ad- ditional computational overhead with larger B remains neg- ligible due to the efficiency of the risk predictive model-its cost is significantly lower than that of agent-environnen
P ((S(T)- - Se Vr E Tz))
increases with B and converges to with B - oc.
5

---

PDTS for Adaptive Decision-Makers in Randomized Environments
As implied, with the increase of B and fixed B, the Top-B operator tends to select the subset in small neighborhood of arg max+ f f(T), which corresponds to the subset concen- tration issue in MPTS, over-optimizes a local region and hampers task space exploration. Due to the use of Top-B operator in batch optimization, we also hypothesize a sim- ilar phenomenon in DRM (Wang et al., 2024e; Lv et al., 2024). Next, we will propose a natural plausible mechanism to enlarge B in RATS without suffering concentration pains.
One Forward Pass: 2t (0 (12a)
l ~ (e((t)) Vi (I...., B] (12b)
J = arg max Ap(TT) + [f T,)M (12c)
5" 1=
As implied in (Russo & Van Roy, 2014), posterior sampling avoids over-exploitation of inaccurate estimated uncertainty, and benefits more from stochasticity. Together with diver- sity regularization, we obtain PDTS as a new RATS method in Eq. (12). Intuitively, diversity penalty and posterior sam- pling suppress inaccurate Top-BB operations and synergize exploration in the task space. PDTS enjoys implementation simplicity, computational efficiency and stochastic optimism in decision-making
Diversity Regularization & Robustness Concept. Our strategy is to encourage the coverage of the task space dur- ing subset selection. Specifically, rather than selecting in- dividual candidate tasks based on their acquisition score, we evaluate task batches based on both adaptation risk and task diversities in the subset. This formulation constructs diversity maximization problem (Wang et al., 2023b):
max J |
A(T²) + 7 [d((T, T3)
(11)
3.4. Overall Implementation
where 8 measures the diversity of the subset T from iden- tifiers' pairwise distances, e.g., Ei, ||T,
Putting the above modules together, we present PDTS in Algorithm I, which is easily wrapped into DR and Meta-RL (See Algorithm 2 and 4).
Though the involvement of the regularization term makes the subset selection problem NP-hard, it can be solved using simple approximate algorithms (Borodin et al., 2017; Wang al., 2023b).
1 Posterior-Diversiy Synergized Task Sampling Input : Task distribution p(T); Task batch size B; Candi- date batch size B; Latest updated { $}: Latest history Hę-1; Iteration number K; Learning rate
Algorithm
Proposition 3.4 (Nearly Worst-Case Optimization with PDTS). When B grows large enough, optimizing the subset from Eq. (11) achieves nearly worst-case optimization.
A2-
Output::elee task identifier batch {T..)=)- // Optimize ze Risk Predictive Module for i = 1 to K do Perform gradient updates given H;-1:
Proposition 3,4 indicates that the regularized acquisition rule in Eq. (11) enables exploration of the worst subset from numerous candidate arms while retaining the subset's coverage of the task space to a certain level. Compared with (Collins et al., 2020; Sagawa et al., 2019), such a regularization will show its effectiveness in experiments.
& + -+ A2 SELBOO b,, Φ) in Eq. (5); 3 + s + da V SELBO (tb, Φ) in Eq. (5);
end
// Simulating After-Adaptation Results 1 Randomly sample { from P(T);
3.3. Practical Sampling with Stochastic Optimism
// Posterior Sampling Outcome
So far, we have presented the i-MAB as a theoretical tool for RATS and interpreted MPTS the UCB-guided solu- tion. However, the UCB acquisition rule requires multiple stochastic forward passes for each task's evaluation, and computations grow with B. It also demands calibration of exploration and exploitation weights in subset search.
Amortized evaluation ( with one stochastic forward
pass for all candidate tasks through executing Eq. (12)a-b;
// Diversity-Guidedd Subset Search
Run approximate algorithms to solve Eq. (12)c; Return the screened subset {
1B for next iteration.
To cut off unnecessary computations and retain the uncer- tainty optimism, we adopt the posterior sampling strategy as the acquisition principle for RATS. The posterior sam- pling (Osband et al., 2013; Asmuth et al., 2012) is exten- sion of Thompson sampling (Thompson, 1933) in solving MDPs with infinite actions. The reward or action value is treated as a randomized function, and each arm's 's value, sampled from the posterior once, serves action selection, i... AP(TT) = Di1 f in Eq. (12).
4. Experiments
This section presents experimental results on Meta-RL and robotics DR involving randomized physics properties.
Risk-Averse Baselines. For fair comparison, we retain the standard setup in (Wang et al., 2025) and use SOTA baselines as described in Section 2.1: ERM (Vapnik et al., 1998), DRM (Wang et al., 2024c; Greenberg al., 2024),
6

---

PDTS for Adaptive Decision-Makers in Randomized Environments
PDTS (Ours)
MPT5
ERM
GDRM
DRM
ReacherPos
Walker2dVel
Walker2aMMassel
HalfCheetahMassel
CVARO.8 Return
-20
-50
175
200
-40
-50
-100
-25
-60
100
250
150
-275
200 Iteration
400
0
200 Iteration
400
200
400
200 Iteration
400
Iteraaioon
Average Return
60
40
120
-20
40
-140
-40
20
160
180
-60
-20
-20
-20
200
400
200
400
200 Iteration
400
2000
400
Keration
iteration
Iteration
Meta Mest-Testing Return
-20
20
-30
160
-40
20
-20
180
-40
40
-30
60
60
200
-60
CVaRo.5
CVaRo.7
CVaR0.9
CVaRO.5
CVaRD.7
CVaR0.9
CVaR0.5
CVaR0.7
CVaRo.9
CVaRo.5
CVa
CVaRD.9
Figure 4. Meta-RL Results. The top depicts the cumulative return curves for CVaRp.9 validation MDPs during meta veta-training, the middle
shows the average cumulative returns curves during meta-training; and the bottom presents the meta-testing results with various a.
GDRM (Sagawa et al., 2019), and MPTS. Following im- plementations in (Wang et al., 2025; Tao et al., 2024), we use MAML (Finn et al., 2017) as the backbone algorithm in Meta-RL and employ TD3 (Fujimoto et al., 2018) and
tably, PDTS's performance advantage increases with higher a values. In the extreme scenario CVaRo.9, PDTS outper- forms others by more than 15% on all benchmarks except HalfCheetahMass-iil
PPO (Schulman et al., 2017) for domain randomization, respectively.
PDTS is agnostic to the meta-learning backbone. PDTS is a plug-and and-play module agnostic to the meta-learning backbone. Here, we integrate PDTS with PEARL (Rakelly et al., 2019) and compare it with MPTS and RoML (Green- berg al., 2024), built on the same backbone. As shown in Table 2, evaluate these methods on two Meta-RL scenar- ios from ROML. PDTS outperforms others in terms of CVaR return, further examining its effectiveness and scalability in risk-averse decision -aaiii.....iin.....
We run each algorithm on seven random seeds and report the average performance with standard error of means. For adaptation robustness, we compute CVaR, values across the validation tasks, with {0.9, 0.7, 0.5). and also report some -ittr---distiittiionn (OOD) results. Importantly, the pseudo batch size is set to B = 64B as default for PDTS.
4.1. Meta Reinforcement Learning
Table 2. PDTS's 's compatibility with PEARL backhone. Baselines are MPTS and RoML (Greenberg et al., 2024).
We consider Meta-RI. continuous control scenarios based on MuJoCo (Todorov et al., 2012): Reacher, Walker2d, and HalfCheetah. These include identifiers: target position, velocity, and mass (Wang et al., 2025).
CVaRa Return
PDTS
MPTS
RoML
PEARL
HalfChectahBody
993+226
945+226
855+35
847+44
HalfCheetahMMass
1296114
1209-445
1197:159
1118151
Fig. 4 reveals PDTS consistent superiority over others in CVaRu.9 validation return, demonstrating excellent adapta- tion robustness. Owing to intrinsic sampling randomness, PDTS and MPTS achieve average performance comparable to ERM, while DRM struggles to improve robustness and sacrifices more average returns. Surprisingly, benefiting from exploration, PDTS boosts both adaptation robustness and task efficiency on ReacherPos. We will further discuss this in Section& 4.2. Regarding meta-testing, PDTS excels across CVaR values, consistent with the learning curves. No-
4.2. Physical Robotics Domain Randomization
We evaluate PDTS on three DR scenarios introduced by Mehta et al. (2020); game scenario, LunarLander, and two robotic arm control scenarios, Pusher and EeghRe
Empirical findings in Fig. 5(a) are consistent with those in Meta-RL, where PDTS dominates the robustness perfor- mance. Notably, PDTS's ' advantage is more pronounced on Pusher and LunarL where the identifiers' dimension is lower than ErgoReacher's. 's. This is likely because low- dimensional task identifiers exacerbate the concentration
7

---

PDTS for Adaptive Decision-Makers in Randomized Environments
PDTS (Ours)
MPTS
ERM
GDRM
DRM
Pus
LunarLander
ErgoReacher
ODD tesing on LunarLander
96
300
CVARO. Return
-30
-60
200
75
70
2000
100
65
Testing
1000
60
-40
0.0
Agent Steps (x20
0.0
Agent Sleps (x109)
3.0
2.0
55
0.5
Agent Steps (x109)
1.0
Mainn Engine Strongth (OOD)
-47.5
300
Average Retumn
1.3x acco oration
Testing Bellum
-50.0
250
250
2.4k accelpration
200
88
200
-35.0
86
150
-57.
84
100
50
0.0
Agentt Steps (x10s)
0.5
10
0.0
0.S
Agent Steps (x10
3.0
2.0
0.0
0.5
Agent Steps (x10f)
1.0
2.0
Main Engine Strength (ID)
10
15
Testing Return
-50
300
90
85
Testing Return
275
300
-55
200
250
150
75
-45
70
200
Average CVARO.5
CVORO.7
CVaRo.9
50
Average
CNaRO.5 CVaRO.7
CVaR0.9
Average
CVARO.5
CV3RO.7
CVaRO.9
Main Engine Strength (OOD)
21
22
(a)
(b)
Figure 5. Physical Robotics DR Results. (a) The top shows the cumulative return curves for CVaRo.9 validation MDPs during training; the middle displays the average cumulative return curves across all validation MDPs during training: and the bottom presents the test results at various CVaR.. (b) We evaluate the trained policies both in- di-ttriibution (ID) and dfitt--diitiitutin (OOD) domains on
LunarLander, reporting the average returns for each sampled task.
issue, limiting the performance of MPTS. GDRM and DRM perform poorly in both average performance and robustness. Given fixed batch size, this weakness may stem from their limited exploration capacity. Regarding final testing per- formance, PDTS excels across CVaR values. Specifically, PDTS outperforms ERM by more than 8% on all bench- marks in CVaRa.9, and by as much as 73% on LunarLander.
4.3. Visual Robotics Domain Randomization
Vision-based robotics control is more challenging and un- derexplored in MPTS. With the latest robotics simulator, ManiSkill3 (Tao et al., 2024), we design two visual DR sce- narios: one randomizing lighting with a table-top two-finger gripper arm robot called LiftPegUprigh Light, and another randomizing goal locations with a quadruped robot called AnymalCReach Goal, as illustrated in Fig. 6(a).
PDTS improves task efficiency in specific scenarios. Be- yond adaptation robustness, PDTS shows the potential to ac- celerate training in specific scenarios. As shown in Fig. 5(a), PDTS achieves average returns comparable to ERM's best performance but with fewer training steps, achieving 2.4x acceleration in Pusher and 1.3x in LunarLander. We at- tribute this acceleration to the efficient exploration of acqui- sition criteria. Hence, PDTS holds promise for achieving robustness with reduced training steps in broader scenarios.
PDTS achieves best robustness performance in Fig. 6(b), same symbolic scenarios. Moreover, PDTS also exhibits trend of accelerated training in terms of average perfor- mance, highlighting its potential for realistic scenarios re- quiring both fast and robust adaptation. Due to insufficient task exploration, DRM performs poorly and probably gets trapped in challenging tasks. ERM obtains the worst fi- nal performance on AnymalCReach Goal. These findings underscore the importance of robust optimization with suffi- cient exploration, which may explain PDTS's s success.
PDTS achieves superior zero-shot policy adaptation in OOD MDPs. On Lundarlander, Fig. 5(b) witnesses PDTS's highest average test returns across sampled tasks. To evaluate zero-shot adaptation in OOD MDPs, we shift the identifier interval т e [4.0, 20.0] to the OOD range E [1.0, 4.0) U (20.0, 23.0]. Notably, PDTS exhibits the smallest performance degradation, particularly in the most challenging OOD tasks (T E [1.0,4.0).
PDTS can well discriminate task difficulties. We evaluate the Pearson Correlation Coefficient (PCC) between the pre- dicted episode returns and the exact values during training. As shown in Fig. 6(c), PDTS inherits MDPs' difficulty scor- ing capability of MPTS, with both PCC values greater than 0.5. In particular, PDTS exceeds MPTS in predicting accu- racies, which may be attributed to the extended exploration and the stochastic optimism in posterior sampling.
8

---

PDTS for Adaptive Decision-Makers in Randomized Environments
PDTS (Ours)
MPTS
ERM
GDRM
DRM
LinPegU pright Light
LitPegUprrghh_.oo Light
URPegUpright Light
LittPegUpright Light
9.6
CloR0.5 Rabo
0.4
1.0
0.5
0.2
0.2
0.2
0.0
Agent Steps (2105) AnymaiCReach Goal
Agent Steps Steps Anyma Rimaacc Goal
Gx10f;
Agent Steps (x10*)
AnymalCReach Goal
(c)
2:25
DRM
Averoge Success
CVORO.3 Success
0.4
Reletive Clock
2.75
0.4
0.3
3.50
0.3
1.25
01
1.00
0.0
0.0
2.2
Agent Steps (x109)
Agent Steps 6410%
Relative Hemory
(a]
(b)
Figure 6. Visual Robotics DR results. (a) Illustrations of two visual DR scenarios. (b) Curves of the average success ratio and the
CVaRo.5 success ratio on validation tasks during training. (c) Training curves of PCC values between predicted and true episode returns. (d) Memory cost and clock time relative to ERM during meta-traie stainiranning.
PDTS offers learning efficiency advantages at minimal cost. Like MPTS, PDTS avoids additional evaluation costs associated with interaction and rendering, making it more beneficial in vision-based RL tasks. From relative clock time and memory usage in Fig. 6(d), we find PDTS retains computational efficiency with others except DRM. As a result, PDTS robust and scalable approach for complex, interactioniinteniivy tasks.
bility, well integrating stronger robust optimization techniques into i-MAB frameworks to develop practical and scalable RATS methods.
Impact Statement
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal con- sequences of our work. including enhancing real-worke ap- plications of generalized robotics and other decision-makern systems. Although this study significantly improves adapta- tion robustness, adaptation failures remain possible, poten- tially leading to serious consequences real-world scenar- ios. Thus, human supervision remains essential.
5. Conclusion
Technical Discussions. This work studies RATS's s use for learning adaptive decision-makers in risk-averse decision- making. We present the i-MAB as a theoretical tool to enable RATS and establish theoretical connections with its latest method (Wang et al., 2025). Built these insights, we propose PDTS as a competitive RATS method to achieve nearly worst-case optimization. Extensive DR and Meta-RL experiments validate the effectiveness of RATS and show its efficiency potential in risk-averse scenarios.
Acknowledgments and Disclosure of Funding
This work is funded by National Natural Science Foun- dation of China (NSFC) with the Number # 62306326 and the National Key R&D Program of China under Grant 201SAAA0102801 We thank all reviewers for their positive comments and constructive suggestions in this work.
Limitations & Extensions. Our approach relies on the risk predictive model for roughly scoring task difficulties and leverages identifier information alongside the inherent smoothness of the adaptation risk function, though these assumptions may not always hold in restricted scenarios.
References
Abbas, M., Xiao, Q., Chen, L., Chen, P.-.., and Chen, T. Sharp-maml: Sharpness-aware model-agnostic meta learning, In International conference machine learn- ing, pp. 10-32, PMLR, 2022.
While robust active learning remains underexplored in se- quential decision-making, reducing computational and sam- ple expenses and improving robustness are crucial for build- ing large-scale decision models. The models and algorithms introduced in this work offer a tractable strategy with strong empirical performance toward achieving RATS goals. Fu- ture explorations can include designing more accurate risk- predictive models to enhance amortized evaluation relia-
Asmuth, J., Li, L., Littman, M. L., Nouri, A., and Wingate, D. A bayesian sampling approach to exploration in re- inforcement learning. arXiv preprint arXiv: 1205. 2et..220...... 2012.

---

PDTS for Adaptive Decision-Makers in Randomized Environments
Auer, P. Finite-time analysis of the multiarmed bandit prob- lem, 2002a.
Chow, Y., Ghavamzadeh, M., Janson, L.,, and Pavone, M. Risk-constrained reinforcement learning with percentile risk criteria. Journal of Machine Learning Research, 18 (167):1-51, 2018.
Auer, P. Using confidence bounds for exploitation- exploration trade-offs. Journal of Machine Learning Research, 3970v):3 2002b.
Chow, Y., Nachum, O., Faust, A., Dueñezz-Guzman, E., and Ghavamzadch, M. Safe policy learning for continuous control. Conference on Robot Learning, pp. 801-821. PMLR, 2021.
Beck, J.. Vuorio, R., Liu, E. Z., Xiong, Z., Zintgraf, L., Finn, C., and Whiteson, S. survey of rentrrreiffre learning, arXiv preprint arXiv:2301. 000008,............11.. 2023.
Cohn, D. A., Ghahramani, Z., and Jordan, M. I. Active learning with statistical models. Journal of artificial intelligence research, 4:129 4:129-145, 1996.
Borodin, A., Jain, A., Lee, H. C., and Ye, Y. Max-sum diver- sification, monotone submodular functions, and dynamic updates. ACM Transactions on Algorithms (TALG), (3):1-25, 2017.
Collins, L., Mokhtari, A., and Shakkottai, S. Task-robust model-agnostic meta-learning. Advances in Neural Infor- mation Processing Systems, 13:18880-1887, 2020.
Broderick, T., Boyd, N., Wibisono, A., Wilson, A. C., and Jordan, M. Streaming variational bayes. Advances in neural information processing systems, 26, 2013.
Coumans, E. Bullet physics simulation. In ACM SIG- GRAPH 2015 Courses, pp. 2015.
Carlini, N., Athalye, A.. Papernot, N., Brendel, W., Rauber, J., Tsipras, D., Goodfellow, I., Madry, A., and Kurakin, A. On evaluating adversarial robustness. arXiv preprint arXiv: 19Xi::9902.0675.. 2019.
Dennis, M., Jaques, N., Vinitsky, E., Bayen, A., Russell, S., Critch, A., and Levine, S. Emergent complexity and zero-shot transfer via unsupervised environment design. Advances in neural information processing systems, 33: 13049-13061, 2020.
Carpentier, A. and Valko, M. Simple regret for infinitely many armed bandits. International Conference on Machine Learning, pp. 1133-1141. PMLR, 2015.
Duan, Y., Schulman, J., Chen, X., Bartlett, L.,, Sutskever, I., and Abbeel, RI2: Fast reinforcement learn- ing via slow reinforcement learning. arXiv preprint 2016.
Carpin, S., Chow, Y.I.., and Pavone, M. Risk aversion in finite markov decision processes using total cost criteria and average value at risk. In 2016 ieee international conference robotics and aulomation (icra), pp. 335- 342. IEEE, 2016.
Finn, C., Abbeel, P., and Levine, S. Model-agnostic meta- learning for fast adaptation of deep networks. In Interna- tional conference on machine learning, pp. 1126-1135. PMLR, 2017.
Catto, E. Box2d: 2d physics engine for games, 2007.
URL http: httt:::/bbx2a0 http///xxddd.rr.goo
Chi, H., Liu, F., Yang, W., Lan, L.,, Liu, T., Han, B., Cheung, W., and Kwok, J. Tohan: A one-step approach towards few-shot hypothesis adaptation. Advances in Neural In- formation Processing Systems, 34:20970-20982. 2021.
Finn, C., Xu, K., and Levine, S. Probabilistic model- agnostic meta-learning. Advances in neural information processing systems, 31, 2018.
Fujimoto, S., Hoof, H., and Meger, D. Addressing function approximation error in actor-critic methods. In Interna- tional conference on machine learning, pp. 1587-1596. PMLR, 2018
Chi, H., Li, H., Yang. W., Liu, F., Lan, L., Ren, X., Liu, T., and Han, B. Unveiling causal reasoning in large language models: Reality or mirage? Advances in Neural Informa- fion Processing Systems, 37:96640-96670 2024a.
Chi, H., Yang, W., Liu, F, Lan, L., Qin, T., and Han, B. Does confusion really hurt novel class discovery? International Journal of Computer Vision, 132(8):311-322) 2024b.
Gal, Y., Islam, R., and Ghahramani, Z. Deep bayesian active learning with image data. In International conference on machine learning, pp. 1183-1192. PMLR, 2017.
Chow, Y. Risk-sensitive and data-drived sequential decision making. PhD thesis, Stanford University, 2017.
Garivier. A. and Moulines, E. On upper-confidence bound policies for non-stationary bandit problems. arXiv preprint arXiv:0805. 2008.
Chow, Y., Tamar, A., Mannor, S., and Pavone, M. Risk- sensitive and robust decision-maker making: a evar optimization approach. Advances in neural information processing systems, 28, 2015.
Garnelo, M., Schwarz, J., Rosenbaum, D., Viola, F,, Rezende, D. J., Eslami, S., and Teh, Y. W. Neural pro- cesses. arXiv preprint 1807.0000 2018.
10

---

PDTS for Adaptive Decision-Makers in Randomized Environments
Gondal, M. W., Gast, J., Ruiz, I. A., Droste, R., Macri, T., Kumar, S., and Staudigl, I.. Domain aligned clip for few-shot classification. In Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision, pp. 5721-5730, 2024.
Levine, S., Kumar, A., Tucker, G., and Fu, J. Offline rein- forcement learning: Tutorial, review, and perspectives open problems. arXiv preprint arXiv:2005.0164: 2020.
Li, L., Yang, R., and Luo, D. Focal: Efficient fully- offline meta-reinforcement learning via distance met- ric learning and behavior regularization. arXiv preprint arXiv:2010.01112, 2020.
Greenberg, I., Chow, Y., Ghavamzadeh, M., and Mannor, S. Efficient risk-averse reinforcement learning. Advances in Neural Information Processing Systems, 35:32639- 32652, 2022.
Linsmeier, T. J. and Pearson. N. D. Value at risk. Financial analysts journal, 56(2):47-67, 2000.
Greenberg, I., Mannor, S., Chechik, G., and Meirom, E. Train hard, fight easy: Robust meta reinforcement learn- ing. Advances in Neural Information Processing Systems, 36, 2024.
Lv, Y., Wang, C., Liang, D., and Xie, Z. Theoretical investi- gations and practical enhancements on tail task risk mini- mization in meta learning. In The Thirty--ighth Annual Conference on Neural Information Processing Systems, 2024. URL https: http:://oppnrevieww net/ forum? id=Mcrz000hwr.
Gupta, A., Mendonca, R., Liu, Y., Abbeel, P., and Levine, S. Meta-reiinooremmen learning of structured exploration strategies. Advances in neural information processing systems, 31, 2018.
Ma, Y. J., Liang, W., Wang, H.-J., Wang, S., Zhu, Y., Fan, L., Bastani, O., and Jayaraman, D. Dreureka: Lan- guage model guided sim-to-real transfer. arXiv preprint arXi::2406.01966 2024.
Hong, Z., Kumar, A., Karnik. S., Bhandwaldar, A., Sri- vastava, A., Pajarinen, J., Laroche, R., Gupta, A., and Agrawal, P. Beyond uniform sampling: Offline reinforce- ment learning with imbalanced datasets. Advances in Neural Information Processing Systems, 36:4985-5000, 2023.
Mahajan, A. and Teneketzis, D. Multi-armed bandit prob- lems. In Foundations and applications of sensor manage- ment, pp. 121-151. Springer, 2008.
Hospedales, T., Antoniou, A., Micaelli, P., and Storkey, A. Meta-learning in neural networks: A survey. IEEE transactions pattern analysis and machine intelligence, 44(9):5149-5169 2021.
Mao, Y., Zhang, H., Chen, C., Xu, Y., and Ji, X. Supported value regularization for offline reinforcement learning. Advances in Neural Information Processing Systems, 36: 40587-40609, 2023a.
Jiang, M., Dennis, M., Parker-Holder J., Foerster, J., Grefen- stette, E., and Rocktäschel, T. Replay-guided adversarial environment design. Advances in Neural Information Processing Systems, 1884-18777 2021.
Mao, Y., Zhang, H., Chen, C., Xu, Y., and Ji, X. Supported trust region optimization for offline reinforcement learn- ing. International Conference on Machine Learning, pp. 23829 28829-23851, PMLR, 2023b.
Kingma, D. P. and Welling, M. Auto-cncoding variational bayes. arXiv preprint arXiv:1312.6114,, 2013.
Mao, Y., Wang, Q., Chen, C., Qu, Y., and Ji, X. Offline reinforcement learning with ood state correction and ood action suppression. arXiv preprint 10.009000 2024a.
Kirsch, A.,
Van Amersfoort, J., and Gal, Y. Batchbald: Ef- ficient and diverse batch acquisition for deep bayesian active learning. Advances in neural information process- ing systems, 32, 2019.
Mao, Y., Wang, Q., Qu, Y., Jiang, Y., and Ji, X. Doubly mild generalization for offline reinforcement learning, arXiv preprint urXiv:2411.079.. 2024b.
Koh, P. W., Sagawa, S., Marklund, H., Xie, S. M., Zhang. M., Balsubramani, A., Hu, W., Yasunaga, M., Phillips, R. L. Gao, I., et al. Wilds: A benchmark of in-the- wild distribution shifts. In International conference on machine learning, pp. 5637-5664. PMLR, 2021.
Mchta, B., Diaz, M., Golemo, F., Pal, J., and Paull, L. Active domain randomization. In Conference on Robot Learning, pp. 1162-1176. PMLR, 2020.
Mukhoti, J., Kirsch, A., van Amersfoort, J., Torr, P. H., and Gal, Y. Deep deterministic uncertainty: A new simple baseline. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 24384 24394, 2023.
Koprulu, C., Simão, T. D., Jansen, N., and Topcu, U. Risk- aware curriculum generation for heavy-tailed task dis- tributions. In Uncertainty in Artificial Intelligence, pp. 1132-1142, PMLR, 2023.

---

PDTS for Adaptive Decision-Makers in Randomized Environments
Muratore, F., Treede, F., Gienger, M., and Peters, J. Domain randomization for simulation-based policy optimization with transferability assessment. In Conference on Robot Learning, Pp. 700-713. PMLR, 2018.
Rempe, D., Philion, J., Guibas, L. J., Fidler, S., and Litany, O. Generating useful accident-promt driving scenarios via a learned traffic prior. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 17305-17315, 2022.
Muratore, F., Eilers, C.. Gienger, M., and Peters, J. Data- efficient domain randomization with bayesian optimiza- tion. IEEE Robotics and Automation Letters, 6(2)9111- 918, 2021.
Ren, P., Xiao, Y., Chang, X., Huang, P.-Y., Li, Z., Gupta, B. B., Chen, X., and Wang. X. A survey deep active learning. ACM computing surveys (CSUR), 54(9):1-44) 2021.
Nguyen, C. V. Li, Y., Bui, I. D., and Turner, R. E. Variational continual learning. arXiv preprint arXiv: 10628, 2017.
Rigter, M., Lacerda, B., and Hawes, N. Risk-averse bayes- adaptive reinforcement learning. Advances in Neural Information Processing Systems, 34:1 142-1144, 2021.
Osband, I., Russo, D., and Van Roy, B. (more) efficient reinforcement learning via posterior sampling. Advances in Neural Information Processing Systems, 26, 2013.
Ritter, S., Wang, J., Kurth-Nelson, Z., Jayakumar, S., Blun- dell, C., Pascanu, R., and Botvinick, M. Been there, done that: Meta-learning with episodic recall. International conference on machine learning, pp. 4354 4363. PMLR, 2018.
Pan, X., Seita, D., Gao, Y., and Canny, J. Risk averse robust adversarial reinforcement learning. In 2019 Inter- national Conference on Robotics and Automation (ICRA),
pp. 8522 8528. IEEE, 2019.
Rockafellar, R. T., Uryasev, S., et al. Optimization of condi- tional value-at-risk. Journal of risk, 2:21 2000.
Puterman, M. L. Markov decision processes: discrete stochastic dynamic programming. John Wiley & Sons, 2014.
Russo, D. and Van Roy, B. Learning to optimize via poste- rior sampling. Mathematics of Operations Research, 39 (4):1221-1243, 2014.
Qi, Y., Ban, Y., Wei, T., Zou, J., Yao, H., and He, J. Meta- learning with neural bandit scheduler. Advances in Neural Information Processing Systems, 36, 2024.
Sagawa, S., Koh, W., Hashimoto, T. B., and Liang, P. Distributionally robust neural networks. In International Conference on Learning Representations, 2019.
Qu, Y., Wang, B., Shao, J., Jiang, Y., Chen, C., Ye, 7., Line, L., Feng, Y., Lai, L., Qin, H., et al. Hokoff: Real game dataset from honor of kings and its offline reinforcement learning benchmarks. Advances in Neural Information Processing Systems, 36:22166 2269229, 2023.
Schulman, J., Wolski, F., Dhariwal, P., Radford, A., and
Klimov, O. Proximal policy optimization algorithms. 2017.
arXiv preprint
Setlur, A., Garg, S., Smith, V., and Levine, S. Prompting is a double-edged sword: Improving worst-group robust- ness of foundation models. In Forty-first International Conference on Machine Learning, 2024.
Qu, Y., Wang, B., Jiang, Y., Shao, J., Mao, Y., Wang, C., Liu, C., and Ji, X. Choices are more important than efforts: Llm enables efficient multi-agent exploration. arXiv preprint ar arXiv:210002511,, 2024.
Shao, J., Qu, Y., Chen, C., Zhang, H., and Ji, X. Coun- terfactual conservative q learning for offline multi-agent reinforcement learning. Advances in Neural Information Processing Systems, 16.77290-77312, 2023a.
Qu, Y., Jiang, Y., Wang, B., Mao, Y., Wang, C., Liu, C., and Ji, X. Latent reward: Lim-empowered credit assignment in episodic reinforcement learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 39, pp. 20095-20103, 2025.
Shao, J., Zhang, H., Qu, Y., Liu, C., He, S., Jiang, Y., and Ji, X. Complementary attention for multi-agent reinforce- ment learning. In International Conference on Machine Learning, pp. 01777630779... PMLR, 2023b.
Rajeswaran, A., Finn, C., Kakade, S. M., and Levine, S. Meta-learning with implicit gradients. Advances in neural information processing systems, 32, 2019.
Sutton, R. S. and Barto, A. G. Reinforcement learning: An introduction. MIT press, 2018
Rakelly, K., Zhou, A., Finn, C., Levine, S., and Quillen, D. Efficient off policy meta-reinforcement learning via prob- abilistic context variables. In International conference on machine learning, pp. 5331-5340. PMLR, 2019.
Tamar, A., Chow, Y., Ghavamzadeh, M., and Mannor, Policy gradient for coherent risk measures. Advances in neural information processing systems, 28, 2015.
12

---

PDTS for Adaptive Decision-Makers in Randomized Environments
Tao, S., Xiang, F., Shukla, A., Qin, Y., Hinrichsen, X., Yuan, X., Bao, C., Lin, X., Liu, Y., Chan, T.-k., et al. Maniskill3; Gpu parallelized robotics simulation and rendering for generalizable embodied ai. arXiv preprint arXiv:2410.00425, 2024.
Wang, Q. Federici, M., and Hoof, H. Bridge the in- ference gaps of neural processes via expectation maxi- mization. In The Eleventh International Conference Learning Representations, 2023a.
Wang, Lv, Y., Xie, Z., Huang, J., et al. simple yet effective strategy robustify the meta learning paradigm. Advances in Neural Information Processing Systems, 36, 2024c.
Thompson, W. R. On the likelihood that one unknown probability exceeds another in view of the evidence of two samples. Biometrika, 25(3-4):285-22 2944, 1933.
Wang, Q. C., Xiao, Z., Mao, Y., Qu, Y., Shen, J., Lv, Y., and Ji, X. Model predictive task sampling for efficient and robust adaptation, 2025. URL https: //axivv. org/ / abs 22002 11039.
Tiboni, G., Klink, P., Peters, J., Tommasi, T., D' D'Eramo, C., and Chalvatzaki, G. Domain randomization via entropy maximization. arXiv preprint urXiv:2311.01885 2023.
Tobin, J., Fong, R., Ray, A., Schneider, J., Zaremba, W., and Abbeel, P. Domain randomization for transferring deep neural networks from simulation the real world. In 20/7 IEEE/RSI international conference on intelligent robots and systems (IROS), pp. 23-30. IEEE, 2017.
Wang. Y., Mathioudakis, M., Li, J., and Fabbri, F. Max- min diversification with fairness constraints: Exact and approximation algorithms. In Proceedings of the 2023 SIAM International Conference on Data Mining (SDM), pp. 91-99. SIAM, 2023b.
Todorov, E., Erez, T., and Tassa, Y. Mujoco: A physics engine for model-based control. In 2012 IEEE/RSJ inter- national conference on intelligent robots and systems, pp. 5026-5033. IEEE, 2012.
Wu, J., Chen, J., and Huang, D. Entropy-based active learning for object detection with progressive diversity constraint. In Proceedings of the IEEF/CVF Conference on Computer Vision and Pattern Recognition, pp. 9397- 9406, 2022.
Vapnik, V. N., Vapnik, V., et al. Statistical learning theory. 1998.
Xu, F., Jiang, S., Yin, H., Zhang, Z., Yu, Li, D., and Liu, W. Enhancing context-based metaareenforreement learning algorithms via an efficient task encoder (student abstract). Proceedings g of the AAAI Conference on Ar- tificial Intelligence, volume 35, pp. 15937-15938, 2021.
Wang, B., Qu, Y., Jiang, Y., Shao, J., Liu, C., Yang, W., and Ji, X. Lim-empowered state representation for re- inforcement learning, arXiv preprint 2rXiv:2407.13237, 2024a.
Yoon, J., Kim, T., Dia, O., Kim, S., Bengio, Y., and Ahn, S. Bayesian model-agnostic meta-learning Advances in neural information processing systems, 31, 2018.
Wang, C., Lv, Y., Mao, Y., Qu, Y., Xu, Y., and Ji, X. Robust fast adaptation from adversarially explicit task distribu- tion generation, arXiv preprint arXiv:2407. 19523,, 2024b.
Zhang, H., Mao, Y., Wang, B., He, S., Xu, Y., and Ji, X. In in-sample actor critic for offline reinforcement learning. In The Eleventh International Conference on Learning Representations, 2023.
Wang, Q. and Van Hoof, H. Doubly stochastic variational in- ference for neural processes with hierarchical latent vari- ables. In International Conference on Machine Learning, pp.
PMLR, 2020.
Zhu, X., Lafferty, J., and Ghahramani, Z. Combining ac- tive learning and semi-supervised learning using gaussian fields and harmonic functions. In ICML 2003 workshop on the continuum from labeled to unlabeled data in ma- chine learning and data mining, volume 3, pp. 58-65, 2003.
Wang, Q. and Van Hoof, H. Learning expressive meta- representations with mixture of expert neural processes. Advances in neural information processing systems, 35: 26242-26255, 2022a.
Wang, Q. and Van Hoof, H. Model-based meta reinforce- ment learning using graph structured surrogate models and amortized policy search. In International Conference on Machine Learning, pp. 23055-23077 PMLR, 2022b.
Zintgraf, I.., Shiarlis, K., Igl, M., Schulze, S., Gal, Y., Hof- mann, K., and Whiteson, S. Varibad: A very good method for adaative deep rl via meta-learning. arXiv preprint arXiv:1910.083488 2019.
Wang, Q. Federici, M., and van Hoof, H. Bridge the in- ference gaps of neural processes via expectation maxi- mization. In The Eleventh International Conference on Learning Representations, 2022.
13

---

PDTS for Adaptive Decision-Makers in Randomized Environments
A. Adaptive Decision-Making in Randomized Environments
A.I. Zero-Shot and Few-Shot Sequential Decision-Making
In this part, we include typical MDPs and distributions, which align with randomized environments for sequential decision- making. Meanwhile, we show examples of DR and Meta-RL this context, which can well specify the form of the task-specific adaptation risk & in zero-shot and few-shot setup.
Definition A.1 (ddnttffierInduuee MDP Distribution). We consider a distribution over MDPs p(M,) induced by the identifier distribution p(r). Here, the MDP as the sequential Pecision-makingg environment can be characterized as M, =< S. R,, where S and A are the state and the action space shared across all MDPs.
MDPs vary in either the transition dynamics P- (8t, az) or the both, specified by corresponding identifiers T. The H-horizon rollout under policy is R(sH-1, aH-1), with its cumulative reward Ei=o y'R(s;, ai) and ry R (so. ag. R(so, an).... ao).
P+(81+1 ( a.) or the reward function R,
With Definition 1. can provide some instantiations of DR and Meta-RI. under specific policy optimization backbones.
Me with Model Agnostic Meta Learning. Meta-RL seeks to train agents capable of rapidly adapting to new tasks or environments by utilizing knowledge from related tasks. We adopt MAML (Finn et al., 2017) as the standard backbone algorithm due to its widespread application in addressing few-shot sequential decision-making problems. This also retains the same configuration in MPTS (Wang et al., 2025). The optimization objective of MAML is mathematically defined as:
min minEppr Ep(r) [ A (D); BE
(13a)
H Σ -==0
Ertre 1)]
min Ep(r) E
E sa ,
0' = - AVo rre
-E DF
Σ rrr
(13b)
where D, denotes episodes collected from MDPs specified by T under either the meta-policy or the fast-adapted policy. The term inside the brackets specifies the adaptation risk E(D? D. 0). which represents either the negative cumulative returns or the negative final rewards. The expression 0 - AV 0) indicates the gradient update for fast task adaptation, where A is the learning rate. Upon completion of meta-training, the resulting meta-policy A can be generalized across the task space.
Robotics Domain Randomization. Robotics domain randomization refers to a training paradigm in which an agent is trained across a collection of diverse environments to develop a generalizable policy. The diversity of these environments enhances the robustness of the resulting policy during deployment. Notably, this approach eliminates the need for few-shot episodes in unseen but similar environments. Mathematically, the optimization objective can be expressed as:
max {(0) e
R. By((r)
Σ Σx'r
(14)
where p(r) denotes the distribution over MDPs specified by t. and fre) represents the stepwise rewards obtained from interacting with a specific MDP, with H as the horizon.
After solving the optimization problem in Eq. (14), the resulting policy mo serves as zero-shot decision-maker in new 2'r. For physical domain
environments. In this context, the adaptation risk can be expressed as E(D?, D 0) randomization, policy optimization utilizes the TD3 algorithm (Fujimoto et al., 2018), an off-policy method known for its sample efficiency and stability. For visual domain randomization, we employ the PPO algorithm (Schulman et al., 2017), an on-policy method recommended by ManiSkill3 (Tao et al., 2024).
A.2. Pseudo Algorithm in Meta-RL and DR
In this part, we show how PDTS is incorporated in DR and Meta-RL. These are Algorithm 2/3 and Algorithm 4/5.
14

---

PDTS for Adaptive Decision-Makers in Randomized Environments
Algorithm 2 PDTS for Domain Randomization (Zero-Shot
Algorithm 3 Posterior-Diversity Synergized Task Sampling
Scenarios)
Input : Task distribution P(T); Task batch size B; Candi- date batch size B; Latest updated {tb, Φ); Latest history H;-1; Iteration number K; Learning rate
Input :Takk distribution p(r); Task batch size B; Learn-
ing rate A1.
:Adapteedapted policy 0.
22-
Set the initial iteration number / 1;
Output Selected task identifier batch {T. // Optimize Risk Predictive Module for i 1 to K do Perform gradient updates given H;-1:
Randomly initialize policy 0:
Randomly initialize risk learner {+,, #3;
while not converged do
Execute Algorithm 3 to access the task batch (T,
Φ -- Φ + X2 SELBOO sb, Φ) in Eq. (5); s + do V &ELBO sb Φ) in Eq. (5);
Sample trajectories for each task with policy 0 to induce (DR,1) J
end
1/ Eval Adaptation Performance
// Simulating After-Adaptation Results
Compute the task specific adaptation risk (4,5 1=
Randomly sample (fr, j=z from p(r);
-E no.2% (2 rr 1
// Posterior Sampling Outcome
Amortized evaluation ( with one stochastic forward
Return He = [[Tr,; f, as the Input to Algorithm
pass for all candidate tasks through executing Eq. (12)a-b;
3:
// Update Policy
// Diversity-Guided Subset Search
Run approximate algorithms to solve Eq. (12)c;
Perform batch gradient updates:
Return the screened subset {Te,) 2 for next iteration.
0++1
+ 0. A
1 E Voli,i
Update the iteration number: t + + + 1;
end
Algorithm 4 PDTS for Model Agnostic Meta Learning
Algorithm 5 Posterior-Diversity Synergized Task Sampling
(Few-Shot Scenarios: Meta-RL, Sinusoid Regression)
Input : Task distribution P(T); Task batch size B; Candi- date batch size B; Latest updated {ı, Φ} Latest history H;-1; Iteration number K; Learning rate
Input Task distribution p(r); Task batch size B; Learn-
ing rates: f
Output: Meta-trained initialization m Set the initial iteration number t = 1;
X2-
Output: :eeected task identifier batch {T..) // Optimizo ze Risk Predictive Module for i= 1 to K do Perform gradient updates given H4-1:
Randomly initialize meta policy m Randomly initialize risk learner {+,, Φ); while not converged do
Execute Algorithm
access the batch ( Sample trajectories for each task with policy gmeta to induce {DS (D.,JE)) 1
Φ Φ A2 S tb, Φ) in Eq. (5); * - +b + A2 V wH (//, Φ) in Eq. (5);
end
// Inner Loop to Fast Adapt for i 1 to K do Compute the task-speeific adaptation
// Simulating After-Adaptation Results
Randomly sample ] from p(T);
risk:
// Posterior Sampling Outcome
((S 7,, 0) = Exo.DE DE,, o
Amortized evaluation { with one stochastic forward pass for all candidate tasks through executing Eq. (12)a-b;
Perform gradient updates as fast adaptation:
0: + 0m \1,1 V (DD
// Diversity-Guidee
Subset Search
Sample trajectories D2 with policy 0 for task T.:
Run approximate algorithms to solve Eq. (12)c;
end
Return the screened subset ) for next iteration.
// Outer Loop to Meta-train Evaluate fast adaptation performance {he
D Tap Te,s [Σ# rr 12::
-E
Return H. = ([Tr,, 4, as the Input to Algorithm 5;
Perform meta initialization updates:
@ure + 0
Σ Vol.,
3
Update the iteration number: & +- ++ 1;
end
15

---

PDTS for Adaptive Decision-Makers in Randomized Environments
A.3. Related Work
Risk-Averse Methods and Robustness Evaluation. Real-world decision-making scenarios are risk-sensitive (Chow et al., 2021), and this makes robust optimization indispensable component of policy optimization (Tamar et al., 2015; Chow et al., 2015; Chow, 2017; Chow et al., 2018; Pan et al., 2019; Rigter et al., 2021: Greenberg et al., 2022). Recent advances turn to some risk-averse measures to improve robustness in the task distribution. These are detailed in Section 2.1, which includes the DRM (Lv et al., 2024; Wang et al., 2024c), GDRM (Sagawa et al., 2019), and MPTS (Wang et al., 2025). The existing worst-case optimization in meta-learning is (Collins et al., 2020), which relies on special relaxation tricks; otherwise, the optimization can be unstable. Besides, other inspiring robust methods in task sampling (Qi et al., 2024) have not been applied in DR or Meta-RL. As MPTS is the latest SOTA RATS method, we mainly compared PDTS with it in experiments. RoML (Greenberg et al., 2024) also considers the CVaR optimization in Meta-RL.; hence, we include some comparisons given the PEARL (Rakelly al., 2019) backbone. In terms of evaluation, the subpopulation shift and out-of- -uttrf-disttiiution scenarios are commonly used in the field (Koh et al., 2021).
Cross-Task Adaptation in Sequential Decision-Making. This work focuses on the zero-shot and few-shot decision-making scenarios. zero-shot decision maaiing. the primary technique is to randomize the environments and train the agent in the distribution over environments, which is called domain randomization (Tobin et al., 2017; Muratore et al., 2018; Mehta et al., 2020; Muratore et al., 2021; Tiboni et al., 2023). Meta-learning is promising paradigm for achieving few-shot adaptation to unseen tasks, avoiding learning from scratch (Hospedales et al., 2021). The secret behind its few-shot adaptation capability is to leverage past experience and consolidate these as the prior for fast moblemm-solving There are three primary types of policy adaptation methods that can be seamlessly incorporated into deep RL scenarios. (i) The optimization based methods aim to seek robust meta-initializztion of the model that can be adapted to unseen scenarios through fine-tuning (Finn et al., 2017; 2018; Abbas et al., 2022; Rajeswaran et al., 2019; Yoon et al., 2018; Gupta et al., 2018; Qi et al., 2024). (ii) The context-based methods mostly adopt an encoder and decoder structure policy networks (Xu et al., 2021; Wang & Van Hoof, 2022a;b; Rakelly et al., 2019; Zintgraf al., 2019; Li et al., 2020). The encoder summarizes the support trajectories into latent variable to guide the policy adaptation. (ii) The recurrent methods mainly employ a recurrent neural network to encode the sequential All of these zero-shot or few-shot learning is performed in task episodic way. which means that a batch of MDPs are resampled in each iteration to train the adaptive policy during DR and Meta-RL.
betisionnmmking episodes as the adaptation prior (Duan al., 2016; Ritter et al., 2018).
Curriculum learning is also a crucial topic related to adaptive decision makiin.. Dennis et al. (2020) develop unsupervised environment design (UED) as a novel paradigm for environment distribution generation and achieve SOTA zero-shot transfer. Jiang et al. (2021) cast prioritized level replay enhance UED and formulate dual curriculum design for improving OOD and zero-shot performance. In (Koprulu et al., 2023), heavy-tailed distributions are incorporated into the automated curriculum, which leads to robustness improvement. Wang et al. (2024b) propose to generate task distributions for meta-RL through adversarial training normalizing flows, which provides a data-driven experimental design pipeline for increasing adaptation robustness. In contrast, our work emphasizes robust task adaptation under a fixed task distribution. Integrating the idea of surrogate evaluation from PDTS into curriculum design could be an interesting direction for future research.
A.4. CVaR, as Risk-Averse Metrics
Definition A.2 (CVaR,)). With 0-parameterized model, e.g., a deep RL policy, we can induce a random variable & := {D? T. D 0) from p(T). Then, we can define the cumulative adaptation risk distribution and its quantile by F(E) and 8o mine ( 2 at. Following (Rockafellar et al., 2000), CVaR at a-level robustness can be expressed as:
CVaR. ((a0)))00)000)))
(AF)
(15)
Accordingly, the normalized tail risk task distribution is
0,
1 1< &a
Fa(E: 0)
(16)
F(E:0)aa 1-a
1 t
with (T; as its probability density function.
Assumption 1 (Lipschitzz Continuity). The adaptation risk function ((::0) is Lipshhttt continuous w.r.t. 0 and Bo-
16

---

PDTS for Adaptive Decision-Makers in Randomized Environments
Lipschitz w.r.t. T, i.e.,
I(D9, - ()))) 0 - 0'1 ((D? D.; ((D), 25; - BellT T/1,
(17)
where v10,01 E + and "(,T') E J.
Assumption 2 (Bounded Functions). Given arbitrary
O and the task T E J, the adaptation risk function {(: 0) satisfies:
max TET
< m
(18)
CVa is a risk-averse measure in computational finance and management science (Linsmeier & Pearson, 2000; Rockafellar et al., 2000), and we write it in the Definition A.2. As for Assumptions and 2, these are commonly seen analysis (Lv et al.. 2024; Wang et al., 2024c) and are also necessary to MPTS and PDTS. These constitute the predictability of the adaptation risk value over iterations the relative task difficulty remains invariant after the model's parameter perturbed a bit.
A.5. Details in Risk Predictive Modules
As introduced Section 2.2, a key feature of RATS is its use of risk predictive models, e.g., p(((T, Hi.1; 0s), as surrogates for expensive evaluations. To the best of our knowledge, MPTS (Wang et al., 2025) is the first develop such model for robust optimization. As introduced in Section 2.2, MPTS leverages generative modeling for adaptation optimization and employs approximate posterior inference to forecast adaptation risk values after one-step optimization. Since small deviations in the machine learner's 's parameters do not alter the relative difficulty scores within the task batch, this provides tractable approach to coarse-grained task difficulty evaluation. The probabilistic graphical model is provided in Fig. 7.
Empirically, series of evaluation on DR and Meta-RL in (Wang et al., 2025) has validated the plausibility of such a schema. Particularly, it achieves high Pearson correlation coefficient (PCC) values between the risk learnrr' predicted adaptation risk values, (c++1, (x H.) (pts (f) (f|T2+1,i H= and the corresponding exact adaptation risk values (f++1,)) indicating the risk predictive model's ability to score difficulty. The Pearson correlation coefficient value is computed as Ex (Er41, Mean[( H0 Mean((fe41.)))
Př.e"
V EE. 1. (eaneee Mean((fe41. HDI HD2 E Mean (f.41..))
0
0r
21
ZT
Lo
tork.....
(-A....6
To,5
T1,1
Figure The probabilistic graphical model the risk predictive model in (Wang et al., 2025). where gray units denote observed variables with the white unobservable ones. The solid directed lines depict the generative model, and the dashed directed lines indicate the recognition model and approximate inference (Kingma & Welling, 2013).
Formulation of ELBO & Stochastic Gradient Estimates For the sake of easy and fast understanding of PDTS for the reviewer, we include the risk predictive model part as follows. These are modified from the MPTS team's open- sourced manuscript in (Wang et al., 2025). Here, the optimization objective of the risk predictive model is in Eq. (5) The risk predictive model uses a latent variable to summarize historical information and quantify uncertainty in predicting
17

---

PDTS for Adaptive Decision-Makers in Randomized Environments
task-specific adaptation risk. The following outlines the steps for deriving the evidence lower bound in optimization.
[/ Pv H1 (
(L)( == In pup Hi H14-1) = In
(19a)
=Im[/
I/ gol (0(
p(zz) 111-111)
Pvb (1g )z)
(19b)
(0 (H))))H)))
(He (z. In
DE (00 p(z)l H1:-1) = GELBO u)
(19c)
Then, the ELBO is derived with the help of the reparameterizatioo trick (Kingma & Welling, 2013). And Wang et al. (2025) further relax the ELBO to derive the 3-VAE version, which corresponds to Eq. (5) this work.
Neural Architecture of the Risk Predictive Model. For practicality, stability and fair comparison, we adopt the same risk predictive model design proposed in MPTS (Wang et al., 2025). As shown in Fig. 8, the risk learner follows an encoder-decoder structure. For consistency across benchmarks,
employ the same neural architecture similar with that of neural processes (Garnelo et al., 2018; Wang et al., 2023a) for all experiments. The encoder consists of an embedding network with four hidden layers, each containing 10 units and using Rectified Linear Unit (ReLU) activations. It encodes the batch {[T.,, 4, J into through mean pooling, subsequently mapping it to [HL, or.. The latent variable is sampled from a normal distribution defined by [1,, or). The decoder a three-layer neural network with nonlinear activation functions, mapping [ =] + to the predicted risk Ve The optimization objective is given in Eq. (5). For further details on the implementation, please refer to our code repository.
Input
Demaa/
Denow(10)
Demoei11
Demae/10)
q3 (2gl
Conditional prior H
Deroee
Danaw/10)
Dansee10;
{4
B
KL
B
MSE
Mean Pooling
qs 2 H.)
14,5)
Approximate posterior
Encoder
Decoder
Figure 8. Illustration of the neural architecture of the risk predictive model in (Wang al., 2025). The risk predictive model follows encoder-decoder structure which encodes the batch [Tr.: L,,i] into latent variable and then decodes it into predicted adaptation risks l,s. We reuse it in PDTS standard risk predictive model backbone.
B. Theoretical Analysis and Proofs
B.1. Preliminaries of MPTS
In MPTS (Wang al., 2025), the streaming adaptation outcome in task episodic learning is characterized as:
Oo eval
oval, (1 opt
O2 eval
opt
opt
Or.
(20)
The most encouraging finding is that these cumulated task identifiers and adaptation risk values can be coupled to learn helpful adaptation prior and serve the evaluation of the task difficulty in a rough granularity. The resulting dataset is used to train the risk predictive model associated with the modified optimization objective as Eq. (5).
B.2. One Secret MDP Steers Adaptation to Many MDPs
The Operator for the State Transition in the Secret MDP. Here, given a smooth and fixed machine learning optimizer, we can define the operator in zero-shot or few-shot adaptation optimization as:
State Transition Operation after Adaptation 9:0 x T + 0',
(21)
18

---

PDTS for Adaptive Decision-Makers in Randomized Environments
which corresponds the deterministic transition function in the secret MDP M. Here, take the DR an example, the above operator F results in the following prohablly transited states:
0 = arg min 0
E pp.(r:)) ((D),D;:0)) ]
B
(22)
++1 { 0:
11 V
1
EDD
D To E T J = B, Jhe S J
B
b=1
where Po (T; 0) corresponds to the probability density function of the (1 - ----le tasks.
Probabilistic Graphical Model of the Secret MDP. Here, we can write the probabilistic form of the constructed MDP
as:
P(T RIT Po.T ToTTITTT HoT--i) = p(O)) ARR1110011
II m (0
II 10==1 (23)
Initial State 1=1
Step Wise Reward
Task Sampler
State Transition
where (Jor7. ro.T, fo.r} records the trajectory information during policy search. The corresponding probabilistic graphical model is Fig. 2.
B.3. Proof Proposition 3.2
Proposition 3.2 (MPTS as a UCB-guided Solution to i-MABs) Executing MPTS pipeline in Eq. (6) is equivalent to approximately solving M with the i-MAB under the UCB principle.
To demonstrate the claim in Proposition 3.2, we revisit fundamental concepts in CVaR, optimization and break the problem into several steps. These include # Step1 the dual form of CVaR, with its Monte Carlo estimates, Lemma B.1 to determine the optimal step-wise action, and Step2 nearly CVaR,, upproximation/UCB-Guuuie
Solution to i-MABs in MPTS.
Stepl. The Greedy Action as the Subset with Top-B Adaptation Risk Values.
Unbiased Monte Carlo Estimate of CVaR.. Here, we can rewrite the optimization objective of CVaR, in the dual form:
min O
CVaR. == ++ 1 a -
| D. -
(24)
where the conditional function means (;; 0) - -c]+ max max D2; - c]+, 0}. And its Monte Carlo sample
average approximation corresponds
B
min CVaR. c ( a ((D -c] CERRER
(25)
As the optimal auxiliary variable satisfies Ç - (a). the Top-B element in the set {l 6 = E(D), DJ 0) can be viewed as the unbiased Monte Carlo estimate w.r.t. Eq. (24) and the equivalent form of Eq. (25). implementation, (Wang al., 2024c; et al., 2024) typically evaluate the machine learner's adaptation performance on candidate tasks, rank their values, and filter (1 a to optimize.
Lemma B.1 (Subset with Top-B Risk Values as the Optimal Arm). Given the secret MDP M in Section 3.1 and its step- wsie reward R(O, J.B ( (O.) (V selecting the subset with average Top-B risk values inherently
maximizes R(O,, TB
Proof. Given the state O, and the available action set A, in M, we need to show that selecting the The E A, brings largest CVaR decrease as CVaR,(()) 1
Let J denote the subset with highest average Top-B risk values in the set {es = EDD) DJ 0) which can be viewed as the Monte Carlo sample, i.e., unbiased estimate of pa (T; 0.). This also in accordance with the dual form of
19

---

PDTS for Adaptive Decision-Makers in Randomized Environments
C in Eq. (25). Meanwhile, we can define its transited state as 0++1- In a similar way, let some arbitrary feasible action, ie., the subset be T+1- we can define the corresponding transited state as 07 Inspired by the notation in Eq. (A.2), we re-express the subset T/41 as the Monte Carlo sample from another task distribution g(7), which differs from Pa (T; 0z).
Next, we can estimate the robustness improvement under different actions from the one-step Taylor expansion trick.
Unbiased Objective [(0)= :- (r.0.) [
Biased Objective (0) = Eyy(r) [[[D] D.; 0)]
(26a)
(00 0/+1 0. -)Vo(()) 0111 = 0, - -Vo((O))
(26b)
0, 0,112) = [(,)) (((((11) ~ (
(26c)
- [(0,)- - - = Vo cos 0q I n/ (0) = ((.)) - 16)
-
(26d)
where we typically assume the norms of stochastic gradients for VoL(o)) and Vec(O) are the same and their angle is ag. As a result, we can see the optimal stochastic gradient should be the unbiased estimate of the Eq. (24), which corresponds to the Top-B tasks in the pseudo batch J with 2 = 1 - a. This completes the proof of Lemma B.1.
Meanwhile, remember that in the task-selection MDP M. the agent will never revisit the previous state 0, due to the nature of the stochastic gradient descent in the operator F in Eq. (21). Finally, we can claim that maximizing the state action value Q(0,, 1 in the i-MAB actually corresponds to maximizing the step-wise reward due to the Bellman optimality in the main paper, and picking up the worst subset secretly maximizes R(0,, J2+1) in a greedy way. This implies 1 where (0 denotes the evaluated adaptation performance of feasible subset T conditioned 0z. In the presence of RATS, the adaptation performance evaluated by the risk predictive model p(E|T, H1:; 02) in an amortized way, which suggests Q(0;, TT.1) 1 with i the predicted value.
that πī arg maxy 4 ST 0(0,JT411 and Q(0,,
B EE 4442,
# Step2. UCB-Guided Solution to i-MABs in MPTS.
Assumption 3 (Randomized Adaptation Risk Value Function). Given the secret MDP M. in Section 3, we assume the distri- bution of the adaptation risk value follows an implicit Gaussian distribution, i.e., p(f++1./ Hi 0₂) = N(P + )-
Note that tasks with their identifiers are sampled an i.i.d. way; this induces the conditional independence and the distribution of their summation as Eq. (27) with the help of Assumption 3.
B
2
p(c B c+1 141:
B
H1:; 01) = II ple
H1:; 0,) = II N(u (411)
1=1
(27)
B
B
T B 141: 0.) N
o4t+)) = N( 8²
2
pl
1e+1,
041)
At the same time, we find in MPTS, the multiple stochastic forward passes in Eq. (7) are performed to obtain the MC estimated distribution parameters {m(()) := je+1,i, (fa) ..+1,i} for each task identifier т; in the batch T This can be further associated with the factorization in the risk predictive model as Eq. (28):
p 1 HiI 0,) ~ / Po 22)9 (z, H1
(28a)
/ ll pll6111 2 2)) H1.)
B
=
(28b)
The last step shows that the worst subset corresponds to the unbiased Monte Carlo estimate of CVaR,,, and its sample average adaptation risk value can be treated as the proxy of the reward. Hence, picking up the subset with each element in the Top-B risk values is doing exploitation in robust fast adaptation. Rethinking MPTS's acquisition function in Eq. (7), we
20

---

PDTS for Adaptive Decision-Makers in Randomized Environments
find the following inequality:
1
0+41
with Vor+1, E R+
(29a)
1=1
8
B
71
1+1
+
1,14+1.5
7 70/41,
(29b)
1=1
1=1
B
(e,) + Σ
m (C,)
rom())) == Av(Jz),
(29c)
UCB
Approximate UCB
which means MPTS actually executes the approximate UCB to balance the exploitation (picking up the subset with the estimated worst performance) and the exploration of the task space (picking the arm with the nearly highest epistemic uncertainty (Wang& & Van Hoof, 2020) captured by the risk predictive model).
The above two steps complete the proof of Proposition 3.2.
B.4. Proof Proposition 3.3
Proposition 3.3 (Concentration Issue in Average Top-B Selection) Let f(T) Fd + R be a unimodal and continuous function, where d E N+ and Rd, with maximum value f(T") at We uniformly sample set of points J (T,) where are i..dd with probability Pe of falling within e-neighborhood Tf as [S(T) [(T*)] S Following MPTS, we select the Top-B samples with the largest function values, i.e.,
Jh = Top-B(J" :f), B,BEN, B
For any >> 0 such that Pe
B-B+2 the concentration probability B+1
P f(=))
increases with B and converges to I with B + oo.
Proof. We define Pe as the probability that random variable uniformly sampled from the domain of definition falls within the neighborhood of the maximum value T", i.e.,
Pe =P((() -(( Se sel Unif(-)).
Next, consider the probability that least random variables from the set T (T. {T.) where B, B E N+ and B B, within the neighborhood of T* Since the 's are i.i.d. and t3 Unif(-), the probability given by
B pBBB - Σ
--i +11 (1- -ppp-- ( B
(111)
Since f' is unimodal and continuous function, we can directly relate the concentration probability as
PP -
VT Vr (T) pB.B
To establish the monotonicity of R. with respect to B, observe that the term p. 8- (1 (1-p)) is monotonically
decreasing in B for fixed i, given that P.
R-i+2 To see this, we compute the ratio of consecutive terms:
B+1
2 ++1 (1 - p. ) (123) 1 B-i+2
>
pt
-- ++2
p..) )
B+1
1-11
(. +1)
21

---

PDTS for Adaptive Decision-Makers in Randomized Environments
Thus, the sequence Pe 13-++1(1) (1-Pe) --p)=-1())) decreases with B, and consequently, pB.3 increases monotonically in when B is fixed, provided that P. B-B+2
Therefore, we conclude that for any c > 0 such that Pe
2 B-B+2 the probability P ((I((T) f(r")(()))) sc TT) 8+1
increases monotonically with B for fixed B.
B.5. Proof of Proposition 3.4
Proposition 3.4 (Nearly Worst-Case Optimization with PDTS) When B grows large enough, optimizing the subset from Eq. (1I) achieves nearly worst-case optimization.
Proof. As the exact size of the subset B is fixed in the optimization, the ratio goes nearly 0 with the increase of B to certain scale. As the consequence, the number of arms grows to C% and the robustness concept is 1V Since the involvement of the diversity regularization perturbs the worst arm selection, this induces the nearly worst-case optimization PDTS.
C. Experimental Setups & Implementation Details
C.1. Risk-Averse Baseline Details
These baselines are SOTA methods published in NeurlPS/ICLR conferences and the latest open-sourced version (Wang et al., 2024c; Lv et al., 2024; Sagawa et al., 2019; Wang et al., 2025; Greenberg et al., 2024).
GDRM (Sagawa et al., 2019; Setlur et al., 2024). As briefly introduced in the main paper, GDRM can be viewed as a max optimization problem. The core concept of enhancing the machine learner's 's robustness involves reallocating more probability mass to the worst-case scenarios in a weighted manner. In each iteration with the optimal pa (7), the optimization problem simplifies to:
min E (o( [(CD). D.: DE
EE((r)
(p(T))
p(T) ((D) (DD
(30)
where we use (F)
palt) p(r)
denote the weight.
In general, for a fixed number of tasks, GDRM organizes tasks heuristically or dynamically into clusters, followed by a reweighting mechanism based on assessed risks. However, in task-episodic learning, no task grouping performed because the task batch is reset after each iteration (Wang al., 2024c; Lv al., 2024). Task-specific weights calculated as
@ 2 0))
(T:) Eie ( D5, .0)" where denotes the temperature parameter, and (7,)²= represents the identifiers of the task batch. Additional implementation details are available at https: /ttp://ggtttttbt com/ wwwppangwwi/grouppp..u.
DRM (Wang et al., 2024;; Lv et al., 2024). As outlined in the main paper, we adopt the standard approach in DRM, where the optimization objective is expressed as (r:00 ((D), D.; 0) A widely used practical strategy involves evaluating the performance of the machine learner, ranking task task-specific adaptation risks, and optimizing over the worst (1 a) proportion of tasks.
Following the setup in Wang et al. (2024c), we select the Top-B tasks during optimization, which corresponds to evaluating
task batch of size BB T To ensure fair comparison with PDTS while preserving computational efficiency, we employ the same Monte Carlo estimator for the risk quantile as in Wang et al. (2024c). For all benchmarks, we set the actual task batch size to B- = 2B, discarding the easiest half before optimizing the machine learner.
MPTS (Wang et al., 2025). We have thoroughly introduced the MPTS pipeline in Section 2.2 and provided details of the risk learner in Section A.5. For additional configurations, we adopt those from the official repository at https: //gi thub. cww/co/////oo//oo htmmttuurri aalbbbbba //MM,, including the heuristic random mixture strategy and the specific values of B.
22

---

PDTS for Adaptive Decision-Makers in Randomized Environments
Table Details of Task Identifiers and Algorithm Backbones Across Benchmarks. Here, we provide detailed information about the task identifiers used to induce task distributions and the algorithm backbones, including MAML (Finn et al., 2017), TD3 (Fujimoto et al., 2018), and PPO (Schulman et al., 2017), employed various benchmarks.
Benchmarks
Identifier Meaning
Identifier Range
Backbone
K-shot sinusoid regression
amplitude and phase (a, b)
[0.1. 5.0) x [0, x]
MAML
Meta-RT.; ReacherPos
goal location (T1)
0.2, 0.2] x 0.2, 0.2]
Meta-RL. Walker2dVel
velocity
[0.2.0]
MAML
Meta-RL: Walker2dMassvel
mass and velocity (m, v)
[0.75, 1.25] x [0, 2.0]
Mcta-RL: HallChoctahMasssell
mass and velocity (m, v)
[0.75, 1.25] x [0,
DR: Pusher
puck friction loss f and puck joint damping d
[0.004, 0.01] x [0.01,, 0.025]
DR: LumarLander
main engine strength s
[4. 20]
TD3
DR: FrgoReacher
joint damping d and max torque (x4 joints)
[0.1, 2.0j x [2, 20]
VisualDR: LifiPegUUpighh.LLight
ambient light 7 (x3 dimensions) goal location (12)
[-0.0, 2.0] [000, 1.0] x (0.0. 1.0)
PPO
VisualDR: AnymalCReach Goal
C.2. Sinusoid Regression
Following standard setups in prior works (Finn et al., 2017; Wang et al., 2025), we define the sinusoid regression problem as toy example of supervised meta-learning The goal of this problem is to predict the wave function sin - b), where the amplitude a and phase b are sampled as task-specific parameters, i.e., task identifier T = (a, b). using a few-shot support dataset. We adopt the 10-shot sinusoid regression setting, where 10 data points are uniformly sampled from the interval [-5.0, 5.0) to form the support dataset for each task.
Implementation Details. The machine learner is a neural network with 2 hidden layers, each of size 40, using the Rectified Linear Unit (ReLU) as the nonlinear activation function. The task batch size is set to 16 for ERM and GDRM, while batch size of 32 is used as the default for DRM. The temperature parameter in GDRM is rj = 0.001, and the learning rates for both the inner and outer loops are fixed at 0.001. The task identifier has a dimensionality of 2. For MPTS and PDTS, the identifier batch size during training is set to 32 (2x) and 512 (32x), respectively. We use the Adam optimizer with a learning rate of x 10-4 to update the risk learner over 15,000 steps. The label for the risk learner is the average MSE loss value for each task. For sinusoid regression and meta-reinforcemmen learning, we use the standard repository provided by MAML (Finn al., 2017). Regarding validation during training, we use separate uniform task sampler with fixed random seed to select 1,000 tasks for validating the training checkpoints of various methods.
C.3. Meta Reinforcement Learning
We adopt four Meta-RL RL Eecnarios-RRaccerPPoo Valker2dVel, Walker2dMass46 and HallChettaMaassod-- MPTS (Wang et al., 2025), which represent distinct MDP distributions based on the Mujoco physics engine (Todorov et al., 2012). These scenarios involve three types of robots (HalfCheetah, Walker2d, and Reacher) and three randomized meta-lcarning objectives: velocity, goal location, and body mass.
The objective in velocity- scenarios, e.g., Velker2dVel, is to train the robot achieve a target velocity, with the reward function defined as the negative absolute difference between the robot's current velocity and the target velocity. This reward is augmented by a control penalty and alive bonus to facilitate learning. As the task distribution is specified by a uniform distribution over the target velocity, T y can be viewed the task identifier.
The body mass and velocity scenarios, e.g., Walker2dMass Vel and HalfCheetah Vel, share the same objective as the velocity scenarios but additionally feature varying robot masses. As the task distribution is specified by a uniform distribution over the body mass and target velocity, (m, v) can be viewed the task identifier.
The goal location scenario, e.g., ReacherPos, requires moving a two-jointed robot arm's end effector close to a target position. Its reward function defined as the negative L, distance between the end effector's position and the target,
23

---

PDTS for Adaptive Decision-Makers in Randomized Environments
a
(b)
(c)
Figure 9. Illustrations of three types of robots in Meta-RL based on the Mujoco. (a) Reacher, (b) HalfCheetah, and (c) Walker2d.
b
(c)
Figure Illustrations three physical robotics domain randomization scenarios: (a) LunarLander, (b) Pusher, and (c) ErgoReacher. The illustration of ErgoReacher is adapted from (Mehta al., 2020).
supplemented by a control cost to encourage robustness. As the task distribution specified by a uniform distribution over the goal location, = (I1, .2) can be viewed the task identifier.
Implementation Details. The machine learner implemented as a neural network with 2 hidden layers, each consisting of 64 units, and utilizes ReLU activations nonlinearity. For ERM and GDRM, the default task batch size is set to 20, while DRM uses batch size of 40. The temperature parameter for GDRM is configured 0.001. Both the inner and outer loop learning rates fixed 0.1. Besides, the identifier batch size during training is 30 (1.5x) for MPTS and 1280 (64x) for PDTS. The risk learner is updated using the Adam optimizer with learning rate of 5 x 10 The label for the risk learner is the negative average reward value at the final step for each task. For validation during training, we uniformly sample 40 tasks from the task space at fixed intervals to validate the training checkpoints of different methods. For meta-testing after training, we uniformly sample 100 tasks from the task space to test the trained models different methods.
C.4. Physical Robotics Domain Randomization
As shown in Fig. 10, we adopt three scenarios for robotics domain randomization from Mehta et al. (2020); Pusher, LunarLander, and ErghRe The task distribution is defined as in (Wang et al., 2025). Specifically:
LunarLander is a 2-degree-of of-eeedm (DoF) environment where the agent must softly land a spacecraft, implemented in Box2D (Catto, 2007). Its reward function provides positive rewards for successful landings, negative rewards for crashes, and penalties for fuel consumption and deviations from the landing pad, thereby promoting efficient and controlled landings. The task distribution is specified by a uniform distribution over the main engine strength which can be viewed as the task identifier =.
Pusher is a 3-DoF robotic arm control environment based on MuJoCo (Todorov et al., 2012), where the agent pushes puck to target. The reward function penalizes the L2 distance between the puck and the target, augmented by a control penalty. The task distribution is specified by uniform distribution over the puck friction loss and puck joint damping d, which can be viewed the task identifier T (f.d)
24

---

PDTS for Adaptive Decision-Makers in Randomized Environments
(a)
Figure Illustrations two visual robotics domain randomization scenarios: controlling (a) a table-top robotic arm and (b) a quadruped robot, operating under randomized lighting conditions and varying goal locations, respectively.
ErgoReacher involves a 4-DoF robotic arm implemented in the Bullet Physies Engine (Coumans, 2015), tasked with reaching goal using its end effector. Its reward function penalizes the distance between the end effector and the target, combined with control penalties. The task distribution is specified by uniform distribution over the joint damping d and max torque across 4 joints, which can be viewed as the task identifier T (d1, d2, d3, da, L1, 12, 13 14).
Details of the randomized task identifier range are presented in Table 3.
Implementation Details. The machine learner is neural network with two hidden layers, each consisting of 10 units, and uses ReLU activation functions. For ERM and GDRM, the task batch size is set to 10, while DRM uses batch size of 20. GDRM employs a temperature parameter of 0.01. We adopt TD3 algorithm (Fujimoto et al., 2018) 1) as the algorithm backbone. The actor and critic learning rates are both set to 3 x 10 For MPTS, the identifier batch size during training is 25 (2.5x) for LunarL.ander, 50 (5x), and 250 (25x) for ErgoReacher. contrast, for PDTS, the identifier batch size during training is 640 (64x) for all environments, with no additional requirements for fine-tuning. The risk learner is updated using the Adam optimizer with a learning rate of 0.005. The label for the risk learner the negative average return for each task. For validation during training, we uniformly sample 100 tasks from the task space validate the training checkpoints of different methods.
C.5. Visual Robotics Domain Randomization
Visual-based robotics control is common and crucial in real-world scenarios. As illustrated in Fig. 11, based the latest robotics simulator, ManiSkill3 (Tao al., 2024), design two scenarios for visual roboties domain randomization: MirightUprighLLiigee and AnymalCReach.GGon These scenarios involve controlling a tabletop two-finger gripper arm robot and quadruped robot, respectively, under randomized lighting conditions and goal locations. Specifically:
UprrPht Light is derived from the LiftPegUprightt-Il scenario in ManiSkill3, where the objective is move peg lying the table to upright position. To emulate the complex lighting conditions found in real-world environments, we modify this scenario to make the ambient lighting controllable. An illustration of this setup is shown in Fig. 6. The task distribution is specified by uniform distribution over the configurations of ambient light, which can be viewed as the task identifier t = (2,12,1a).
AnymalCReach Goal is adapted from the LnymalC-Reach-v scenario in ManiSkill3. The task is to control the AnymalC robot to reach a target location in front of it. Inspired by point-robot navigation scenarios in prior works(Finn et al., 2017), we randomize the goal location, which remains visible to the robot. The task distribution is specified by a uniform distribution over the goal location, which can be viewed as the task identifier (x),
The detailed randomization configurations are provided in Table 3.
Implementation Details. We use the PPO algorithm (Schulman et al., 2017), along with its hyperparameters and network architecture, as provided in the official ManiSkil13 codebase. To accommodate GPU memory limitations, we set the
25

---

PDTS for Adaptive Decision-Makers in Randomized Environments
task batch size to 16 for ERM and GDRM and 32 for DRM, with 500 steps per iteration on Upriihe For AnymalCReach Goal, the task batch size is 128 for ERM and GDRM and 256 for DRM, with 200 steps per iteration. For MPTS, the identifier batch size during training is 32 for LiftPeg Light and 256 for AnymalCReach PDTS, the identifier batch size during training is 64x the default for all environments. The risk learner is updated using the Adam optimizer with a learning rate of 0.005. The label for the risk learner is the negative average reward for each task. For validation during training, we uniformly sample 64 tasks from the task space to validate the training checkpoints of different methods.
(2.5x). For
D. Additional Experiment Results
PDTS (Ours)
MPTS
ERM
GDRM
DRM
Sinusoid Regression fix) asin(x
Validation Curves
Meta-Testing Performance
1.0
1.5
MSE
-122
n/2
1.0
0.5
0.2
[a. b]
Iteration (x10") (b)
Average CVaRO.5 CV3RO.7 CVaRO.9
(a)
(c)
Figure 12. Few-Shot Sinusoid Regression Results. (a) Illustration of the sinusoid regression problem, where the task identifier T consists of the amplitude and phase (a, bj. (b) Curves of averaged MSEs on the validation task set for all methods during meta-training. (c) MSE values for all methods at various a levels of CVaR. during meta-testing.
D.I. Beyond Decision-Making: Robust Supervised Meta-Learning
This work primarily focuses on risk-averse sequential Decision-making However, PDTS is readily extendable to other risk-averse scenarios, such as supervised meta-learning. As shown in Fig. 12, we evaluate PDTS on sinusoid regression, a commonly-used toy example introduced in Finn et al. (2017). which involves adapting quickly to new functions using only 10 samples. Consistent with the results observed in decision-making PDTS achieves superior performance compared to all baselines, demonstrating faster average performance and more robust adaptation. Moreover, as demonstrated in MPTS (Wang et al., 2025), the RATS paradigm has broad applicability. including image classification (Gondal et al., 2024). It is believed to be promising in other interesting areas, such as LLM-guided decision-making (Ma et al., 2024; Wang et al., 2024a; Qu et al., 2024; 2025), multi-agent systems (Shao et al., 2023b;a; Qu et al., 2023), and data sampling in offline reinforcement learning (Levine et al., 2020; Zhang et al., 2023; Mao et al., 2023a;b; Hong et al., 2023; Mao et al., 2024a:b).
D.2. Ablation Studies and Additional Analysis
In this section, we perform additional experiments to carry out ablation studies on the hyperparameters and components of PDTS, demonstrate the presence of the concentration issue in MPTS, and validate the effectiveness of PDTS in addressing it. For computational efficiency, we use sinusoid regression as the testbed.
Ablation Study on Diversity Regularization Weight 7. As shown in Fig. 13()), we evaluate the effect of varying values of the diversity regularization weight 7. It evident that diversity regularization plays a crucial role in PDTS, as its absence (7 0) leads dramatic performance drop. PDTS demonstrates robustness to the choice of y within a certain range (c.g.. [1, 2). However, excessively large values of ry degrade performance, emphasizing the importance of balancing diversity and robust optimization. In most cases, setting y to I or nearby value secures superior enough performance. For scenarios with an extremely low-dimensional task identifier, increasing 7 appropriately may improve performance.
Ablation Study on Key Components: Posterior Sampling and Diversity Regularization. In simple terms, PDTS can be viewed as the combination of MPTS and diversity regularization, with UCB replaced by posterior sampling. We conduct an ablation study to highlight the significance of each component. As shown in Fig. 13(b), we replace the UCB in MPTS with posterior sampling (MPTS+P) and incorporate diversity regularization into MPTS (MPTS+D). The results demonstrate that each component contributes significantly to the superiority of PDTS.
26

---

PDTS for Adaptive Decision-Makers in Randomized Environments
Average CVaRo.5
Average CVaR0.9
Tosks
Average
CVaR0.9
5
Selected Tasks
NPTS_CVARD CVARO.9
PDTS Average
cveno.9
MSE
Phase
1.0
1 2
MPTS MPTS+P MPTS+D POTS
Amplitude (c)
2x
Bx
32x
(b)
(d)
Figure 13. (a) Meta-testing results trained with different hyperparameter, "/. (b) Ablation studies key components, including posterior
sampling (P) and diversity regularization (D). (c) Visualization of the distribution of candidate tasks (gray points) and selected tasks (red
points) in MPTS, highlighting the concentration issue. (d) Comparison of the average and CVaRo. meta-testing performance of PDTS and MPTS with increasing pseudo-batch sizes.
The Presence of the Concentration Issue in MPTS. As introduced in Sec. 3.2, theoretically prove the presence of concentration issue in MPTS. To empirically demonstrate the concentration issue and its impact, we evaluate MPTS on sinusoid regression. Fig. 13(c) shows that, as the candidate batch size increases, MPTS tends to select tasks concentrated within a small region-a phenomenon referred to as the concentration issue in the main paper. As illustrated in Fig. 13(d), this concentration issue leads to catastrophic performance degradation in MPTS as the candidate batch size increases.
PDTS Addresses the Concentration Issue and Benelits from Improved Coverage. In contrast to MPTS, Fig. 13(c) demonstrates that by incorporating diversity regularization, our method, PDTS, avoids the concentration issue and does not experience performance collapse as the candidate batch size increases. More impressively, the performance of PDTS improves with increasing candidate batch size, demonstrating the benefits of encouraging broader coverage of the task space during subset selection as proposed by PDTS.
Ablation Study on the Risk Predictive Model. We conducted an ablation study to analyze the impact of the risk predictive model. We designed two variants, PDTS-Deep and PDTS-Shallooe by increasing and decreasing the number of encoder-decoder layers, respectively. Additionally, we replaced the encoder-decoder structure with an MLP create PDTS-MLP. Results Walker2dVel are summarized in Table 4. From these results, we observe: (1) All PDTS variants achieve better task robust adaptation, confirming the effectiveness of PDTS and its generality across different risk predictive models. (2) The encoder-decoder architecture generally outperforms MLP-based models, supporting the rationale behind this design. (3) Deeper networks may introduce a performance-robusten trade-off in the current setting, which we plan to further investigate more complex scenarios. (4) Weaker risk prediction models degrade overall performance, due to poorer difficult MDP identification.
Table Comparison of methods with different risk predictive models on Walker2dVel.
Methods
CVaRo.9
CVaRo.7
CVaRo.5
Average
ERM
-69.77+7.62
-31.73++.82
-6.66
38.88-t-4.73
PDTS
-22.42-13113
2.8641.00
16.5742.93
40.40:13.07
PDTS-Deep
-30.5147.11
0.55:15.69
17.99
44.42.13.69
PDTS-Shallow
-41.2414.74
-11.3444.5
3.92+4.33
33.64::..01
PDTS-MLP
-41.07+4.88
-1..04+5.06
3.38
32.96++4.58
E. Other Discussions
Relation with Traditional Active Learning. Traditional active learning (Ren al., 2021) aims at reducing the sampling redundancies during optimization and exploiting historical optimization information to improve learning efficiency, such as annotations and computations. The active query strategies (Zhu et al., 2003; Gal et al., 2017; Kirsch et al., 2019; Wu et al., 2022; Mukhoti et al., 2023) also rely on some predictive models and utilize principles like uncertainty, diversity, etc. RATS, such as MPTS (Wang et al., 2025) and PDTS in this work, stresses the importance robustness during active sampling. Hence, the predictive model requires scoring the task difficulties without exact evaluation.
27

---

PDTS for Adaptive Decision-Makers in Randomized Environments
When MPTS Meets Diversity Regularization. The diagnosis of the concentration issue in Sec. 3.2 identifies the diversity regularization as plausible solution encourage the exploration of the task space and bring more worst-case robust solutions as more arms are constructed by increasing B. Actually, we also examine this part and include the ablation studies on the sinusoid regression in Fig. 13(b). For other evaluations in the main paper, we assume that MPTS's group (Wang et al., 2025) adopts the optimal hyper-parameter configurations. Hence, their setup is adopted to produce MPTS results. Meanwhile, the posterior sampling's advantage lies (i) no extra hyperparameter adjustment, unlike UCB used MPTS, and (ii) stochastic optimism when the uncertainty is difficult to estimate.
PDTS Agnostic the Risk Predictive Model. As RATS in this work is rarely investigated concept in the field, limited methods have been developed score task difficulties, particularly MDPs' difficulties under a policy. The risk predictive model in MPTS has approximately achieved the purpose. Hence, we reuse their module as the backbone. In reality, our theoretical analysis and PDTS will be compatible with other riseppredi models in the future.
F. Computational Platform & Software
This research project conducts experiments using NVIDIA 3090 GPUs in computation, and Pytorch works as the deep learning toolkit in implementation. The software requirement list can be found in the open-source code from the project website. Please refer to technical blog from team website at ht htsps: / /www. thui dm. com/.
28


Reviews:
[{"review_id": "U7buVbR1iP", "reviewer": "Reviewer_xw4Y", "summary": "This paper studies a robust active task sampling (RATS) paradigm, models it as an infinitely many-armed bandit (i-MAB) problem, and proposes a novel method called Posterior and Diversity Synergized Task Sampling (PDTS). PDTS mitigates the task concentration issues in an existing approach, Model Predictive Task Sampling (MPTS), by incorporating a diversity regularized acquisition function, replacing the upper confidence bound acquisition function in MPTS. As a result, PDTS enables exploration in a broader range of tasks and improves robustness for nearly worst-cases. The authors conduct extensive experiments in various meta RL settings, and show that PDTS improves CVaR robustness, sample-efficiency for average return, and zero-shot adaptation in out-of-distribution tasks.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "_**We sincerely thank Reviewer xw4Y's efforts and thoughtful feedback. Below, we provide our responses.**_\n___\n\n**1. Discussion of curriculum learning methods**\n\nThank you for your valuable suggestion. We'll cite these important works and discuss them in the revised manuscript, adding contents below to Line 667:\n\n> 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.\n\n**2. Writing optimization**\n\nThank 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.\n\n**3. Trade-off between worst-case performance and the computational costs**\n\nThanks 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.\n\n**4. Limitations discussion**\n\nThank 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.\n\n**5. Meta-RL agent have access to task identifiers**\n\nApologies 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.\n\n___\n\n_**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.**_\n\n___\n\n**References:**\\\n[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.\\\n[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.\\\n[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.\\\n[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.\\\n[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.", "claims_and_evidence": "Claims in terms of mitigating the task concentration issue in MPTS, computational efficiency, versatility of the framework, improved robustness, and zero-shot adaptation are well supported through theoretical and empirical evidence. Although theoretically justified, empirical evidence supporting proposition 3.4 is missing. An interesting ablation would be observing the trade-off between improved nearly worst-case robustness and computational efficiency of PDTS.", "methods_and_evaluation_criteria": "They make sense when evaluating CVaR robustness, sample-efficiency, and zero-shot adaptation in meta RL.", "theoretical_claims": "Yes, I checked the correctness of proofs of propositions 3.2, 3.3, and 3.4. I haven't seen any issues.", "experimental_designs_or_analyses": "I checked the experimental designs behind emta RL, domain randomization, and visual domain randomization settings. I haven't seen any issues.", "relation_to_broader_scientific_literature": "The i-MAB model provides a theoretical framework for active task sampling, which is useful for risk-averseness in many decision-making problems, such as reinforcement learning and robotics, where robust adaptation is key for real-world applicability.", "essential_references_not_discussed": "The paper discusses various task sampling methods for providing risk-averseness and robustness, yet it does not mention curriculum learning at all. For example, in the paper that presents RoML (Greenberg et. al.), which the authors evaluate in their documents, curriculum learning methods are also studied and discussed. I believe a discussion of curriculum learning methods would broaden the audience of this work. Here are a few of those instances:\n\nDennis, 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.\n\nJiang, 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.\n\nKoprulu, 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.", "other_strengths_and_weaknesses": "Strengths:\n- The paper clearly structures and explains the motivation, the problem, the proposed method and the experimental results.\n- Introduction of i-MABs as a model for RATS provides a versatile theoretical framework.\n- PDTS mitigates key issues in an existing method, MPTS, and provides computational efficiency.\n- Theoretical results are clearly explained, and empirical evaluation presents strong evidence in favor of the proposed method in improving risk-averseness, robustness and generalization capabilities of meta RL agents.\n\nWeaknesses:\n\n- The flow of Section 2 and 3 degrades as more notation is used.\n- There is no ablation study to justify nearly worst-case optimization, as proposition 3.4 suggests. Although the authors prove the proposition, an interesting study would be on the trade-off between worst-case performance of PDTS and the computational costs of increasing the cardinality of the set of candidate tasks.\n- There is no discussion on the limitations of the introduced model and the proposed method.", "other_comments_or_suggestions": "I didn't see any typos. However, I highly recommend that the authors refine their use of symbols. Section 2, and most importantly, Section 3, is very hard to read.", "questions_for_authors": "Does the meta RL agent have access to task identifiers? I assume it does not. But then I'm confused how they can be utilized in Eq. (11) to measure the diversity of candidate tasks. As the task identifiers are assumed to be partially observable in meta RL, I think a clarification is needed here.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "VHaFQzy208", "reviewer": "Reviewer_b9NR", "summary": "This paper tackles robust active task sampling (RATS) in domain randomization or meta-RL for worst-case performance. Tasks are viewed as arms in an infinite multi-armed bandit, but the existing MPTS can over-concentrate on top-B tasks. The authors propose PDTS, which replaces UCB-based acquisition with posterior sampling and adds a diversity term. Experiments show PDTS achieves faster, more robust adaptation than baselines (ERM, DRM, GDRM, MPTS) on MuJoCo and domain-randomized robotics tasks.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "_**We sincerely appreciate Reviewer b9NR's efforts and constructive feedback. Below, we provide our responses.**_\n___\n\n**1. Additional computational overhead**\n\nThanks 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.\n\n**2. Ablation study on the risk predictive model**\n\nThank 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:\n\n|Methods|$\\text{CVaR}_{0.9}$|$\\text{CVaR}_{0.7}$|$\\text{CVaR}_{0.5}$|$\\text{Average}$|\n|-|-|-|-|-|\n|ERM|-69.77$\\pm$7.62|-31.73$\\pm$7.82|-3.78$\\pm$6.66|38.88$\\pm$4.73|\n|PDTS|**-22.42$\\pm$3.13**|**2.86$\\pm$3.04**|16.57$\\pm$2.93|40.40$\\pm$3.07|\n|PDTS-Deep|-30.51$\\pm$7.11|0.55$\\pm$5.69|**17.99$\\pm$4.8**|**44.42$\\pm$3.69**|\n|PDTS-Shallow|-41.24$\\pm$4.74|-11.34$\\pm$4.5|3.92$\\pm$4.33|33.64$\\pm$4.01|\n|PDTS-MLP|-41.07$\\pm$4.88|-12.04$\\pm$5.06|3.38$\\pm$4.94|32.96$\\pm$4.58|\n\nFrom these results, we observe:\n- All PDTS variants achieve better task robust adaptation, confirming the effectiveness of PDTS and its generality across different risk predictive models.\n- The encoder-decoder architecture generally outperforms MLP-based models, supporting the rationale behind this design.\n- Deeper networks may introduce a performance-robustness trade-off in the current setting, which we plan to further investigate in more complex scenarios.\n- Weaker risk prediction models degrade overall performance, due to poorer difficult MDP identification.\n\n**3. Scaling to high dimensions**\n\nThank 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.\n\n**4. Sensitivity of hyperparameter $\\gamma$ and tuning practices**\n\nWe 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**\n- In most cases, setting $\\gamma$ to 1 or a nearby value secures superior enough performance.\n- For scenarios with an extremely low-dimensional task identifier, increasing $\\gamma$ appropriately may improve performance.\n\nWe will incorporate these recommendations into the revised manuscript and specify the $\\gamma$ values used in practice in the released code to aid practitioners.\n\n**5. Would you expect the best $\\gamma$ to scale with $\\hat{\\mathcal{B}}$ or with dimension $d$?**\n\nWe 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. \n\nAdditionally, 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).\n\n\n___\n\n_**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.**_\n\n___\n\n**References:**\\\n[1] Miller, D. L. (2024). WordLlama: Recycled token embeddings from large language models.  https://github.com/dleemiller/wordllama \\\n[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.", "claims_and_evidence": "Overall, the claims have coherent theoretical backing and empirically strong results across multiple benchmarks.", "methods_and_evaluation_criteria": "The evaluation criteria make sense for risk-averse policy adaptation. The chosen suite of MuJoCo and robotic tasks is widely used in domain randomization and meta-RL research, so the evaluation is aligned with standard practice.", "theoretical_claims": "All theoretical results are plausible and consistent with standard concepts in bandit theory and set diversification. There is no obvious flaw in these short formal statements. Given the scope of the paper, the claims appear correct, and the sketches/logic are standard enough that no glaring issues stand out.", "experimental_designs_or_analyses": "The experimental design is thorough and appropriate for the proposed method, supporting the authors’ conclusions about robust adaptation performance and sample efficiency.", "relation_to_broader_scientific_literature": "The paper situates itself among:\n\n- Risk-averse RL frameworks (DRM/CVaR, GDRM, etc.).\n\n- Meta-RL approaches.", "essential_references_not_discussed": "No major omissions jump out. The paper cites standard domain-randomization, risk-averse RL, and meta-learning literature.", "other_strengths_and_weaknesses": "- Strengths:\n\n1. The i-MAB perspective is novel and successfully integrates RATS with risk-averse RL under a cohesive theoretical argument.  \n\n2. The PDTS method is simple but effectively addresses subset concentration by blending posterior sampling and diversity.  \n\n3. The empirical evaluation is robust, spanning multiple benchmarks (both symbolic and visual), and consistently demonstrates performance improvements.\n\n- Weaknesses:\n\n1. Relying on a large pseudo batch size $\\hat{B}$ for nearly worst-case coverage can introduce computational overhead. \n\n2. The new diversity regularization parameter $\\gamma$ may require careful tuning. \n\n3. The method heavily depends on a risk-predictive model. If that model’s performance degrades, PDTS coverage might fail to accurately capture the most challenging tasks. While correlation results are encouraging, model reliability remains a potential concern.", "other_comments_or_suggestions": "- Clarifying the best practices for choosing the diversity weight $\\gamma$ or the approximate search method might help practitioners. \n\n- Additional ablations on how the quality of the risk-predictive model $p(\\ell|\\tau)$ influences PDTS’s final performance could further strengthen the discussion.", "questions_for_authors": "1. Scaling to high dimensions: How does PDTS handle extremely high-dimensional task spaces? Do you have any heuristic or projection strategy if the dimension is large? This would clarify how widely PDTS can be applied in large-scale real-world DR or meta-RL scenarios.\n\n2. Diversity Regularization: In practice, how sensitive is the method to the diversity weight $\\gamma$? Would you expect the best $\\gamma$ to scale with $\\hat{B}$ or with dimension $d$? If so, how?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "xbLhPO9y3G", "reviewer": "Reviewer_v1xi", "summary": "The paper focuses on adaptation robustness, addressing scenarios where a risk-predictive model is utilized to mitigate intense evaluation requirements. It formulates the robust active task sampling (RATS) problem as a partially observable Markov decision process (POMDP), providing theoretical insights into the problem. Empirically, the paper demonstrates that the proposed method, Posterior-Diversity Synergized Task Sampling, achieves stronger performance in vision-based reinforcement learning tasks compared to baseline methods.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "_**We sincerely appreciate Reviewer v1xi's efforts and positive feedback. Below, we provide our responses.**_\n___\n\n**1. Clarification on the simplification of configurations**\n\nApologies for any confusion. The simplification of configurations lies in two aspects: \n- 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.\n- 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. \n\nWe will highlight these in the revised manuscript.\n\n**2. Discussion on assumptions failure**\n\nThank 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.\n\n___\n\n**_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._**\n\n___\n**References:**\\\n[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.\\\n[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.", "claims_and_evidence": "The claims are supported by empirical results.\n\nOne question I have is regarding the introduction, where the authors state that their method requires less complex configurations compared to prior works. Could the authors clarify what is meant by ‘configurations’? Providing additional context and specific examples would help in understanding this claim", "methods_and_evaluation_criteria": "The methods and evaluation protocol are mainly based on meta-RL literature and robust RL literature, which makes sense to me.", "theoretical_claims": "I did not evaluate the theoretical analysis as I am not familiar with this field.", "experimental_designs_or_analyses": "The experimental design appears sound to me.\n\nA general question I have is regarding the limitations of the proposed method. What are the potential failure cases if the assumptions made in the analysis do not hold? A discussion on these aspects would help in understanding the robustness and applicability of the approach.", "relation_to_broader_scientific_literature": "Robust adaption is very important in real-world applications, especially robotics applications.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "N/A", "other_comments_or_suggestions": "N/A", "questions_for_authors": "Please see previous sections.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "jjspVqV8n3", "reviewer": "Reviewer_beLW", "summary": "The authors study the problem of robust reinforcement learning through the lens of meta-RL in which the trained (meta) agent receives a few task-specific samples (can be zero) using which it adapts to a new task. The objective is to maximize the expected performance conditioned on the sampled task being in a specific bottom quantile w.r.t. performance (CVaR). The proposed approach is based on the RATS framework and an existing approach called MPTS which continuously learns a model (along with the meta-agent) that can be used to predict the expected reward (or loss) of the agent on a new task given historical information about the samples used to train the agent. This predictive model is then used to select new tasks for the next iteration of the meta-RL algorithm. The main contributions of this paper are (i) formulation of the task selection problem (which tasks to use in each iteration of the meta-RL algorithm) as a higher-level reinforcement learning problem in which the actions are subsets of tasks, (ii) connecting the task-selection algorithm of MPTS to UCB, and (iii) proposing a new task-selection procedure that includes diversity of selected tasks in its objective.", "questions": "", "limitations": "", "rating": 5, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**_We sincerely appreciate Reviewer beLW's efforts and recognition of our work. Below, we improve the manuscript based on beLW's feedback._**\n___\n\n**1. Terminology clarity about secret MDP**\n\nThank 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.\n\n**2. Comparison betwen PDTS and MPTS**\n\nThank 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.\n\nIn the revised manuscript, we will involve more RATS background and further highlight the significance of PDTS for more general readers.\n\n___\n\n**_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._**", "claims_and_evidence": "Yes. The theoretical claims seem sound (though I haven't checked proofs in Appendix carefully) and the experiments show improved performance over existing approaches. The authors do not claim anything that is not backed by evidence in the form of experiments and theory.", "methods_and_evaluation_criteria": "Yes. The authors use standard evaluation methods and use standard RL environments in the experiments.", "theoretical_claims": "I did not reach the proofs in the Appendix (only briefly glanced over them).", "experimental_designs_or_analyses": "The experiments appear sound and there is no major issue with the evaluations.", "relation_to_broader_scientific_literature": "- Overall, framing the task-selection problem as an RL problem appears novel and could potentially provide a framework for future work on meta learning. The authors also showed how prior work like MPTS fits into their framework.\n- Insight on task diversity and how the selection criterion can lead to selecting tasks from a narrow range of values is interesting and applicable more broadly.\n- The proposed approach is agnostic to the specific meta-RL algorithm used and the authors show this using experiments based on different meta-RL algorithms in the literature.", "essential_references_not_discussed": "None that I am aware of.", "other_strengths_and_weaknesses": "### Additional Strengths\n\n- The experimental results appear convincing. The authors show clear improvements in both (i) the final performance and (ii) the sample-efficiency of learning. \n- The approach is shown to be lightweight and offers performance improvements with minimal impact to cost.\n- The paper is well-written and prior concepts related to MPTS are explained well.\n\nThere does not appear to be any major weakness that would impact acceptance. Some minor comments are below.", "other_comments_or_suggestions": "- I think the term \"secret MDP\" is a bit confusing as it suggests that there is a secret environment in which we want the agent to perform well. Something like task-selection MDP or meta-MDP is maybe less confusing.\n- As someone with no prior background on MPTS, it was a little hard to figure out which parts are \"new\" in this paper. I had to look at the MPTS paper to clearly figure this out. It would be great if this can be clarified in the paper.", "questions_for_authors": "N/A", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: All reviewers agree that the paper is worth publication at ICML. For this reason, I recommend acceptance. The authors are encouraged to revise the manuscript by incorporating all reviewer suggestions, including empirical findings, limitations, and statistical results, while also polishing the writing for clarity.
Decision: Accept (poster)