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.
Figures & tables
Figure 1: Task Risk Distribution Shift over Iterations. Using the risk generalization results at θt tends to cause the rank-flipping at θt+1 , i.e., sgn[ℓt(τi)−ℓt(τj)]=sgn[ℓt+1(τi)−ℓt+1(τj)] .
Figure 2: Left : The risk landscape evolves from θt to θt+1 , with historical prompt-loss pairs Ht modeling temporal risk dynamics and predicting future risk using an attention mechanism. Right : The Transformer-based ICRPM tokenizes prompt-loss sequences with a target prompt at time t+1 , and a regression head estimates the future risk ℓ^t+1 by minimizing the discrepancy with ground-truth ℓt+1 through gradient optimization.
Figure 3: Few-shot sinusoid regression results (10 Runs, Mean ± Standard Error of the Mean (SEM)). From left to right: (a) validation CVaR 0.9 MSE curves during meta-training; (b) validation average MSE curves during meta-training; (c) SRC of predicted vs. ground-truth difficulty during training; (d) meta-testing CVaR α MSE at multiple robustness levels. ICRPM (ours) achieves consistently lower tail-risk error and maintains higher SRC throughout training.
Figure 4: DR (columns 1–2) and Meta-RL (columns 3–5) learning curves and correlation trends (10 Runs, Mean ± SEM). Rows from top to bottom: CVaR 0.9 returns, average returns, and SRC of predicted vs. ground-truth task risk. ICRPM consistently attains the highest tail-risk performance while maintaining competitive average returns, alongside superior rank correlation throughout training.
Method
HalfCheetahVel
Walker2dVel
ReacherPos
LunarLander
ErgoReacher
CVaR 0.5
CVaR 0.9
CVaR 0.5
CVaR 0.9
CVaR 0.5
CVaR 0.9
CVaR 0.5
CVaR 0.9
CVaR 0.5
CVaR 0.9
ERM [ 30 ]
-153.6 ± 3.9
-209.9 ± 7.6
23.5 ± 15.6
-34.2 ± 20.7
-17.5 ± 1.9
-31.1 ± 5.4
243.0 ± 14.2
185.5 ± 30.4
83.8 ± 0.7
70.1 ± 1.4
DRM [ 17 , 32 ]
-151.1 ± 6.2
-195.5 ± 8.4
26.1 ± 12.2
-26.0 ± 14.7
-16.0 ± 2.2
-26.5 ± 5.6
184.2 ± 25.0
96.8 ± 36.6
85.3 ± 0.4
72.9 ± 1.0
GDRM [ 31 ]
-149.8 ± 4.3
-206.4 ± 8.0
16.1 ± 13.3
-42.8 ± 16.6
-19.6 ± 4.3
-30.7 ± 6.1
207.9 ± 18.4
93.7 ± 38.0
84.4 ± 0.7
70.2 ± 1.7
OHTM [ 33 ]
-160.0 ± 7.3
-194.8 ± 10.3
7.6 ± 21.5
-40.8 ± 21.3
-16.3 ± 1.4
-28.9 ± 5.6
258.7 ± 6.4
215.7 ± 17.9
83.8 ± 0.7
72.6 ± 2.0
MPTS [ 22 ]
-143.6 ± 4.6
-184.0 ± 8.3
24.6 ± 16.2
-27.9 ± 19.1
-20.1 ± 3.4
-33.2 ± 8.0
268.5 ± 3.4
251.8 ± 4.1
84.7 ± 0.5
75.7 ± 1.1
Table 1: Evaluation return on five robotic control tasks. (10 Runs, Mean ± SEM) The best results are bold, with the runner-up underlined.
Figure 5: CVaR 0.9 return vs. relative runtime on Walker2dVel.
Figure 7
Figure 7: Ablation studies on Lunar-Lander (10 Runs). Left : Impact of acquisition principles. Two ICRPM variants (worst-case and diversity-aware) versus PDTS [ 23 ] . Right : Architecture ablation comparing ICRPM with MLP or Transformer backbone, vanilla MPTS, and MPTS with history augmentation.
Figure 8: Sensitivity to W and γ on Lunar-Lander (10 Runs). Left : Sensitivity to look-back window W . CVaR 0.9 returns, average returns under W∈{1,2,3,8} . Right : Sensitivity to decay factor γ . CVaR 0.9 returns, average returns under γ∈{0.5,0.7,0.9,1.0} .
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.
Qi Wang, Zehao Xiao, Yixiu Mao +4
Department of Automation, Tsinghua University · Informatics Institute, University of Amsterdam
Reinforcement learning (RL) has become a central post-training paradigm for eliciting reasoning capabilities in large language models, yet uniform task sampling allocates compute without regard to differences in how tasks respond to optimization. Existing task-valuation methods mostly rely on snapshot-based signals such as current pass rate or reward, which estimate how solvable a task is under the current policy. However, tasks with similar current solvability can still differ substantially in how positively they respond to further training. We study this residual axis as task learnability: a regime-conditional measure of expected positive response to continued training under a fixed RL post-training regime. By analyzing per-task reward trajectories, we find that learnability is reproducible across independently sampled training contexts and predictive of downstream utility. To make this signal practical before training begins, we propose TrajVal, a lightweight probe-based estimator that approximates per-task learnability from a short probe run and two endpoint evaluations. TrajVal can be used either as a standalone static prior for task sampling or as a multiplicative prior for existing online schedulers. Experiments on mathematical and logical reasoning benchmarks across multiple model scales show that TrajVal improves data efficiency over uniform sampling and provides complementary gains when combined with online scheduling methods.
Reinforcement learning (RL) finetuning has become a key technique for enhancing the reasoning abilities of large language models (LLMs). However, its effectiveness critically depends on the selection of training data. Recent advances underscore the importance of online prompt selection methods, which typically concentrate training on partially solved or moderately challenging examples under the current policy, thereby yielding more effective model updates. While significantly accelerating RL finetuning in terms of training steps, they also incur substantial computational overhead by requiring extensive LLM rollouts over large candidate batches to identify informative samples, an expense that can outweigh the finetuning process itself. To address this challenge, this work proposes Dynamics-Predictive Sampling (DPS), which online predicts and selects informative prompts by inferring their learning dynamics prior to costly rollouts. Specifically, we introduce a new perspective by modeling each prompt's solving progress during RL finetuning as a dynamical system, where the extent of solving is represented as the state and the transition is characterized by a hidden Markov model. Using historical rollout reward signals, we perform online Bayesian inference to estimate evolving state distributions, and the inference outcome provides a predictive prior for efficient prompt selection without rollout-intensive filtering. Empirical results across diverse reasoning tasks, including mathematics, planning, and visual geometry, demonstrate that DPS substantially reduces redundant rollouts, accelerates the training process, and achieves superior reasoning performance.