Foundation models have revolutionized general-purpose problem-solving, offering rapid task adaptation through pretraining, meta-training, and finetuning. Recent crucial advances in these paradigms reveal the importance of challenging task prioritized sampling to enhance adaptation robustness under distribution shifts. However, ranking task difficulties over iteration as a preliminary step typically requires exhaustive task evaluation, which is practically unaffordable in computation and data-annotation. This study provides a novel perspective to illuminate the possibility of leveraging the dual importance of adaptation robustness and learning efficiency, particularly in scenarios where task evaluation is risky or costly, such as iterative agent-environment interactions for robotic policy evaluation or computationally intensive inference steps for finetuning foundation models. Firstly, we introduce Model Predictive Task Sampling (MPTS), a framework that bridges the task space and adaptation risk distributions, providing a theoretical foundation for robust active task sampling. MPTS employs a generative model to characterize the episodic optimization process and predicts task-specific adaptation risk via posterior inference. The resulting risk predictive model amortizes the costly evaluation of task adaptation performance and provably approximates task difficulty rankings. MPTS seamlessly integrates into zero-shot, few-shot, and supervised finetuning settings. Empirically, we conduct extensive experiments in pattern recognition using foundation models and sequential decision-making. Our results demonstrate that MPTS significantly enhances adaptation robustness for tail risk or out-of-distribution (OOD) tasks and improves learning efficiency compared to state-of-the-art (SoTA) methods. The code is available at the project site https://github.com/thu-rllab/MPTS.
Figures & tables
Figure 1: Framework of MPTS in Adaptation Learning. a. The left shows standard random sampling for generating candidate tasks. The middle in blue denotes costly evaluation of B^ tasks (e.g., agent–environment interaction or foundation model forward pass) in DRM to select Top- B worst ones. The middle in green depicts MPTS, which predicts task difficulty via a lightweight generative model, avoiding expensive evaluation. The right illustrates the standard optimization pipeline in Meta-Learning, DR, or SFT (Snow: frozen models; Fire: updated models). b. MPTS samples candidate tasks, ranks their difficulty via the risk predictive model’s predictions for subset selection, and updates the learner, approximating CVaRα ’s Monte Carlo optimization in a predictive manner. The gathered risk signals further update RPMs online. c. The RPM utilizes the risk history H1:t to train under a streaming VI framework. d. The RPM simulates adaptation outcomes p(ℓ∣τ,H1:t;θt) for B^ candidate identifiers, computes acquisition scores, and selects the Top- B identifiers for the (t+1) -th iteration.
Figure 2: Fundamental Concepts: Task Identifiers, Episodic Learning and Probabilistic Graphical Models. a. The task distribution is uniform and defined over meaningful identifiers τ . For example, the amplitude and the phase [a,b] specifies a sinusoid curve to complete with K-shot observed data points. Robots like Half-Cheetahs are trained to accomplish different locomotion tasks with varying masses and velocities. Some multimodal pattern recognition tasks’ identifiers are implicit but can be described from a reference model, e.g., text encoders in CLIP 67 . b. The tail task generalization corresponds to CVaRα , i.e., the integral of tail task risk values in red. In OOD generalization, this work prompt-tunes CLIP on ImageNet 77 to test on ImageNet-S 92 . c. Here, the generative model includes grey units as observed variables and white ones as unobservable. The solid directed lines describe the generative model 90 . We use the dash-directed lines to indicate the recognition model and approximate inference within autoencoding variational Bayes 48 .
Figure 3: K-shot Sinusoid Regression Results (10 Runs). Note that a - i are results with MAML as the backbone, while j - n reports the results with Reptile as the backbone. a. Shown are curves of averaged MSEs with standard error of the mean (SEM) on the validation task set during meta-training for all methods. b. Curves illustrate the CVaR0.9 MSEs with SEM on the validation task set during the meta-training process. c. The meta-trained machine learners are tested on a fixed task set, reporting the average MSEs and CVaR values with SEM. d. Displayed are meta-testing results with SEM using MPTS machine learners trained by various γ1/γ0 ratios. e. Displayed are meta-testing results using MPTS machine learners trained in various pseudo batch sizes, i.e., B^={1B,2B,4B,8B} . f. The PCC values with SEM are tracked during meta-training. g. At a specific iteration, the statistical correlation between predicted and exact adaptation risk values of the task batch is visualized with overall ρℓˉ,ℓ=0.669 . h. The required relative run-time is computed for all methods during meta-training with ERM as the anchor. i. At some meta-training time-step, we visualize the subset selection from the pseudo batch under the RPM. j. We illustrate the curves of averaged MSEs with SEM on the validation task set during Reptile meta-training for all relevant baselines. k. We track the corresponding CVaR0.9 MSEs with SEM on the validation task set throughout the Reptile meta-training. l. Reported are the tested average MSEs and CVaR values with SEM using the meta-trained machine learners on a fixed task set. m. The PCC values with SEM are tracked for MPTS during Reptile meta-training. n. We compute the relative run-time for all methods during meta-training with ERM as the anchor.
Dataset
Metrics
ERM
DRM
GDRM
OHTM
DATS
TDPS
EMTL
MPTS (Ours)
Δ↑
ImageNet-CG 39
CVaR0.9
76.87 ± 0.31
77.47 ± 0.24
77.56 ± 0.19
76.52 ± 0.14
77.55 ± 0.21
77.28 ± 0.17
77.59 ± 0.25
78.16 ± 0.25
+1.29
CVaR0.7
82.04 ± 0.21
82.47 ± 0.20
82.62 ± 0.20
81.34 ± 0.17
82.31 ± 0.18
82.33 ± 0.24
82.56 ± 0.14
82.95 ± 0.26
+0.91
CVaR0.5
84.62 ± 0.24
84.99 ± 0.17
85.16 ± 0.22
84.06 ± 0.29
85.19 ± 0.14
85.29 ± 0.23
85.14 ± 0.17
85.33 ± 0.26
+0.71
Avg
89.22 ± 0.22
89.54 ± 0.16
89.49 ± 0.18
88.59 ± 0.17
89.45 ± 0.12
89.69 ± 0.11
89.34 ± 0.21
89.86 ± 0.14
+0.64
ImageNet-CI 39
CVaR0.9
80.40 ± 0.18
80.49 ± 0.22
80.26 ± 0.19
77.08 ± 0.18
78.86 ± 0.15
78.07 ± 0.19
80.41 ± 0.16
80.94 ± 0.19
+0.54
CVaR0.7
84.98 ± 0.17
85.05 ± 0.19
84.90 ± 0.22
81.60 ± 0.11
83.66 ± 0.22
83.16 ± 0.11
85.04 ± 0.18
85.70 ± 0.24
+0.72
Table 1: 5-Way 1-Shot Meta-Testing Classification Results on Various Datasets and Efficiency Comparison (10 Runs). We report testing CVaR0.9 , CVaR0.7 , CVaR0.5 and average accuracies with corresponding SEMs evaluated by the meta-trained machine learner on ID and OOD datasets. The best results are in bold with the runner-up underlined, and MPTS’s performance gains over ERM Δ↑ are marked in blue. With experiments on ImageNet-A as an example, we report the memory cost and runtime relative to ERM during meta-training.
Figure 4: Meta-RL Results on Mujoco Environments (10 Runs). a. The cumulative returns with SEM belonging to CVaR0.9 validation MDPs are displayed during meta-training. b. We compute the average cumulative returns with SEMs on validation MDPs during meta-training. c. Tracked are the RPM’s PCC values with SEMs over training iterations. d. The relative clock time quantifies the computational complexity for all methods on Walker2dVel, where ERM’s runtime works as the anchor. e. We report CVaRα returns with SEM of meta-testing MDPs. f. The box-plot reports results averaged over meta-testing MDPs. g. With PEARL 70 as the Meta-RL backbone, we illustrate the learning curves and meta-testing results with SEM on HalfCheetahBody and HalfCheetahMass from RoML 33 baseline.
Figure 5: DR Results on Ergo-Reacher and Lunar-Lander (10 Runs). a. In Ergo-Reacher, the CVaR0.9 , CVaR0.7 , CVaR0.5 , and average cumulative returns with SEM on validation MDPs are reported. The RPM’s PCC values with SEM are tracked during DR training. b. In Lunar-Lander, the cumulative returns on validation MDPs with SEM are illustrated. The RPM’s PCC values with SEM are tracked during DR training. c. We test the DR-trained policies on the fixed MDP set and report the CVaRα cumulative returns with SEM. d. The returns averaged over DR-testing MDPs are illustrated. e. The required runtime is computed for all methods on Lunar-Lander. f. In Lunar-Lander, shown are frequencies of sampled identifiers using MPTS during DR training. g. In Lunar-Lander, we test the trained policies in both in-distribution (ID) domains and out-of-distribution (OOD) domains to report each task’s average returns with SEM.
Dataset
Metrics
ERM
DRM
GDRM
OHTM
MPTS (Ours)
Δ↑
ImageNet 77
CVaR0.9
31.70 ± 0.12
32.23 ± 0.15
31.37 ± 0.18
31.53 ± 0.21
32.52 ± 0.16
+0.82
(ID)
CVaR0.7
42.87 ± 0.14
44.10 ± 0.14
42.97 ± 0.18
43.28 ± 0.18
44.28 ± 0.14
+1.41
CVaR0.5
51.43 ± 0.10
52.61 ± 0.11
51.75 ± 0.19
51.96 ± 0.16
52.72 ± 0.15
+1.29
Avg
70.80 ± 0.08
70.90 ± 0.10
71.00 ± 0.08
70.80 ± 0.10
71.20 ± 0.09
+0.40
ImageNet-A 40
CVaR0.9
15.35 ± 0.21
15.50 ± 0.25
15.62 ± 0.22
15.44 ± 0.25
18.38 ± 0.22
+3.03
(OOD)
CVaR0.7
22.82 ± 0.23
23.05 ± 0.21
23.13 ± 0.19
23.00 ± 0.23
24.05 ± 0.20
+1.23
Table 2: Testing Classification Results after Prompt-Tuning on ImageNet and Efficiency Comparison (10 Runs). We report testing CVaR0.9 , CVaR0.7 , CVaR0.5 and average accuracies together with corresponding SEMs evaluated by the prompt-tuned machine learner on ID and OOD datasets. Evaluation on OOD datasets corresponds to the domain generalization setting. The best results are in bold with the runner-up underlined, and MPTS’s performance gains over ERM Δ↑ are marked in blue. During prompt-tuning ImageNet, we report the memory cost and runtime relative to ERM for all methods.
Appendix figures & tables11 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 8Figure 9Figure 10
Figure S1 : Risk Predictive Model as the World Model to Guide Conditional Interaction. This shows the explanation of MPTS in the presence of active inference theory.
Figure S2 : Risk Predictive Module in MPTS for Active Subset Selection. MPTS adopts a predict-then-optimize strategy and uses a predictive module in green to approximately score the task subset difficulty and obtain the preferred task subset. While the traditional method in blue exhausts B^ tasks in construction and evaluation to filter preferred subset.
Figure S3 : Comparison between Traditional Multi-Task Learning and Meta Learning and Domain Randomization. Traditional multi-task learning aims to exploit the correlation between heterogeneous tasks, and task IDs in a batch do not change 36 . In comparison, meta learning and domain randomization typically consider the homogeneous task family and resample task IDs in a batch to optimize. The latter is our studied scenario, referred to as task episodic optimization. [Same types of colors mean the homogeneous tasks with depth to denote the various episode.]
ERM
DRM
GDRM
OHTM
MPTS (Ours)
computation
O(∣θ∣TERM)
O(1−α^1∣θ∣TDRM)
O(∣θ∣TGDRM)
O(∣θ∣TOHTM)
O((∣ϕ∣+∣ψ∣+∣θ∣)TMPTS)
task eval
O(BTERM)
O(1−α^BTDRM)
O(BTGDRM)
O(BTOHTM)
O(BTMPTS)
Appendix
Table S1 : Computational Complexities using Different Methods. Here, we drop out the ranking or reweighting computational complexity as the model complexity of the machine learner considered in this analysis is major, such as the multimodal foundation models. T refers to the required iteration steps until the convergence for separate methods.
Benchmarks
Identifier Meaning
Identifier Range
K-shot sinusoid regression
amplitude and phase (a,b)
[0.1,5.0]×[0,π]
Meta-RL: HalfCheetahMassVel
mass and velocity (m,v)
[0.75,1.25]×[0,2.0]
Meta-RL: HalfCheetahVel
velocity v
[0,2.0]
Meta-RL: ReacherPos
goal location (x1,x2)
[−0.2,0.2]×[−0.2,0.2]
Meta-RL: Walker2dMassVel
mass and velocity (m,v)
[0.75,1.25]×[0,2.0]
Meta-RL: Walker2dVel
velocity v
[0,2.0]
Appendix
Table S2 : Benchmarks with Explicit Task Identifiers. Here, we list the detail information about the task identifier to induce the task distribution.
Figure S4 : The Encoder-Decoder Neural Network to Paramterize the RPM.
Figure S5 : Rank Flipping after One-Step Update. We number tasks from the sampled pseudo task batch Tt+1B^ for the (t+1) -th subset selection. The left figure partitions the set into U∪V as illustrated in the Proof of Lemma 8 , with the blue ones, i.e., 1, 2, 3 as Top- B tasks evaluated under θt and the red ones as the complementary set. Note that tasks 2 and 8 are swapped in relative Top- B ranking after one-step model update; hence, the right figure depicts the Top- B tasks as 1,3, and 8. The gray lines indicate possible order-flipped cross-pairs between U and V . The blue and the red lines are probably flipped cross-pairs in separate partitioned sets’ induced graphs.
Benchmarks
Adaptation
Backbone
Expensive Part
K-shot sinusoid regression
few-shot
MAML/Reptile
computations
N-way K-shot image classification
few-shot
MaPLe
computation/memory
Meta-RL
few-shot
MAML/PEARL
interactions
Robotic DR
zero-shot
TD3
interactions
SFT
many-shot
MaPLe
computation/memory
Appendix
Table S3 : A Summary of the Considered Benchmarks. Here, we list the primary expensive part in task episodic learning for each scenario together with backbone methods. Also note that N-way K-shot image classification and SFT requires implicit task identifiers while others can directly access explicit task identifiers as the lower dimensional task representation.
Reinforcement learning (RL) is essential for adaptive decision-making, e.g., robust Meta-RL and domain randomization (DR), and efficient post-training of large language models (LLMs), where tasks are commonly structured as Markov decision processes or prompts. A central bottleneck in all these settings is the cost of agent-environment interactions required for task selection and policy optimization. Model predictive task sampling (MPTS) has emerged as a promising approach to improve sample efficiency by querying informative tasks via a risk predictive model (RPM) trained on optimization history. However, existing RPMs suffer from two failure modes: (i) sparse optimization histories, which limit RPMs' generalization across the task space, and (ii) a one-step lag in risk approximation, which causes biased difficulty estimation as the policy evolves. This work formalizes these failure modes, derives lagged-risk decompositions that expose the resulting temporal bias, and recasts risk prediction as an in-context sequence modeling problem. The developed RPM integrates a one-step look-ahead mechanism with a temporal-weighted sliding window to mitigate data non-stationarity and scarcity, without requiring additional rollouts. Empirically, our approach yields more reliable difficulty estimates and consistent performance gains across robust Meta-RL, DR, and prompt curriculum RL post-training of LLMs.
Jiarong Wen, Qi Tao, Kaiyu Zhang +5
College of Science, National University of Defense Technology · Department of Automation, Tsinghua University
Generalist robot policies carry broad manipulation priors from large-scale data, but specializing them to a new task remains the deployment bottleneck. This requires eliciting task-specific behavior from limited demonstrations without degrading their broad capabilities. We introduce Proxy Policy Steering (PPS), an inference-time adaptation method that resolves this challenge by training two lightweight proxy policies whose calibrated velocity-space difference steers the frozen base sampler. A reference proxy models the frozen base's behavior on target-task observations, and a task proxy, initialized from the reference, captures how this behavior changes under task supervision. Their difference forms a calibrated velocity-space residual that steers the frozen base sampler at every denoising step. We identify the conditions under which this residual isolates the change induced by task supervision, and validate them empirically. Because the base is never directly modified, its broad capabilities remain available at inference, including behaviors such as recovery from failure that the demonstrations themselves do not exercise. Adaptation requires only forward velocity predictions from the base, making PPS lightweight to train and applicable even without access to the base's parameters. On 8 real-world and 4 simulation manipulation tasks, PPS lifts the state-of-the-art pi 0.5 base policy by 53% absolute success rate on average, with zero-to-one gains on tasks the base never solves, while preserving the base's broad capabilities. PPS outperforms LoRA fine-tuning, from-scratch specialists, residual policies, and prior inference-time steering methods.
Multi-task reinforcement learning (MTRL) aims to train a single agent to efficiently optimize performance across multiple tasks simultaneously. However, jointly optimizing all tasks often yields imbalanced learning: agents quickly solve easy tasks but learn slowly on harder ones. While prior work primarily attributes this imbalance to conflicting task gradients and proposes gradient manipulation or specialized architectures to address it, we instead focus on a distinct and under-explored challenge: imbalanced data allocation. Standard MTRL allocates an equal number of environment interactions to each task, which over-allocates data to easy tasks that require relatively few interactions to solve and under-allocates data to hard tasks that require substantially more experience to solve. To address this challenge, we introduce Distributionally Robust Adaptive Task Sampling (DRATS), an algorithm that adaptively prioritizes sampling tasks furthest from being solved. We derive DRATS by formalizing MTRL as a feasibility problem from which we derive a minimax objective for minimizing the worst-case return gap, the difference between a desired target return and the agent's return on a task. In benchmarks like MetaWorld-MT10 and MT50, DRATS improves data efficiency and increases worst-task performance compared to existing task sampling algorithms.
Nicholas E. Corrado, Wenyuan Huang, Josiah P. Hanna
Computer Sciences Department · University of Wisconsin – Madison