Bridging Risk Approximation Gaps in Model Predictive Task Sampling via In-Context Modeling

Jiarong Wen; Qi Tao; Kaiyu Zhang; Yun Qu; Lizhou Cai; Heming Zou; Yuhang Jiang; Qi Wang

cs.AIOct 10, 2026

Bridging Risk Approximation Gaps in Model Predictive Task Sampling via In-Context Modeling

Authors: Jiarong Wen, Qi Tao, Kaiyu Zhang, Yun Qu, Lizhou Cai, Heming Zou, Yuhang Jiang, Qi Wang

Organizations: College of Science, National University of Defense Technology · Department of Automation, Tsinghua University

Abstract

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 $\bm{\theta}_{t}$ tends to cause the rank-flipping at $\bm{\theta}_{t+1}$ , i.e., $\text{sgn}[\ell_{t}(\bm{\tau}_{i})-\ell_{t}(\bm{\tau}_{j})]\neq\text{sgn}[\ell_{t+1}(\bm{\tau}_{i})-\ell_{t+1}(\bm{\tau}_{j})]$ .

Figure 2: Left : The risk landscape evolves from $\bm{\theta}_{t}$ to $\bm{\theta}_{t+1}$ , with historical prompt-loss pairs $H_{t}$ 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 $\hat{\ell}_{t+1}$ by minimizing the discrepancy with ground-truth $\ell_{t+1}$ through gradient optimization.

Figure 3: Few-shot sinusoid regression results (10 Runs, Mean

\pm

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

\pm

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 $\pm$ 3.9	-209.9 $\pm$ 7.6	23.5 $\pm$ 15.6	-34.2 $\pm$ 20.7	-17.5 $\pm$ 1.9	-31.1 $\pm$ 5.4	243.0 $\pm$ 14.2	185.5 $\pm$ 30.4	83.8 $\pm$ 0.7	70.1 $\pm$ 1.4
DRM [ 17 , 32 ]	-151.1 $\pm$ 6.2	-195.5 $\pm$ 8.4	26.1 $\pm$ 12.2	-26.0 $\pm$ 14.7	-16.0 $\pm$ 2.2	-26.5 $\pm$ 5.6	184.2 $\pm$ 25.0	96.8 $\pm$ 36.6	85.3 $\pm$ 0.4	72.9 $\pm$ 1.0
GDRM [ 31 ]	-149.8 $\pm$ 4.3	-206.4 $\pm$ 8.0	16.1 $\pm$ 13.3	-42.8 $\pm$ 16.6	-19.6 $\pm$ 4.3	-30.7 $\pm$ 6.1	207.9 $\pm$ 18.4	93.7 $\pm$ 38.0	84.4 $\pm$ 0.7	70.2 $\pm$ 1.7
OHTM [ 33 ]	-160.0 $\pm$ 7.3	-194.8 $\pm$ 10.3	7.6 $\pm$ 21.5	-40.8 $\pm$ 21.3	-16.3 $\pm$ 1.4	-28.9 $\pm$ 5.6	258.7 $\pm$ 6.4	215.7 $\pm$ 17.9	83.8 $\pm$ 0.7	72.6 $\pm$ 2.0
MPTS [ 22 ]	-143.6 $\pm$ 4.6	-184.0 $\pm$ 8.3	24.6 $\pm$ 16.2	-27.9 $\pm$ 19.1	-20.1 $\pm$ 3.4	-33.2 $\pm$ 8.0	268.5 $\pm$ 3.4	251.8 $\pm$ 4.1	84.7 $\pm$ 0.5	75.7 $\pm$ 1.1

Table 1: Evaluation return on five robotic control tasks. (10 Runs, Mean

\pm

SEM) The best results are bold, with the runner-up underlined.

Figure 5: CVaR 0.9 return vs. relative runtime on Walker2dVel.

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

\gamma

on Lunar-Lander (10 Runs). Left : Sensitivity to look-back window

W

. CVaR 0.9 returns, average returns under

W\in\{1,2,3,8\}

. Right : Sensitivity to decay factor

\gamma

. CVaR 0.9 returns, average returns under

\gamma\in\{0.5,0.7,0.9,1.0\}

Appendix figures & tables3 assets

Supplementary material from the paper’s appendix.

Appendix

Hyperparameter	Value
Meta-training iterations	500
Meta-batch size $B$	20 tasks
Validation batch size	20 tasks (uniform)
Inner adaptation steps	1
Fast batch size (trajectories / task)	20
Max KL divergence	0.01

Appendix

Table 3: Shared meta-learning hyperparameters (Meta-RL).

Hyperparameter	Value
Base model	Qwen3-4B-Base
Training steps	100
Prompt batch size $B$	128 (4-GPU)
Responses per prompt $K$	8
Max response length	1,024 tokens
Policy optimiser	AdamW, lr $1\times 10^{-6}$

Appendix

Table 4: RLVR training hyperparameters.

Experiment	Hardware	Common	Slower
Sinusoid	RTX 4090 24GB	$\sim$ 12.2 min	DRM: $\sim$ 16.4 min
DR	RTX 4090 24GB	$\sim$ 3.0 h	DRM: $\sim$ 4.2 h
Meta-RL	A10 24GB	$\sim$ 8.7 h	DRM: $\sim$ 11.6 h
RLVR	4 $\times$ A100 80GB	$\sim$ 11.1 h	GRESO: $\sim$ 17.8 h

Appendix

Table 5: Approximate single-run compute resources.

Sample-Efficient RL

Sample-efficient reinforcement learning studies how agents learn effective policies from fewer environment interactions, demonstrations, or costly rollouts.

See all 248 papers

Model Predictive Task Sampling for Efficient and Robust Adaptation
Jan 19, 2025Qi Wang, Zehao Xiao, Yixiu Mao +4Adaptive Sampling Distribution Shift Robustness
Department of Automation, Tsinghua University · Informatics Institute, University of Amsterdam
Beyond Solvability: Task Learnability as a Static Prior for LLM RL Post-Training
Aug 10, 2026Ting Zhou, Zhenqing Ling, Daoyuan Chen +4Reinforcement Learning RL for Language Model Reasoning
Sun Yat-Sen University · Alibaba Group
Dynamics-Predictive Sampling for Active RL Finetuning of Large Reasoning Models
Mar 11, 2026Yixiu Mao, Yun Qu, Qi Wang +2Prompt Selection for RL RL for Language Model Reasoning
Department of Automation, Tsinghua University

Jan 19, 2025J/K move · Enter open · S save

Model Predictive Task Sampling for Efficient and Robust Adaptation

Qi Wang, Zehao Xiao, Yixiu Mao, Yun Qu+3

Department of Automation, Tsinghua University · Informatics Institute, University of Amsterdam

Adaptive Sampling Distribution Shift Robustness

Open paper AskPDF arXiv

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.

Abstract

Sample-Efficient RL

Explore similar work

Model Predictive Task Sampling for Efficient and Robust Adaptation

Beyond Solvability: Task Learnability as a Static Prior for LLM RL Post-Training

Dynamics-Predictive Sampling for Active RL Finetuning of Large Reasoning Models