Paper deep dive
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
Intelligence
Status: not_run | Model: - | Prompt: - | Confidence: 0%
Entities (0)
Relation Signals (0)
No relation signals yet.
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
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.
Tags
Links
- Source: https://arxiv.org/abs/2610.14913v1
- Canonical: https://arxiv.org/abs/2610.14913v1
Trouble viewing inline? Open PDF directly →
Full Text
91,619 characters extracted from source content.
Expand or collapse full text
Bridging Risk Approximation Gaps in Model Predictive Task Sampling via In-Context Modeling Jiarong Wen1,† Qi Tao1,† Kaiyu Zhang1,† Yun Qu2 Lizhou Cai2 Affiliation: Heming Zou2, Yuhang Jiang2, Qi Wang2 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 (i) 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. 1 Introduction Over the past decade, deep reinforcement learning (RL) has achieved remarkable progress in complex sequential decision-making [1, 2, 3, 4, 5]. This work focuses on two settings where RL operates over structured task distributions. (i) In robotics and physical control, agents must adapt to new Markov decision processes (MDPs) in zero-shot or few-shot scenarios under safety constraints [6, 7]. Robust domain randomization (DR) [8] and robust meta-reinforcement learning (Meta-RL) [9] address this by optimizing over MDP distributions and prioritizing risky environments. (i) In foundation model post-training, RL has become central to incentivizing reasoning in large language models (LLMs) [10, 11], where identifying informative prompts and rollouts is critical for stable optimization [12, 13, 14, 15]. Optimization over prompts further requires policies to explore complex logical search spaces through long-horizon token generation. Evaluation Bottlenecks in Task Prioritization. Despite their different objectives, both settings share the same bottleneck: expensive policy evaluation. In robust DR and Meta-RL, assessing task risk requires extensive agent-environment trajectories on selected tasks [16, 17, 18, 19]; in LLM post-training, each update requires long Chain-of-Thought (CoT) rollouts [20], making prompt difficulty evaluation costly [12, 13, 21]. Thus, exhaustive task evaluation is impractical, motivating efficient task prioritization. Conditional Agent-Environment Interactions via Model Predictive Strategies. Model predictive task sampling (MPTS) [22] and its variants [13, 14, 23, 24] reduce unnecessary interactions by using optimization histories, e.g., task or prompt identifiers paired with losses, to train an RPM that approximates task risk online. This enables conditional interactions: policy updates target selected tasks without exhaustive rollouts. However, RPM fidelity remains underexplored despite directly governing task selection and policy evolution. Vanilla RPMs [22, 23] suffer from two key limitations: (i) temporal lag bias, where outdated policy snapshots misrepresent the current risk landscape; and (i) sparse-history generalization, where limited interaction data restricts generalization across the task space. Analogous issues appear in prompt curriculum learning for LLM post-training, e.g., PCL [14], where temporal lag and sparse evaluations distort curricula and reduce sample efficiency. Enhancing Risk Predictive Models via In-Context Sequence Modeling. We analyze vanilla RPMs and quantify extrapolation errors caused by temporal lag and data scarcity under small task batches. To address these issues, we propose two remedies: (i) a one-step look-ahead mechanism that conditions risk prediction on recent optimization signals to reduce extrapolation bias; and (i) a temporal-weighted sliding window that augments training with recent task trajectories while exponentially downweighting stale data. Together, they instantiate an in-context model of the evolving task risk function, improving prediction under non-stationary policy dynamics with negligible additional rollout cost. Primary Contributions and Empirical Findings. Focusing on RPM design for the MPTS family, our contributions are two-fold: 1. Theoretical Analysis of MPTS Vulnerabilities: We derive a diagnostic lagged-risk decomposition that exposes finite-sample reconstruction error and a policy-update-induced (η)O(η) mismatch, explaining why interpolation-based RPMs can destabilize task-risk rankings. 2. In-Context Modeling of the Evolving Risk Function: We introduce in-context RPM (ICRPM) to characterize optimization histories, with one-step look-ahead prediction as the core design and a temporal-weighted sliding window as lightweight data reuse to reduce extrapolation bias and alleviate data scarcity without modifying the policy optimization loop. Experiments within the MPTS framework show that ICRPM improves sample efficiency and tail-risk robustness on robust DR and Meta-RL benchmarks, and accelerates curriculum-based RL post-training for LLMs, demonstrating better generalization across continuous-control and LLM reasoning tasks. 2 Research Background Notation and Setup. We consider a unified task space T that encompasses both continuous control and LLM post-training regimes. Each task ∈ τ is formalized as a MDP [1]. In DR and Meta-RL, τ is represented as a continuous descriptor in ℝdR^d, whereas in reinforcement learning with verifiable rewards (RLVR), τ corresponds to a discrete reasoning prompt. At optimization step t, the agent is parameterized by t∈Θ θ_t∈ , representing either a control policy or an autoregressive model πt _ θ_t. Over the course of training adaptive decision-makers or RL finetuning LLMs, we sample a batch of tasks tB=t,ii=1B⊂T_t^B=\ τ_t,i\_i=1^B for the agent to interact. Let ℒ:Θ×→ℝL: ×T denote the policy-dependent risk, and define ℓt():=ℒ(t,) _t( τ):=L( θ_t, τ) as the risk of policy πt _ θ_t on task τ at iteration t. Concretely, ℓt() _t( τ) can be the negative episodic return in control settings, ℓt()=−πt[∑i=0H−1λ¯irt,i] _t( τ)=-E_ _ θ_t\! [ _i=0^H-1 λ^ir_t,i τ ] under post-adapted policies πt _ θ_t, or the expected penalty in RLVR, e.g., ℓt()=1−[rt] _t( τ)=1-E[r_t τ]. In the MPTS setup [22, 23, 13, 14], the per-iteration optimization batch is Ht=(t,i,ℓt(t,i))i=1BH_t=\( τ_t,i, _t( τ_t,i))\_i=1^B, and the accumulated form H0:t=⋃k=0tHkH_0:t= _k=0^tH_k tracks the temporal optimization history. We further denote the space of complete historical trajectories by ℋH and define the RPM as a functional ℓ^t:×ℋ→ℝ _t:T×H , which approximates the non-stationary task risk landscape conditioned on H0:tH_0:t. General MPTS Pipeline. MPTS [22, 23, 14] executes a predict-then-select pipeline, which can instantiate both active prompt selection and robust task sampling. For prompt curriculum learning in RLVR [13, 14], the RPM predicts the sampled success rate ρ^t ρ_ τ^t, and prompts with intermediate difficulty, i.e., those close to a target threshold ρ∗ρ^*, typically around 0.50.5, are preferred [25]. For robust task sampling, the same rule is applied to the predicted risk by replacing ρ^t ρ_ τ^t with ℓ^t() _t( τ) and setting ρ∗=max∈tB^ℓ^t()ρ^*= _ τ _t B _t( τ), which encourages near-worst-case task selection. At each iteration, MPTS proceeds in three steps: ① Online RPM update. The RPM is refined with the latest history Ht−1H_t-1 by maximizing the likelihood of risk prediction [14] or its approximation [22], e.g., the ELBO in Eq. (2). ② Amortized evaluation. Let B B denote candidate-pool size, a candidate pool tB^∼p()T_t B p( τ) is sampled, and the forward passes through the RPM produce predictive risk for all candidate tasks. ③ Subset selection. The training task batch is constructed by selecting B tasks with predicted risk closest to the target threshold: Amortized Train RPM: max∈Ψ,ϕ∈ΦℒRPM(,ϕ,Ht−1), _ ψ∈ ,\, φ∈ L_RPM( ψ, φ;H_t-1), (1a) Online Task Evaluation: p(ℓ∣^i,Ht−1;t)→MCρ^^it,m(ℓi),σ(ℓi)i=1B^, p_ ψ\! ( τ_i,H_t-1; θ_t ) MC \ ρ_ τ_i^t,\,m( _i),\,σ( _i) \_i=1 B, (1b) Active Subset Selection: tB=argtop−B∈tB^(−‖ρ^t−ρ∗‖22) _t^B= *arg\,top-B_ τ _t B (- \| ρ_ τ^t-ρ^* \|_2^2 ) (1c) In Eq. (1b), m(ℓi)m( _i) and σ(ℓi)σ( _i) are the predictive mean and standard deviation of the task outcome estimated via Monte Carlo sampling, which are then mapped to the acquisition scalar ρ^^it ρ_ τ_i^t to quantify task priority. With the queried batch in Eq. (1c), the policy model π _ θ is accordingly updated in the next iteration, e.g., t+1=t−η1B∑b=1B∇ℓt(b) θ_t+1= θ_t-η 1B _b=1^B _ θ _t( τ_b) in DR. By repeating these steps for T rounds, MPTS achieves either prompt curriculum learning in LLMs or robust Meta-RL and DR. RPM and Variational Inference. Existing MPTS frameworks typically instantiate the RPM as a VAE-based latent generative predictor [26], where a latent state tz_t summarizes historical task-risk observations. Using streaming variational inference [27, 28], the RPM maximizes the per-step ELBO: ELBO(,ϕ)=qϕ(t|Ht)[∑i=1Blogp(ℓt(t,i)|t,i,t)]−βKL(qϕ(t|Ht)||qϕ¯(t|Ht−1)),G_ELBO( ψ, φ)=E_q_ φ(z_t|H_t) [ _i=1^B p_ ψ( _t( τ_t,i)| τ_t,i,z_t) ]- (q_ φ(z_t|H_t)||q_ φ(z_t|H_t-1) ), (2) where ϕ¯ φ represents the parameters from the previous iteration, kept fixed to provide a stable prior. For task evaluation, RPMs use Monte Carlo forward passes to estimate p(ℓ|,H0:t;t)≈qϕ(t|Ht)[p(ℓ|,t;t)].p( | τ,H_0:t; θ_t) _q_ φ(z_t|H_t) [p_ ψ( | τ,z_t; θ_t) ]. These predictions are then used by the acquisition rule in Eq. (1c) for robust [22, 23] or curriculum optimization [13, 14]. Thus, vanilla RPMs rely on generalization at t θ_t to approximate task difficulty after the policy update to t+1 θ_t+1, creating the temporal mismatch analyzed next. 3 From Theoretical Investigation to Practical Enhancement Figure 1: Task Risk Distribution Shift over Iterations. Using the risk generalization results at t θ_t tends to cause the rank-flipping at t+1 θ_t+1, i.e., sgn[ℓt(i)−ℓt(j)]≠sgn[ℓt+1(i)−ℓt+1(j)]sgn[ _t( τ_i)- _t( τ_j)] [ _t+1( τ_i)- _t+1( τ_j)]. This section analyzes the theoretical bottlenecks of RPMs, introduces the In-Context Predictive Model (ICRPM), and discusses its implementation within MPTS. 3.1 Vulnerabilities of Vanilla RPM Although MPTS provides a sample-efficient surrogate for task risk estimation, the vanilla RPM overlooks a common failure mode, i.e., the continual evolution of both the policy and the task risk landscape induces rank-flipping (Fig. 1) in task risk that cannot be corrected post-hoc. Instability of Task Risk Rankings. MPTS relies on accurate task-risk ranking, but task risk is policy-dependent: after updating from t θ_t to t+1 θ_t+1, the ordering induced by ℓt() _t( τ) may differ from that of ℓt+1() _t+1( τ). As shown in Fig. 1, this policy-induced drift creates a non-stationary difficulty landscape. Prior analysis [22] shows that ranking stability requires sufficiently small learning rates and smooth loss landscapes, conditions often violated in practice. Hence, even a perfectly accurate RPM at t θ_t may fail to preserve the correct ranking at t+1 θ_t+1, which we call the task risk distribution shift. Lagged Risk Prediction and Reconstruction Bound. The vanilla RPM, a VAE-based risk generative model [22, 23], is trained on the latest history batch HtH_t of size B in each per-step update, where supervision targets are evaluated at the latest yet lagged parameter t θ_t. The induced structured risk distribution shift cannot be resolved by increasing sample size alone: even under standard capacity and smoothness assumptions, the prediction error includes a temporal mismatch term induced by the policy update from t θ_t to t+1 θ_t+1. Theorem 1 (Per-step Lagged Prediction Bound for Vanilla RPMs). Assume that ℓ()=ℒ(,)∈[0,M] ( τ)=L( θ, τ)∈[0,M] for all ∈Θ θ∈ , ∈ τ , and that ℓ() ( τ) is L-Lipschitz in θ. The policy update uses a normalized average loss computed by i.i.d. samples and ‖∇ℓt(i)‖≤G\| _ θ _t( τ_i)\|≤ G. Let ℓ^,ϕ,t() _ ψ, φ,t( τ) denote the predictive risk function induced by the vanilla RPM at iteration t. Define ϵapprec:=inf,ϕ∼p[(ℓ^,ϕ,t()−ℓt())2]. _app^rec:= _ ψ, φE_ τ p [ ( _ ψ, φ,t( τ)- _t( τ) )^2 ]. Then, for any δ∈(0,1)δ∈(0,1), with probability at least 1−δ1-δ over the sampling of HtH_t, ∼p[|ℓ^t()−ℓt+1()|]≤ϵapprec+ϵgenrec(B,δ)⏟ 1⃝Reconstruction Approximation Error+LηG⏟ 2⃝Temporal Lagging Error.E_ τ p [| _t( τ)- _t+1( τ)| ]≤ _app^rec+ _gen^rec(B,δ)_ 1Reconstruction Approximation Error + Lη G_ 2Temporal Lagging Error. (3) If the reconstruction loss class rec=(,y)↦(ℓ^,ϕ,t()−y)2:∈Ψ,ϕ∈ΦG^rec=\( τ,y) ( _ ψ, φ,t( τ)-y)^2: ψ∈ , φ∈ \ has Rademacher complexity ℜB(rec) R_B(G^rec) and is bounded by M2M^2, then ϵgenrec(B,δ)=4ℜB(rec)+2M2ln(2/δ)2B _gen^rec(B,δ)=4 R_B(G^rec)+2M^2 (2/δ)2B. Theorem 1 decomposes the lagged prediction error of vanilla RPMs into 1⃝ finite-sample reconstruction and approximation error from sparse optimization histories, and 2⃝ policy-update-induced temporal mismatch from applying risk estimates at t θ_t to task selection under t+1 θ_t+1. Thus, interpolation-based vanilla RPMs [22, 14] incur an (η)O(η) lagging bias, which may distort task-risk rankings and weaken task-space generalization. 3.2 Context-based Extrapolative Risk Prediction Having identified temporal misalignment and data scarcity as the main limitations, we introduce ICRPM with two components: (i) a look-ahead mechanism that conditions risk prediction on recent optimization context; and (i) temporal-weighted trajectory augmentation with exponential decay weighting to improve efficiency under tight sampling budgets. Figure 2: Left: The risk landscape evolves from t θ_t to t+1 θ_t+1, with historical prompt-loss pairs HtH_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+1t+1, and a regression head estimates the future risk ℓ^t+1 _t+1 by minimizing the discrepancy with ground-truth ℓt+1 _t+1 through gradient optimization. In-Context Look-Ahead Risk Estimation. To overcome the single-step temporal misalignment in interpolation-based evaluation, ICRPM adopts an extrapolative formulation pϕ(ℓt+1()∣,Ht)p_φ( _t+1( τ) τ,H_t), where the latest batch HtH_t serves as the inference-time optimization context. By directly targeting the future policy state t+1 θ_t+1, ICRPM predicts next-step task risk rather than reconstructing the lagged risk under t θ_t, thereby alleviating the (η)O(η) temporal bias identified in Theorem 1. Exponentially Decayed Sliding Window Augmentation. The sliding window is a data-reuse mechanism for training the look-ahead predictor under limited task evaluations. ICRPM maintains the W most recent history batches and trains on consecutive context–target pairs (t−k+1,i,ℓt−k+1(t−k+1,i),Ht−k)i=1Bk=1W \ \( τ_t-k+1,i, _t-k+1( τ_t-k+1,i),H_t-k) \_i=1^B \_k=1^W, weighted by γkγ^k. This reuse mitigates the sparse-history issue in Theorem 1: recent batches provide more training pairs, while exponential decay gives less weight to outdated histories. Optimization Objective. Combining the look-ahead formulation and historical augmentation, ICRPM minimizes a temporal-weighted negative log-likelihood over in-context prediction targets: minϕ(ϕ)=−∑k=1Wγk(1B∑i=1Blnpϕ(ℓt−k+1(t−k+1,i)|t−k+1,i,Ht−k)) _ φJ( φ)=- _k=1^Wγ^k ( 1B _i=1^B p_ φ ( _t-k+1( τ_t-k+1,i)| τ_t-k+1,i,H_t-k ) ) (4) where ϕ φ denotes model parameters, W is the look-back window and γ∈(0,1]γ∈(0,1] controls temporal decay. The objective in Eq. (4) trains ICRPM on already observed one-step transitions: at round t, it uses consecutive history pairs Hs−1→HsH_s-1\!→ H_s with s≤ts≤ t from the buffer. The unobserved batch Ht+1H_t+1 is not used when selecting t+1BT_t+1^B. Only after the selected batch is evaluated, the resulting Ht+1H_t+1 is appended to the buffer and used as a target in later RPM updates. At the inference stage, i.e., task evaluation under the coming t+1 θ_t+1, ICRPM takes in the latest batch HtH_t as the context set tC_t and the candidate identifier ∈t+1B τ∈T_t+1 B to predict risk ℓ^t+1() _t+1( τ). Algorithm 1 ICRPM Training & Prompt and Robust Predictive Sampling (PPS and RPS) Input: Task/Prompt prior p()p( τ); history buffer =HiS=\H_i\; ϕ φ; B,B^B, B; W; γ; α; mode∈PPS,RPS mode∈\ PPS, RPS\; target ρ∗ρ^*. Output: Selected batch t+1BT_t+1^B // Online ICRPM Optimization for each new batch HtH_t do // FIFO ←∪HtS ∪\H_t\; if ||>W+1|S|>W+1 then Drop oldest batch in S; (ϕ)←0J( φ)← 0; for k=1k=1 to min(W,||−1) (W,|S|-1) do // Context vs. look-ahead Hctx,Htgt=Ht−k,Ht−k+1H_ctx,H_tgt=H_t-k,H_t-k+1; ℒk←−1B∑jlnpϕ(ytgt,j∣tgt,j,Hctx)L_k←- 1B _j p_ φ(y_tgt,j τ_tgt,j,H_ctx); (ϕ)+=γkℒkJ( φ) +=γ^kL_k; // Update ϕ φ by minimizing (ϕ)J( φ) ϕ φ = ϕ φ - α∇ϕ(ϕ)α _ φJ( φ); // Predictive Task/Prompt Selection Sample candidate pool: t+1B^=^t+1,ii=1B^∼p()T_t+1 B=\ τ_t+1,i\_i=1 B p( τ); for i=1i=1 to B B do // One-step look-ahead y^it+1=Epϕ(y∣^t+1,i,Ht)[y] y_i^\,t+1=E_p_ φ(y τ_t+1,i,H_t)[y]; if mode=PPS mode= PPS then ct+1∗←ρ∗c_t+1^*←ρ^* ; // y y is success rate else ct+1∗←maxy^t+1c_t+1^*← _ τ y_ τ^\,t+1 ; // y y is risk t+1B=argtop−B∈t+1B^(−‖y^t+1−ct+1∗‖22)T_t+1^B= *arg\,top-B_ τ _t+1 B (- \| y_ τ^\,t+1-c_t+1^* \|_2^2 ); return t+1BT_t+1^B First-order interpretation. The look-ahead objective in Eq. (4) aligns RPM supervision with the policy state used for task selection. The following proposition gives a lightweight formalization: if recent optimization history enables the in-context predictor to capture the local risk drift, then the next-step prediction error is controlled by a context approximation term and a second-order smoothness residual. Proposition 2 (Look-ahead reduces temporal mismatch). Let t=t+1−t _t= θ_t+1- θ_t with ‖t‖≤ηG\| _t\|≤η G, and assume that ℓ(,) ( θ, τ) is β-smooth in θ. For analysis, define the first-order look-ahead risk as mt()=ℓt()+∇ℓ(t,)⊤tm_t( τ)= _t( τ)+ _ θ ( θ_t, τ) _t. If the ICRPM predictor satisfies ∼p|ℓ^t+1IC()−mt()|≤ϵICE_ τ p| ^IC_t+1( τ)-m_t( τ)|≤ _IC, then ∼p[|ℓ^t+1IC()−ℓt+1()|]≤ϵIC+β2η2G2.E_ τ p [| ^IC_t+1( τ)- _t+1( τ)| ]≤ _IC+ β2η^2G^2. (5) This result provides a first-order explanation for the advantage of look-ahead risk prediction. A lagged RPM estimates ℓt _t for task selection under t+1 θ_t+1 and therefore retains the O(η)O(η) temporal mismatch identified in Theorem 1. In contrast, ICRPM directly targets the next-step risk: when the context predictor captures the local risk drift, the leading lagging bias is replaced by the in-context approximation error ϵIC _IC and a second-order residual. 3.3 Implementation and Neural Architecture Algorithm 1 summarizes the implementation; the selected batch is passed to downstream optimizers such as robust Meta-RL, DR, or LLM RLVR. ICRPM performs the extrapolative formulation p(ℓ|τ,Ht;t+1)p( |τ,H_t; θ_t+1). At each iteration, the context set t−1=(i,ℓt−1(i))i=1BC_t-1=\( τ_i, _t-1( τ_i))\_i=1^B consists of task-risk pairs evaluated under t−1 θ_t-1, while each training target is a pair (,ℓt())( τ, _t( τ)) under the last policy state t θ_t. This yields B supervised training episodes per iteration, one for each j∈[B]j∈[B], without any additional environment interaction. Risk Sequence Tokenization and Cross-Attention. ICRPM has a simple design. As illustrated in Fig. 2, we tokenize the task-risk pairs and adopt the Transformer as the default neural architecture to capture the correlation between the target and past risk sequences. This way encodes the context into a permutation-invariant representation and produces p(ℓ|τ,Ht;t+1)p( |τ,H_t; θ_t+1) via cross-attention. 4 Results and Analysis This section empirically examines ICRPM’s influence within MPTS. Motivated by the failure modes diagnosed in Section 3.1, the experiments are designed to answer two Research Questions (RQs): 1. Do extrapolative RPMs enhance task evaluation reliability over iteration? 2. How does task evaluation reliability influence performance in downstream applications? We track Spearman Rank Correlation (SRC) between predicted task difficulty and ground-truth task risk as a diagnostic measure of rank reliability, and report CVaRα [29] for robust optimization scenarios and test average performance for RLVR. Full experimental details are provided in Appendix E. Benchmarks and Baselines. We evaluate ICRPM on few-shot sinusoid regression, Meta-RL, and DR, and LLM post-training via RLVR. (i) For regression and continuous-control, we retain the standard MPTS protocol [22], e.g., the same task distribution, batch size, candidate pool, selection budget, and downstream optimizer; and baselines include ERM [30], GDRM [31], DRM [17, 32], OHTM [33], MPTS [22]. (i) For LLM RL post-training, we adopt the curriculum-based setup of [14], and baselines are uniform sampling and SoTA predictive sampling methods such as PCL [14], GRESO [34], and MoPPS [13]. For both cases, ICRPM replaces only the interpolation-based RPM with a parameter-matched context-conditioned one-step extrapolative predictor. 4.1 Robust Few-Shot Sinusoid Regression We first validate ICRPM on the classic few-shot sinusoid regression benchmark [9], which provides a controlled setting to isolate the effect of improved task risk estimation. Tasks are parameterised by amplitude and phase, and MAML is used as the meta-learning backbone. Figure 3: Few-shot sinusoid regression results (10 Runs, Mean ± Standard Error of the Mean (SEM)). From left to right: (a) validation CVaR0.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. Few-Shot Regression Analysis. In Fig. 3, ICRPM consistently outperforms all baselines. It achieves lower average MSE (Fig. 3b) and CVaR0.9 MSE (Fig. 3a) with faster convergence, while maintaining higher SRC with ground-truth task risk throughout training (Fig. 3c). At meta-test time (Fig. 3d), ICRPM achieves the lowest CVaRα MSE across all α∈0.5,0.7,0.9α∈\0.5,0.7,0.9\, suggesting that more reliable task prioritisation is accompanied by stronger tail-task generalization. Vanilla MPTS underperforms in both convergence stability and worst-case performance. 4.2 Robust Meta-RL and Domain Randomization We evaluate ICRPM on Meta-RL over three MuJoCo environments [9, 35] and zero-shot DR on Ergo-Reacher and Lunar-Lander [36]. Both settings require generalisation across a distribution of MDPs parameterised by continuous task identifiers, where robustness to tail or out-of-distribution tasks is critical. 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: CVaR0.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. Table 1: Evaluation return on five robotic control tasks. (10 Runs, Mean ± SEM) The best results are bold, with the runner-up underlined. Method HalfCheetahVel Walker2dVel ReacherPos LunarLander ErgoReacher CVaR0.5 CVaR0.9 CVaR0.5 CVaR0.9 CVaR0.5 CVaR0.9 CVaR0.5 CVaR0.9 CVaR0.5 CVaR0.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 ICRPM (Ours) -140.5± 3.5 -180.2± 6.4 43.2± 6.7 -1.5± 7.2 -15.2± 0.9 -27.8± 3.8 285.3± 2.6 265.8± 2.1 86.5± 0.4 78.3± 0.7 Meta-RL Analysis. Fig. 4 (columns 3–5) reports training dynamics on HalfCheetahVel, Walker2dVel, and ReacherPos. ICRPM achieves the highest CVaR0.9 validation returns on all three benchmarks while preserving competitive average performance — a balance DRM fails to maintain, notably degrading on Walker2dVel. The SRC curves (row 3) confirm that the Transformer-based RPM with temporal-weighted augmentation sustains markedly higher rank correlation than MPTS throughout training. At meta-test time (Table 1), ICRPM attains the best CVaRα returns across nearly all environments and robustness levels, with especially large gains on Walker2dVel (CVaR0.9: −1.5-1.5 vs. −27.9-27.9 for MPTS) and substantially reduced variance. Figure 5: CVaR0.9 return vs. relative runtime on Walker2dVel. Domain Randomization Analysis. On Ergo-Reacher and Lunar-Lander (Fig. 4, columns 1–2), ICRPM consistently outperforms all baselines in CVaR0.9 returns while remaining competitive in average performance. Notably, on Lunar-Lander ICRPM improves CVaR0.9 at convergence by over 40% relative to ERM, effectively handling the long-tail of adversarial engine-strength configurations. The SRC advantage of ICRPM over MPTS persists in both DR environments, demonstrating that mitigating temporal lagging bias extends beyond episodic few-shot learning to the zero-shot continuous-control regime. Computational Overhead. Fig. 5 plots CVaR0.9 return against relative runtime on Walker2dVel (normalised to ERM). ICRPM incurs only 8% overhead (1.08×), comparable to MPTS (1.07×) and far below DRM (1.34×), whose per-iteration importance-weight computation dominates cost. GDRM and OHTM add minimal time (1.03× and 1.01×) but lag considerably in tail-risk performance. The scatter clearly shows that ICRPM occupies the ideal upper-left region with the highest CVaR0.9 return at near-baseline computational cost, confirming that its performance gains come at negligible additional runtime budget. 4.3 Sample-Efficient RLVR on Countdown To evaluate the effectiveness of ICRPM in accelerating LLM post-training, we conduct experiments on the Countdown symbolic reasoning task. Countdown requires composing given numbers with arithmetic operations to reach a target value, offering verifiable binary rewards that enable straightforward RLVR training. Experimental Setup. We use Qwen3-4B-Base [37] as the policy backbone. Following the RLVR paradigm, we employ Group Relative Policy Optimization (GRPO) [38] to optimize the model with verifiable rewards that consider both mathematical correctness and format adherence. Training is performed on a 20k-instance subset of Countdown-34 (CD34); evaluation is conducted on both CD34 and the harder Countdown-4 (CD4) variant. Figure 6: Test accuracy (left) and SRC (right) on Countdown 4B. ICRPM demonstrates significantly faster convergence and higher correlation with ground-truth difficulty throughout RL post-training. Method CD4 CD34 Avg. Qwen3-4B [37] 1.3 3.5 2.4 +Uniform 51.1 73.8 62.5 +PCL [14] 51.0 72.8 61.9 +GRESO [34] 53.8 73.8 63.8 +MoPPS [13] 52.9 73.9 63.4 +ICRPM 55.7 75.1 65.4 Table 2: Evaluation results on Countdown. ‘+’ denotes finetuning with the corresponding method. Bold: best; underline: runner-up. Curriculum RLVR Analysis. As summarised in Table 2 and Fig. 6, ICRPM consistently outperforms all baselines. On CD4, ICRPM achieves 55.7% accuracy, surpassing GRESO by 1.9 points and MoPPS by 2.8 points, with similar gains on CD34. The training curve (Fig. 6, left) shows faster convergence and superior sample efficiency, while the SRC curve (right) indicates more reliable online prompt-difficulty rankings as the policy evolves. This addresses the sparse-history and temporal-lag issues discussed in Section 3.1: unlike prompt-specific models such as MoPPS, ICRPM shares information across prompts and better tracks non-stationary difficulty during RL post-training. 4.4 Ablation Studies 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. We ablate the subset acquisition principles, RPM backbones, window size W, and decay factor γ on Lunar-Lander. Fig. 7 shows that ICRPM improves over PDTS under both worst-case and diversity-aware acquisition. For architecture and history augmentation, ICRPM with a parameter-matched MLP still outperforms vanilla MPTS, while MPTS augmented with the same three-step window and γ=0.9γ=0.9 remains below ICRPM and ICRPM (MLP). This indicates that the gains mainly come from context-conditioned look-ahead risk prediction, rather than a specific acquisition rule, Transformer backbone, or history-reuse heuristic. Together, these ablations match the decomposition in Theorem 1: look-ahead prediction addresses temporal lag, while windowed reuse mitigates sparse history. Fig. 8 shows that W has a stronger effect than γ: we choose W=3W=3 because it gives the best CVaR0.9 performance in the ablation, and use γ=0.9γ=0.9 across all benchmarks as a robust default that mildly discounts older histories. Figure 8: Sensitivity to W and γ on Lunar-Lander (10 Runs). Left: Sensitivity to look-back window W. CVaR0.9 returns, average returns under W∈1,2,3,8W∈\1,2,3,8\. Right: Sensitivity to decay factor γ. CVaR0.9 returns, average returns under γ∈0.5,0.7,0.9,1.0γ∈\0.5,0.7,0.9,1.0\. 5 Conclusion This work identifies a structural limitation of MPTS: vanilla RPMs estimate risks under the current policy but use them for post-update task selection, inducing first-order lagging bias and unstable task rankings. We propose ICRPM, which reframes risk prediction as context-conditioned one-step extrapolation with temporal-weighted history augmentation. Empirically, ICRPM improves risk ranking, sample efficiency, and tail-risk performance across regression, robust Meta-RL, domain randomization, and RLVR-based LLM post-training, suggesting that task samplers should model policy-induced risk dynamics rather than rely on static or lagged difficulty estimates. References [1] R. S. Sutton and A. G. Barto (2018) Reinforcement learning: an introduction. MIT press. Cited by: §1, §2. [2] V. Mnih, K. Kavukcuoglu, D. Silver, A. A. Rusu, J. Veness, M. G. Bellemare, A. Graves, M. Riedmiller, A. K. Fidjeland, G. Ostrovski, et al. (2015) Human-level control through deep reinforcement learning. nature 518 (7540), p. 529–533. Cited by: §1. [3] Y. Qu, B. Wang, J. Shao, et al. (2023) Hokoff: real game dataset from honor of kings and its offline reinforcement learning benchmarks. In NeurIPS, Cited by: §1. [4] J. Shao, Y. Qu, C. Chen, H. Zhang, and X. Ji (2024) Counterfactual conservative q learning for offline multi-agent reinforcement learning. NeurIPS 36. Cited by: §1. [5] Q. Wang and H. Van Hoof (2022) Model-based meta reinforcement learning using graph structured surrogate models and amortized policy search. In International Conference on Machine Learning, p. 23055–23077. Cited by: §1. [6] S. Carpin, Y. Chow, and M. Pavone (2016) Risk aversion in finite markov decision processes using total cost criteria and average value at risk. In 2016 ieee international conference on robotics and automation (icra), p. 335–342. Cited by: §1. [7] D. Rempe, J. Philion, L. J. Guibas, S. Fidler, and O. Litany (2022) Generating useful accident-prone driving scenarios via a learned traffic prior. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, p. 17305–17315. Cited by: §1. [8] F. Muratore, F. Treede, M. Gienger, and J. Peters (2018) Domain randomization for simulation-based policy optimization with transferability assessment. In Conference on Robot Learning, p. 700–713. Cited by: Appendix A, Appendix A, §1. [9] C. Finn, P. Abbeel, and S. Levine (2017) Model-agnostic meta-learning for fast adaptation of deep networks. In International conference on machine learning, p. 1126–1135. Cited by: Appendix A, Appendix A, §C.2, §1, §4.1, §4.2. [10] L. Ouyang, J. Wu, X. Jiang, D. Almeida, C. Wainwright, P. Mishkin, C. Zhang, S. Agarwal, K. Slama, A. Ray, et al. (2022) Training language models to follow instructions with human feedback. Advances in neural information processing systems 35, p. 27730–27744. Cited by: Appendix A, §1. [11] J. Achiam, S. Adler, S. Agarwal, L. Ahmad, I. Akkaya, F. L. Aleman, D. Almeida, J. Altenschmidt, S. Altman, S. Anadkat, et al. (2023) Gpt-4 technical report. arXiv preprint arXiv:2303.08774. Cited by: §1. [12] Q. Yu, Z. Zhang, R. Zhu, Y. Yuan, X. Zuo, Y. Yue, W. Dai, T. Fan, G. Liu, L. Liu, et al. (2025) Dapo: an open-source llm reinforcement learning system at scale. arXiv preprint arXiv:2503.14476. Cited by: §1, §1. [13] Y. Qu, Q. Wang, Y. Mao, V. T. Hu, B. Ommer, and X. Ji (2026) Can prompt difficulty be online predicted for accelerating rl finetuning of reasoning models?. In Proceedings of the 32nd ACM SIGKDD Conference on Knowledge Discovery and Data Mining V. 1, p. 1240–1250. Cited by: §E.6, §1, §1, §1, §2, §2, §2, §4, Table 2. [14] Z. Gao, J. Kim, W. Sun, T. Joachims, S. Wang, R. Y. Pang, and L. Tan (2026) Prompt curriculum learning for efficient LLM post-training. In The Fourteenth International Conference on Learning Representations, External Links: Link Cited by: Appendix A, §E.6, §1, §1, §2, §2, §2, §2, §3.1, §4, Table 2. [15] Y. Qu, Q. Wang, Y. Mao, H. Zou, Y. Jiang, W. Liu, C. Bai, K. Yang, Y. Chen, S. Yang, et al. (2026) Small generalizable prompt predictive models can steer efficient rl post-training of large reasoning models. arXiv preprint arXiv:2602.01970. Cited by: §1. [16] T. Haarnoja, A. Zhou, P. Abbeel, and S. Levine (2018) Soft actor-critic: off-policy maximum entropy deep reinforcement learning with a stochastic actor. In International conference on machine learning, p. 1861–1870. Cited by: §1. [17] Q. Wang, Y. Lv, Z. Xie, J. Huang, et al. (2024) A simple yet effective strategy to robustify the meta learning paradigm. Advances in Neural Information Processing Systems 36. Cited by: Appendix A, §1, §4, Table 1. [18] Y. Lv, C. Wang, D. Liang, and Z. Xie (2024) Theoretical investigations and practical enhancements on tail task risk minimization in meta learning. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, External Links: Link Cited by: Appendix A, §1. [19] Q. Wang, Y. Lv, Y. Mao, Y. Qu, Y. Xu, and X. Ji (2025) Robust fast adaptation from adversarially explicit task distribution generation. In Proceedings Of The 31st ACM SIGKDD Conference On Knowledge Discovery And Data Mining V. 1, p. 1481–1491. Cited by: §1. [20] J. Wei, X. Wang, D. Schuurmans, M. Bosma, B. Ichter, F. Xia, E. Chi, Q. Le, and D. Zhou (2023) Chain-of-thought prompting elicits reasoning in large language models. External Links: 2201.11903, Link Cited by: §1. [21] Y. Qu, Q. Wang, Y. Mao, H. Zou, Y. Jiang, Y. Li, W. Xu, L. Cai, W. Liu, C. Bai, et al. (2026) Listwise policy optimization: group-based rlvr as target-projection on the llm response simplex. arXiv preprint arXiv:2605.06139. Cited by: §1. [22] Q. C. Wang, Z. Xiao, Y. Mao, Y. Qu, J. Shen, Y. Lv, and X. Ji (2025) Model predictive task sampling for efficient and robust adaptation. External Links: 2501.11039, Link Cited by: Appendix A, item 2., §E.1, §E.6, §1, §2, §2, §2, §2, §3.1, §3.1, §3.1, §4, Table 1. [23] Y. Qu, Q. C. Wang, Y. Mao, Y. Lv, and X. Ji (2025) Fast and robust: task sampling with posterior and diversity synergies for adaptive decision-makers in randomized environments. arXiv preprint arXiv:2504.19139. Cited by: Appendix A, item 1., §1, §2, §2, §2, §3.1, Figure 7, Figure 7. [24] Y. Mao, Y. Qu, Q. Wang, H. Zou, and X. Ji (2026) Dynamics-predictive sampling for active rl finetuning of large reasoning models. In International Conference on Learning Representations, Vol. 2026, p. 25242–25270. Cited by: §1. [25] S. Bae, J. Hong, M. Y. Lee, H. Kim, J. Nam, and D. Kwak (2026) Online difficulty filtering for reasoning oriented reinforcement learning. External Links: 2504.03380, Link Cited by: §2. [26] D. P. Kingma and M. Welling (2013) Auto-encoding variational bayes. arXiv preprint arXiv:1312.6114. Cited by: §2. [27] A. Ganguly and S. W. F. Earp (2021) An introduction to variational inference. External Links: 2108.13083, Link Cited by: §2. [28] D. J. Rezende, S. Mohamed, and D. Wierstra (2014) Stochastic backpropagation and approximate inference in deep generative models. In International conference on machine learning, p. 1278–1286. Cited by: §2. [29] R. T. Rockafellar S. Uryasev et al. (2000) Optimization of conditional value-at-risk. Journal of risk 2, p. 21–42. Cited by: Appendix A, §4, Definition 4. [30] V. N. Vapnik (1998) Statistical learning theory. Cited by: Appendix A, §E.6, §4, Table 1. [31] S. Sagawa, P. W. Koh, T. B. Hashimoto, and P. Liang (2019) Distributionally robust neural networks. In International Conference on Learning Representations, Cited by: Appendix A, §E.6, §4, Table 1. [32] I. Greenberg, S. Mannor, G. Chechik, and E. Meirom (2024) Train hard, fight easy: robust meta reinforcement learning. Advances in Neural Information Processing Systems 36. Cited by: Appendix A, §4, Table 1. [33] R. Kumar, T. Deleu, and Y. Bengio (2023) The effect of diversity in meta-learning. In Proceedings of the AAAI Conference on Artificial Intelligence, Vol. 37, p. 8396–8404. Cited by: §E.6, §4, Table 1. [34] H. Zheng, Y. Zhou, B. R. Bartoldson, B. Kailkhura, F. Lai, J. Zhao, and B. Chen (2025) Act only when it pays: efficient reinforcement learning for llm reasoning via selective rollouts. arXiv preprint arXiv:2506.02177. Cited by: §E.6, §4, Table 2. [35] E. Todorov, T. Erez, and Y. Tassa (2012) Mujoco: a physics engine for model-based control. In 2012 IEEE/RSJ international conference on intelligent robots and systems, p. 5026–5033. Cited by: §4.2. [36] B. Mehta, M. Diaz, F. Golemo, C. J. Pal, and L. Paull (2020) Active domain randomization. In Conference on Robot Learning, p. 1162–1176. Cited by: Appendix A, §E.3, §4.2. [37] A. Yang, A. Li, B. Yang, et al. (2025) Qwen3 technical report. External Links: 2505.09388, Link Cited by: §4.3, Table 2. [38] Z. Shao, P. Wang, Q. Zhu, R. Xu, J. Song, X. Bi, H. Zhang, M. Zhang, Y. Li, Y. Wu, et al. (2024) Deepseekmath: pushing the limits of mathematical reasoning in open language models. arXiv preprint arXiv:2402.03300. Cited by: Appendix A, §C.3, §4.3. [39] Y. Chow, A. Tamar, S. Mannor, and M. Pavone (2015) Risk-sensitive and robust decision-making: a cvar optimization approach. Advances in neural information processing systems 28. Cited by: Appendix A. [40] Y. Chow (2017) Risk-sensitive and data-driven sequential decision making. Ph.D. Thesis, Stanford University. Cited by: Appendix A. [41] Y. Chow, M. Ghavamzadeh, L. Janson, and M. Pavone (2018) Risk-constrained reinforcement learning with percentile risk criteria. Journal of Machine Learning Research 18 (167), p. 1–51. Cited by: Appendix A. [42] Q. Tao, J. Wen, J. Yang, G. Wu, Z. Kaiyu, Y. Lv, W. Du, X. Liang, and Q. Wang (2026) Latent space robust optimization of neural processes with aligned stratified order-statistic loss reduction. In Forty-third International Conference on Machine Learning, Cited by: Appendix A. [43] H. Li, H. Chi, M. Liu, W. Huang, L. Xu, and W. Yang (2025) Transformer-based spatial-temporal counterfactual outcomes estimation. In International Conference on Machine Learning, Vol. 267, p. 34185–34208. Cited by: Appendix A. [44] A. Tamar, Y. Chow, M. Ghavamzadeh, and S. Mannor (2015) Policy gradient for coherent risk measures. Advances in neural information processing systems 28. Cited by: Appendix A. [45] X. Pan, D. Seita, Y. Gao, and J. Canny (2019) Risk averse robust adversarial reinforcement learning. In 2019 International Conference on Robotics and Automation (ICRA), p. 8522–8528. Cited by: Appendix A. [46] M. Rigter, B. Lacerda, and N. Hawes (2021) Risk-averse bayes-adaptive reinforcement learning. Advances in Neural Information Processing Systems 34, p. 1142–1154. Cited by: Appendix A. [47] I. Greenberg, Y. Chow, M. Ghavamzadeh, and S. Mannor (2022) Efficient risk-averse reinforcement learning. Advances in Neural Information Processing Systems 35, p. 32639–32652. Cited by: Appendix A. [48] J. Tobin, R. Fong, A. Ray, J. Schneider, W. Zaremba, and P. Abbeel (2017) Domain randomization for transferring deep neural networks from simulation to the real world. In 2017 IEEE/RSJ international conference on intelligent robots and systems (IROS), p. 23–30. Cited by: Appendix A, Appendix A. [49] J. Beck, R. Vuorio, E. Z. Liu, Z. Xiong, L. Zintgraf, C. Finn, and S. Whiteson (2023) A survey of meta-reinforcement learning. arXiv preprint arXiv:2301.08028. Cited by: Appendix A. [50] Q. Wang and H. Van Hoof (2022) Learning expressive meta-representations with mixture of expert neural processes. Advances in neural information processing systems 35, p. 26242–26255. Cited by: Appendix A. [51] L. Collins, A. Mokhtari, and S. Shakkottai (2020) Task-robust model-agnostic meta-learning. Advances in Neural Information Processing Systems 33, p. 18860–18871. Cited by: Appendix A. [52] F. Muratore, C. Eilers, M. Gienger, and J. Peters (2021) Data-efficient domain randomization with bayesian optimization. IEEE Robotics and Automation Letters 6 (2), p. 911–918. Cited by: Appendix A. [53] G. Tiboni, P. Klink, J. Peters, T. Tommasi, C. D’Eramo, and G. Chalvatzaki (2023) Domain randomization via entropy maximization. arXiv preprint arXiv:2311.01885. Cited by: Appendix A. [54] W. Yang, H. Chi, Y. Zhan, B. Hu, X. Ren, D. Tao, and L. Lan (2025) NT-fan: a simple yet effective noise-tolerant few-shot adaptation network. Artificial Intelligence 346, p. 104363. Cited by: Appendix A. [55] H. Chi, F. Liu, W. Yang, L. Lan, T. Liu, B. Han, G. Niu, M. Zhou, and M. Sugiyama (2022) Meta discovery: learning to discover novel classes given very limited data. In International Conference on Learning Representations, Cited by: Appendix A. [56] H. Chi, W. Yang, F. Liu, L. Lan, T. Qin, and B. Han (2024) Does confusion really hurt novel class discovery?. International Journal of Computer Vision 132 (8), p. 3191–3207. Cited by: Appendix A. [57] H. Chi, H. Li, W. Yang, F. Liu, L. Lan, X. Ren, T. Liu, and B. Han (2024) Unveiling causal reasoning in large language models: reality or mirage?. Advances in Neural Information Processing Systems 37, p. 96640–96670. Cited by: Appendix A. [58] T. Hospedales, A. Antoniou, P. Micaelli, and A. Storkey (2021) Meta-learning in neural networks: a survey. IEEE transactions on pattern analysis and machine intelligence 44 (9), p. 5149–5169. Cited by: Appendix A. [59] C. Finn, K. Xu, and S. Levine (2018) Probabilistic model-agnostic meta-learning. Advances in neural information processing systems 31. Cited by: Appendix A. [60] M. Abbas, Q. Xiao, L. Chen, P. Chen, and T. Chen (2022) Sharp-maml: sharpness-aware model-agnostic meta learning. In International conference on machine learning, p. 10–32. Cited by: Appendix A. [61] A. Rajeswaran, C. Finn, S. M. Kakade, and S. Levine (2019) Meta-learning with implicit gradients. Advances in neural information processing systems 32. Cited by: Appendix A. [62] K. Rakelly, A. Zhou, C. Finn, S. Levine, and D. Quillen (2019) Efficient off-policy meta-reinforcement learning via probabilistic context variables. In International conference on machine learning, p. 5331–5340. Cited by: Appendix A. [63] L. Zintgraf, K. Shiarlis, M. Igl, S. Schulze, Y. Gal, K. Hofmann, and S. Whiteson (2019) Varibad: a very good method for bayes-adaptive deep rl via meta-learning. arXiv preprint arXiv:1910.08348. Cited by: Appendix A. [64] Q. Wang, M. Federici, and H. van Hoof (2022) Bridge the inference gaps of neural processes via expectation maximization. In The Eleventh International Conference on Learning Representations, Cited by: Appendix A. [65] F. Xu, S. Jiang, H. Yin, Z. Zhang, Y. Yu, M. Li, D. Li, and W. Liu (2021) Enhancing context-based meta-reinforcement learning algorithms via an efficient task encoder (student abstract). In Proceedings of the AAAI Conference on Artificial Intelligence, Vol. 35, p. 15937–15938. Cited by: Appendix A. [66] S. Ritter, J. Wang, Z. Kurth-Nelson, S. Jayakumar, C. Blundell, R. Pascanu, and M. Botvinick (2018) Been there, done that: meta-learning with episodic recall. In International conference on machine learning, p. 4354–4363. Cited by: Appendix A. [67] M. Dennis, N. Jaques, E. Vinitsky, A. Bayen, S. Russell, A. Critch, and S. Levine (2020) Emergent complexity and zero-shot transfer via unsupervised environment design. Advances in neural information processing systems 33, p. 13049–13061. Cited by: Appendix A. [68] M. Jiang, M. Dennis, J. Parker-Holder, J. Foerster, E. Grefenstette, and T. Rocktäschel (2021) Replay-guided adversarial environment design. Advances in Neural Information Processing Systems 34, p. 1884–1897. Cited by: Appendix A. [69] C. Koprulu, T. D. Simão, N. Jansen, and U. Topcu (2023) Risk-aware curriculum generation for heavy-tailed task distributions. In Uncertainty in Artificial Intelligence, p. 1132–1142. Cited by: Appendix A. [70] C. Wang, Y. Lv, Y. Mao, Y. Qu, Y. Xu, and X. Ji (2024) Robust fast adaptation from adversarially explicit task distribution generation. arXiv preprint arXiv:2407.19523. Cited by: Appendix A. [71] Y. Bai, A. Jones, K. Ndousse, A. Askell, A. Chen, N. DasSarma, D. Drain, S. Fort, D. Ganguli, T. Henighan, et al. (2022) Training a helpful and harmless assistant with reinforcement learning from human feedback. arXiv preprint arXiv:2204.05862. Cited by: Appendix A. [72] Y. Mao, Y. Qu, Q. Wang, H. Zou, and X. Ji (2026) RLVR without ineffective samples: group prioritized off-policy optimization for llm reasoning. arXiv preprint arXiv:2606.01281. Cited by: Appendix A. [73] Y. Qu, Q. Wang, Y. Mao, H. Zou, Y. Jiang, Y. Li, W. Xu, L. Cai, W. Liu, C. Bai, et al. (2026) Listwise policy optimization: group-based rlvr as target-projection on the llm response simplex. arXiv preprint arXiv:2605.06139. Cited by: Appendix A. [74] Y. Mao, Y. Qu, Q. Wang, H. Zou, and X. Ji (2026) RLVR without ineffective samples: group prioritized off-policy optimization for llm reasoning. arXiv preprint arXiv:2606.01281. Cited by: Appendix A. [75] H. Zou, Q. Wang, Y. Qu, Y. Jiang, L. Cai, Y. Mao, R. Peng, X. Xu, W. Liu, K. Yang, et al. (2026) TRACE: a unified rollout budget allocation framework for efficient agentic reinforcement learning. arXiv preprint arXiv:2606.11119. Cited by: Appendix A. [76] S. Fujimoto, H. Hoof, and D. Meger (2018) Addressing function approximation error in actor-critic methods. In International conference on machine learning, p. 1587–1596. Cited by: §C.2. Appendix A Related Work Risk-sensitive optimization for adaptive sequential decision-making. Risk considerations have long been central to sequential decision-making, where poor adaptation may lead to severe downstream failures [39, 40, 41, 42, 43]. This has motivated a broad line of work on incorporating risk-averse objectives into policy optimization [44, 45, 46, 47]. In randomized environments, including DR [48, 8] and Meta-RL [9, 49, 50], several optimization principles have been adopted to improve robustness. Empirical risk minimization (ERM) [30] optimizes the average loss under the training task distribution and is therefore commonly used as a standard baseline. Distributionally robust risk minimization (DRM) instead emphasizes high-risk tasks via conditional value-at-risk (CVaR) objectives [29, 17, 18, 32], while group distributionally robust optimization (GDRM) mitigates subpopulation shifts by reweighting group-level risks [31]. Worst-case meta-learning has also been investigated [51], but typically requires careful relaxations to maintain stable training. A shared limitation of these approaches is their reliance on extensive policy evaluation over large task spaces, which makes purely evaluation-driven robust optimization expensive in practice. Robust active task sampling and risk-predictive models. Robust active task sampling (RATS) addresses the above efficiency bottleneck by replacing repeated exhaustive evaluation with a learned RPM [22]. Instead of executing the current policy on every candidate task, RATS estimates task risk through the RPM and then selects an informative or high-risk subset for policy optimization. MPTS [22] instantiates this idea by modeling the optimization trajectory as a generative sequence and using streaming variational inference to predict adaptation risk on unseen tasks. More recently, PDTS [23] reinterprets RATS from a task-selection perspective, formulates subset selection as an infinitely many-armed bandit problem, and combines posterior sampling with diversity regularization to alleviate concentration under large candidate batches. These methods substantially improve the efficiency of robust task selection; however, their predictive modules remain challenged by two factors: early-stage histories provide only sparse supervision, and risk estimates are inherently delayed because predictions are made before the next policy update is observed. Our method focuses on these RPM-level limitations by introducing in-context sequence modeling and a one-step look-ahead prediction mechanism. Cross-task adaptation and curriculum learning. Zero-shot and few-shot adaptation are two central evaluation regimes for adaptive decision-makers. DR aims to obtain zero-shot generalization by training over randomized environments [48, 8, 36, 52, 53, 54, 55, 56, 57], whereas meta-learning seeks reusable priors that enable rapid adaptation from limited experience [58]. Representative Meta-RL methods include optimization-based algorithms that learn initializations [9, 59, 60, 61], context-based approaches that infer latent task representations from support trajectories [62, 63, 64, 65], and recurrent architectures that internalize adaptation through memory [66]. Curriculum learning provides a complementary mechanism by changing which tasks are presented during training. For example, unsupervised environment design and its variants generate tasks with increasing challenge [67, 68, 69], while adversarial task-distribution methods construct curricula through learned generative distributions [70]. Our work is orthogonal to such curriculum-design strategies: rather than modifying the task distribution itself, we improve the accuracy and timeliness of risk estimation under a given distribution, making the approach potentially compatible with existing curricula. Reinforcement learning for LLM post-training. Reinforcement learning from human feedback (RLHF) has become a standard approach for aligning large language models with human preferences [10, 71]. Recent post-training pipelines further use reinforcement learning to improve reasoning and task-solving capabilities [72, 73], where prompts can be viewed as task instances and responses are scored by reward models or verifiers [38]. Prompt-level curricula that emphasize informative or challenging examples have also been explored to improve training efficiency [14, 74, 75]. A practical obstacle is that estimating prompt difficulty usually requires costly rollouts from the current model. Our framework can be naturally mapped to this setting by treating prompts as task identifiers and using the RPM to estimate prompt-level difficulty, thereby supporting curriculum construction for LLM post-training without additional generation overhead. Appendix B Expanded Proofs for Theoretical Analysis This appendix contains the full proofs of all theoretical results stated in the main text. We first collect notation, then prove Theorem 1 and Proposition 2. Notation. T: task space; Θ : parameter space; t∈Θ θ_t∈ : agent parameters at step t; L(,)L( θ, τ): true risk of task ∈ τ under θ; Ht=(t,i,ℓt,i)i=1BH_t=\( τ_t,i, _t,i)\_i=1^B: per-iteration history with batch size B; H0:t=⋃k=0tHkH_0:t= _k=0^tH_k: cumulative trajectory; ℱF: RPM hypothesis class; ℜn(ℱ) R_n(F): Rademacher complexity with n samples. The vanilla RPM ℓ^vanilla _vanilla is trained on HtH_t (size B). ICRPM ℓ^ICRPM=fϕ(,H0:t) _ICRPM=f_φ( τ,H_0:t) is trained on W consecutive batches (size N=WBN=WB). Theorem 1 (Per-step Lagged Prediction Bound for Vanilla RPMs) Assume that ℓ()=ℒ(,)∈[0,M] ( τ)=L( θ, τ)∈[0,M] for all ∈Θ θ∈ , ∈ τ , and that ℓ() ( τ) is L-Lipschitz in θ. The policy update uses a normalized average loss computed by i.i.d. samples and ‖∇ℓt(i)‖≤G\| _ θ _t( τ_i)\|≤ G. Let ℓ^,ϕ,t() _ ψ, φ,t( τ) denote the predictive risk function induced by the vanilla RPM at iteration t. Define ϵapprec:=inf,ϕ∼p[(ℓ^,ϕ,t()−ℓt())2]. _app^rec:= _ ψ, φE_ τ p [ ( _ ψ, φ,t( τ)- _t( τ) )^2 ]. Then, for any δ∈(0,1)δ∈(0,1), with probability at least 1−δ1-δ over the sampling of HtH_t, ∼p[|ℓ^t()−ℓt+1()|]≤ϵapprec+ϵgenrec(B,δ)⏟ 1⃝Reconstruction Approximation Error+LηG⏟ 2⃝Temporal Lagging Error.E_ τ p [| _t( τ)- _t+1( τ)| ]≤ _app^rec+ _gen^rec(B,δ)_ 1Reconstruction Approximation Error + Lη G_ 2Temporal Lagging Error. (6) If the reconstruction loss class rec=(,y)↦(ℓ^,ϕ,t()−y)2:∈Ψ,ϕ∈ΦG^rec=\( τ,y) ( _ ψ, φ,t( τ)-y)^2: ψ∈ , φ∈ \ has Rademacher complexity ℜB(rec) R_B(G^rec) and is bounded by M2M^2, then ϵgenrec(B,δ)=4ℜB(rec)+2M2ln(2/δ)2B. _gen^rec(B,δ)=4 R_B(G^rec)+2M^2 (2/δ)2B. Proof. Assume the RPM uses a fixed-variance Gaussian decoder p(ℓ|,)=(μ(,),σ2)p_ ψ( | τ, z)=N( _ ψ( τ, z),σ^2), and define the predictive mean as ℓ^,ϕ,t():=qϕ(t|Ht)[μ(,t)] _ ψ, φ,t( τ):=E_q_ φ( z_t|H_t)[ _ ψ( τ, z_t)]. For such a decoder, the reconstruction term in the negative log-likelihood is proportional to squared reconstruction error: −logp(ℓt()|,t)=12σ2(μ(,t)−ℓt())2+Cσ,- p_ ψ( _t( τ)| τ, z_t)= 12σ^2 ( _ ψ( τ, z_t)- _t( τ) )^2+C_σ, (7) where CσC_σ is independent of (,ϕ)( ψ, φ). Therefore, the reconstruction component of the ELBO in Eq. (2) induces the squared reconstruction objective analyzed below, while the KL term acts as a latent regularizer and is absorbed into the approximation and optimization error. Let f,ϕ,t():=ℓ^,ϕ,t()f_ ψ, φ,t( τ):= _ ψ, φ,t( τ) and define the empirical and population squared reconstruction risks as ℛ^t(f,ϕ,t):=1B∑i=1B(f,ϕ,t(t,i)−ℓt(t,i))2, R_t(f_ ψ, φ,t):= 1B _i=1^B (f_ ψ, φ,t( τ_t,i)- _t( τ_t,i) )^2, (8) and ℛt(f,ϕ,t):=∼p[(f,ϕ,t()−ℓt())2].R_t(f_ ψ, φ,t):=E_ τ p [ (f_ ψ, φ,t( τ)- _t( τ) )^2 ]. (9) By the standard Rademacher uniform convergence bound applied to recG^rec, with probability at least 1−δ1-δ, uniformly over (,ϕ)∈Ψ×Φ( ψ, φ)∈ × , ℛt(f,ϕ,t)≤ℛ^t(f,ϕ,t)+2ℜB(rec)+M2ln(2/δ)2B.R_t(f_ ψ, φ,t)≤ R_t(f_ ψ, φ,t)+2 R_B(G^rec)+M^2 (2/δ)2B. (10) Let (^,ϕ^)( ψ, φ) denote the parameters obtained by minimizing the empirical squared reconstruction objective, and let f,ϕ,t∗f^*_ ψ, φ,t be a population minimizer in the class. Applying the above inequality to the empirical minimizer gives ℛt(ℓ^t) _t( _t) ≤ℛ^t(ℓ^t)+2ℜB(rec)+M2ln(2/δ)2B ≤ R_t( _t)+2 R_B(G^rec)+M^2 (2/δ)2B (11) ≤ℛ^t(f,ϕ,t∗)+2ℜB(rec)+M2ln(2/δ)2B. ≤ R_t(f^*_ ψ, φ,t)+2 R_B(G^rec)+M^2 (2/δ)2B. Applying the uniform convergence bound once more to f,ϕ,t∗f^*_ ψ, φ,t yields ℛ^t(f,ϕ,t∗)≤ℛt(f,ϕ,t∗)+2ℜB(rec)+M2ln(2/δ)2B. R_t(f^*_ ψ, φ,t) _t(f^*_ ψ, φ,t)+2 R_B(G^rec)+M^2 (2/δ)2B. (12) Combining the last two inequalities and using the definition of ϵapprec _app^rec gives ℛt(ℓ^t) _t( _t) ≤ϵapprec+4ℜB(rec)+2M2ln(2/δ)2B ≤ _app^rec+4 R_B(G^rec)+2M^2 (2/δ)2B (13) =ϵapprec+ϵgenrec(B,δ). = _app^rec+ _gen^rec(B,δ). Thus, ∼p[(ℓ^t()−ℓt())2]≤ϵapprec+ϵgenrec(B,δ).E_ τ p [( _t( τ)- _t( τ))^2 ]≤ _app^rec+ _gen^rec(B,δ). (14) It remains to account for the temporal shift from t θ_t to t+1 θ_t+1. By the L2L_2 triangle inequality and Jensen’s inequality, ∼p[|ℓ^t()−ℓt+1()|] _ τ p [| _t( τ)- _t+1( τ)| ] (15) ≤(∼p[(ℓ^t()−ℓt+1())2])1/2 ≤ (E_ τ p[( _t( τ)- _t+1( τ))^2] )^1/2 ≤(∼p[(ℓ^t()−ℓt())2])1/2+(∼p[(ℓt()−ℓt+1())2])1/2. ≤ (E_ τ p[( _t( τ)- _t( τ))^2] )^1/2+ (E_ τ p[( _t( τ)- _t+1( τ))^2] )^1/2. The first term is controlled by the reconstruction bound above. For the second term, since ℓ() ( τ) is L-Lipschitz in θ, |ℓt()−ℓt+1()|≤L‖t+1−t‖.| _t( τ)- _t+1( τ)|≤ L\| θ_t+1- θ_t\|. (16) The normalized policy update and the gradient bound imply ‖t+1−t‖ \| θ_t+1- θ_t\| =‖η1B∑i=1B∇ℓt(i)‖ = \|η 1B _i=1^B _ θ _t( τ_i) \| (17) ≤η1B∑i=1B‖∇ℓt(i)‖≤ηG. ≤η 1B _i=1^B\| _ θ _t( τ_i)\|≤η G. Therefore, (∼p[(ℓt()−ℓt+1())2])1/2≤LηG. (E_ τ p[( _t( τ)- _t+1( τ))^2] )^1/2≤ Lη G. (18) Substituting the two bounds gives ∼p[|ℓ^t()−ℓt+1()|] _ τ p [| _t( τ)- _t+1( τ)| ] ≤ϵapprec+ϵgenrec(B,δ)+LηG, ≤ _app^rec+ _gen^rec(B,δ)+Lη G, (19) which completes the proof. ∎ B.1 Proof of Proposition 2 Proposition 2 (Look-ahead reduces temporal mismatch, restated). Let t=t+1−t _t= θ_t+1- θ_t with ‖t‖≤ηG\| _t\|≤η G, and assume that ℓ(,) ( θ, τ) is β-smooth in θ. For analysis, define the first-order look-ahead risk as mt()=ℓt()+∇ℓ(t,)⊤tm_t( τ)= _t( τ)+ _ θ ( θ_t, τ) _t. If the ICRPM predictor satisfies ∼p|ℓ^t+1IC()−mt()|≤ϵICE_ τ p| ^IC_t+1( τ)-m_t( τ)|≤ _IC, then ∼p[|ℓ^t+1IC()−ℓt+1()|]≤ϵIC+β2η2G2.E_ τ p [| ^IC_t+1( τ)- _t+1( τ)| ]≤ _IC+ β2η^2G^2. (20) Proof. For any fixed task τ, denote ℓt+1()=ℓ(t+1,) _t+1( τ)= ( θ_t+1, τ) and ℓt()=ℓ(t,) _t( τ)= ( θ_t, τ). Since ℓ(,) ( θ, τ) is β-smooth in θ, the standard first-order remainder bound gives |ℓt+1()−ℓt()−∇ℓ(t,)⊤t|≤β2‖t‖2≤β2η2G2. | _t+1( τ)- _t( τ)- _ θ ( θ_t, τ) _t |≤ β2\| _t\|^2≤ β2η^2G^2. (21) By the definition of mt()m_t( τ), this implies |ℓt+1()−mt()|≤β2η2G2| _t+1( τ)-m_t( τ)|≤ β2η^2G^2. Therefore, using the triangle inequality and taking expectation over ∼p τ p, ∼p[|ℓ^t+1IC()−ℓt+1()|]≤∼p[|ℓ^t+1IC()−mt()|]+β2η2G2≤ϵIC+β2η2G2.E_ τ p [| ^IC_t+1( τ)- _t+1( τ)| ] _ τ p [| ^IC_t+1( τ)-m_t( τ)| ]+ β2η^2G^2≤ _IC+ β2η^2G^2. (22) This completes the proof. ∎ Appendix C Additional Preliminaries and Training Protocols This appendix summarizes the task formulations used in our experiments, covering zero-shot domain randomization, few-shot meta-reinforcement learning, and RL post-training with verifiable rewards. Although these settings differ in their interaction modalities, they share a common structure: at each iteration, a subset of tasks or prompts is selected, evaluated by the current policy, and appended to the optimization history for subsequent predictive sampling. C.1 Task Distributions in Randomized Environments Definition 3 (Identifier-induced MDP family). Let p()p( τ) be a distribution over task identifiers. Each identifier τ specifies an MDP ℳ=⟨,,P,R,λ¯⟩M_ τ= ,A,P_ τ,R_ τ, λ , where S and A are shared state and action spaces, while the transition kernel P(′∣,)P_ τ( s s, a) and/or the reward function R(,)R_ τ( s, a) may vary with τ. For an H-horizon trajectory ξ=(0,0,r0,…,H)ξ=( s_0, a_0,r_0,…, s_H), the discounted return is G(ξ)=∑h=0H−1λ¯hrh.G_ τ(ξ)= _h=0^H-1 λ^hr_h. (23) This formulation covers continuous-control tasks in robotics, where τ may encode physical parameters, target velocities, masses, lighting conditions, or goal positions. It also provides the task-level abstraction used by robust task sampling, where the sampler selects identifiers whose predicted difficulty is informative for policy improvement. C.2 Zero-Shot and Few-Shot Adaptive Decision-Making Domain randomization. In zero-shot domain randomization, the policy is trained over MDPs sampled from p()p( τ) and is directly evaluated on unseen environments without an adaptation phase. The standard objective is maxDR()=∼p()ξ∼π,ℳ[∑h=0H−1λ¯hrh]. _ θ\;J_DR( θ)=E_ τ p( τ)E_ξ _ θ,M_ τ [ _h=0^H-1 λ^hr_h ]. (24) For risk-based sampling, we use the negative return as the task risk, ℓ(,)=−ξ∼π,ℳ[∑h=0H−1λ¯hrh]. ( τ; θ)=-E_ξ _ θ,M_ τ [ _h=0^H-1 λ^hr_h ]. (25) In our experiments, TD3 is used for physical domain randomization [76]. Meta-RL with MAML. For few-shot adaptation, we use MAML [9] as the meta-learning backbone. Each task τ provides support trajectories SD_ τ^S for fast adaptation and query trajectories QD_ τ^Q for post-adaptation evaluation. The inner update is ′=−ηin∇ℓ(S,), θ _ τ= θ- _in _ θ (D_ τ^S; θ), (26) and the meta-objective is min∼p()[ℓ(Q,′)]. _ θ\;E_ τ p( τ) [ (D_ τ^Q; θ _ τ) ]. (27) Here, ℓ is instantiated as the negative query return after fast adaptation. Thus, DR can be viewed as the zero-shot case with no support set, while Meta-RL uses SD_ τ^S before evaluating QD_ τ^Q. C.3 RL Post-Training with Verifiable Rewards We also consider RL post-training for LLM reasoning. Here, each identifier τ denotes a prompt, or its embedding representation, sampled from a prompt pool T. At iteration t, the policy πt _ θ_t generates K responses t=yt,ii=1K y τ_t=\y τ_t,i\_i=1^K for each selected prompt. A verifier assigns rewards t=r(,yt,i)i=1K r τ_t=\r( τ,y τ_t,i)\_i=1^K. The latent success rate is ρt=[rt], _t τ=E[r_t τ], (28) which measures the probability that the current policy solves prompt τ. The optimization history is written as Ht=tB,ℛtB,ℛtB=t,jj=1B,H_t= \T_t^B,R_t^B \, _t^B= \ r τ_t,j_t \_j=1^B, where each i,j r_i τ_i,j denotes the reward vector over K responses for prompt i,j τ_i,j, and B is the number of selected prompts in the batch. RLVR optimizes the expected verifiable reward: max∼,y∼π(⋅∣)[r(,y)]. _ θ\;E_ τ ,\,y _ θ(· τ) [r( τ,y) ]. (29) In practice, we use GRPO [38], which replaces the value network with a group-normalized advantage. For responses generated from the same prompt, the advantage is A^i=ri−mean(rii=1K)std(rii=1K), A_i= r_i-mean(\r_i\_i=1^K)std(\r_i\_i=1^K), (30) and the policy is updated with a clipped importance-ratio objective regularized by the KL divergence to a reference policy. Online prompt selection. Prompt selection is needed because exact rollout-based evaluation is expensive. Evaluation-based methods over-sample a candidate pool tB^T_t B and keep prompts with non-degenerate reward signals, for example those satisfying std(rii=1K)>0std(\r_i τ\_i=1^K)>0. Such filtering is effective but requires extra generations. Prediction-based methods reduce this cost by estimating prompt difficulty, typically through a predictive model of p(ρt∣,t)p( _t τ τ, θ_t) or its history-conditioned variant. However, prompt-wise independent estimators may fail to share information across prompts and may lag behind the changing policy. C.4 Unified Predictive Selection View Both robust task sampling and active prompt selection can be expressed as selection from predicted outcomes. Let y^t y^t_ τ denote the model-predicted outcome for candidate τ. For active prompt selection, this outcome is the predicted success rate, y^t=ρ^t y^t_ τ= ρ^t_ τ, and the target is an intermediate difficulty level ρ∗ρ^*. For robust task sampling, the outcome is the predicted risk, y^t=ℓ^t() y^t_ τ= _t( τ), and the target is the largest predicted risk in the candidate pool. The selected batch is tB=argtop−B∈tB^(−‖y^t−ct∗‖22),T_t^B= *arg\,top-B_ τ _t B (- \| y^t_ τ-c_t^* \|_2^2 ), (31) where ct∗=ρ∗c_t^*=ρ^* for prompt selection and ct∗=max∈tB^ℓ^t()c_t^*= _ τ _t B _t( τ) for robust task sampling. Appendix D Risk-Aware Evaluation and Predictive Ranking Metrics This section summarizes the metrics and regularity conditions used to evaluate tail-task robustness and risk-prediction quality. We write the task risk induced by a model parameterized by θ as ℓ():=ℓ(Q,S,), _ θ( τ):= (D_ τ^Q,D_ τ^S; θ), where ∼p() τ p( τ) is a task identifier. Depending on the benchmark, ℓ() _ θ( τ) is instantiated as MSE for regression, negative return for control, or failure probability for RL post-training. D.1 CVaRαCVaR_α for Tail-Task Robustness Definition 4 (CVaRαCVaR_α). Let F(z)=Pr∼p()(ℓ()≤z)F_ θ(z)= _ τ p( τ) ( _ θ( τ)≤ z ) be the cumulative distribution function of the task risk, and define its α-quantile as qα()=infz:F(z)≥α.q_α( θ)= \z:F_ θ(z)≥α\. Following [29], the α-level Conditional Value-at-Risk is the expected risk over the upper tail: CVaRα()=[ℓ()|ℓ()≥qα()].CVaR_α( θ)=E [ _ θ( τ)\; |\; _ θ( τ)≥ q_α( θ) ]. (32) Equivalently, for a continuous risk distribution, CVaRα()=∫zdFα(z),Fα(z)=0,z<qα(),F(z)−α1−α,z≥qα(),CVaR_α( θ)= z\,dF_ θ^α(z), F_ θ^α(z)= cases0,&z<q_α( θ),\\[2.84526pt] F_ θ(z)-α1-α,&z≥ q_α( θ), cases (33) where FαF_ θ^α is the normalized tail-risk distribution. We denote the corresponding tail-task density by pα(,)p_α( τ; θ). Under this convention, CVaR0.9CVaR_0.9 evaluates the expected risk on the hardest 10%10\% of tasks. Thus, lower CVaR indicates stronger tail-task robustness. For control tasks, where risk is defined as negative return, reporting −CVaRα-CVaR_α gives the corresponding CVaR return, for which higher values are better. D.2 Spearman Rank Correlation for Risk Prediction While CVaRαCVaR_α measures the robustness of the learned policy, it does not directly assess whether the risk-predictive model correctly identifies hard tasks. Since predictive samplers mainly rely on the relative ordering of candidate tasks, we additionally report Spearman Rank Correlation (SRC) between predicted and observed risks. Given a candidate or validation set ii=1N\ τ_i\_i=1^N, let ℓ^i=ℓ^(i) _i= ( τ_i) be the predicted risk and ℓi=ℓ(iQ,iS,) _i= (D_ τ_i^Q,D_ τ_i^S; θ) be the observed risk. Denoting the rank operator by rk(⋅)rk(·), SRC is computed as SRC=1−6∑i=1N(rk(ℓ^i)−rk(ℓi))2N(N2−1).SRC=1- 6 _i=1^N (rk( _i)-rk( _i) )^2N(N^2-1). (34) SRC lies in [−1,1][-1,1]. Values close to 11 indicate that the predicted risk order matches the observed task-difficulty order. This metric is particularly relevant for predictive task sampling: even if the absolute risk values are not perfectly calibrated, a high SRC implies that the sampler can still prioritize informative or high-risk tasks. In contrast, low SRC indicates rank-flipping in the predicted difficulty landscape, which may lead to suboptimal task selection and weaker tail-task performance. Appendix E Experiment Details For all baselines, we follow the default configurations reported in their original papers whenever available, including model architectures, optimizer settings, task-selection rules, and hyperparameters. When a baseline is adapted to our shared MPTS-style evaluation protocol, we keep the training budget, candidate-pool size, batch size, number of seeds, and evaluation schedule consistent with ICRPM for fair comparison. E.1 Meta-Reinforcement Learning Experiments Shared Setup We evaluate on three MuJoCo Meta-RL benchmarks (HalfCheetahVel, Walker2dVel, ReacherPos) from the MAML task-distribution suite. Tasks are parameterised by continuous descriptors τ (target velocity or goal position) sampled from p()p( τ) at each iteration. All methods share the same meta-learner: MAML with TRPO as the inner optimiser. Key shared hyperparameters are listed in Table 3. Table 3: Shared meta-learning hyperparameters (Meta-RL). 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 Discount factor 0.99 GAE λ 1.0 Random seeds 10 Evaluation protocol. All methods are evaluated on a held-out validation set of 20 tasks sampled uniformly from p()p( τ) at each iteration; validation episodes are excluded from the environment interaction count. ICRPM Configuration for Meta-RL ICRPM uses Transformer backbone RPM that learns a predictive distribution p(ℓ|,Ct)p( | τ,C_t) conditioned on the latest context set Ct=(t,i,ℓt(t,i))i=1BC_t=\( τ_t,i, _t( τ_t,i))\_i=1^B. Architecture Details. The encoder maps context pairs (i,ℓi)i∈Ct\( τ_i, _i)\_i∈ C_t through a depth-4 MLP embedding, a Transformer encoder (self-attention over context), and a cross-attention module where each target task τ attends to the encoded context. The decoder outputs Gaussian parameters (μ,σ)(μ,σ) with σ=0.1+0.9⋅softplus(⋅)σ=0.1+0.9·softplus(·). The transformer-based ICRPM is parameter-matched to the vanilla RPM for a fair comparison (≈ 0.65 M parameters). Training objective. At each iteration t, one gradient step is taken on the weighted NLL objective (Eq. (4) in the main text): minϕ(ϕ)=−∑k=1Wγk(1B∑i=1Blnpϕ(ℓt−k+1(t−k+1,i)|t−k+1,i,Ht−k)), _ φ\;J( φ)=- _k=1^Wγ^k ( 1B _i=1^B p_ φ\! ( _t-k+1( τ_t-k+1,i)| τ_t-k+1,i,\,H_t-k ) ), (35) where Ht−k=(t−k,i,ℓt−k,i)i=1BH_t-k=\( τ_t-k,i, _t-k,i)\_i=1^B is the per-step history at time t−kt-k, W=3W=3 is the sliding-window size, γ=0.9γ=0.9 is the temporal decay, and Adam (lr =5×10−3=5× 10^-3) is used. Task selection. This retains the predictive selection setup in MPTS [22]. At each iteration, 1.5B1.5B candidate tasks are drawn from p(τ)p(τ). ICRPM predicts the mean risk y^i=pϕ(y|τ^i,Ht)[y] y_i=E_p_φ(y| τ_i,H_t)[y] for each candidate. For robust task sampling, we set c∗=maxiy^ic^*= _i y_i and select the B candidates with the largest score −‖y^i−c∗‖22-\| y_i-c^*\|_2^2, which is equivalent to selecting the top-B candidates by predicted risk in the scalar-risk case. E.2 RLVR Countdown Experiment Setup We use the Countdown mathematical reasoning task (20k training instances, Countdown-3to4) with Qwen3-4B-Base as the policy backbone π _ θ, trained via GRPO+ inside the veRL framework (v0.3.1) on 4×A100-80 GB GPUs. Key training hyperparameters are listed in Table 4. Table 4: RLVR training hyperparameters. 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×10−61× 10^-6 Entropy coefficient 0.001 KL coefficient 0.0 PPO clip ratio (low / high) 0.2 / 0.28 The reward function follows the format-aware Countdown setting: r=1r=1 (correct), r=0.1r=0.1 (incorrect but well-formatted), r=0r=0 (otherwise). Evaluation reports pass@1 accuracy on held-out CD34 and the harder CD4 benchmark. ICRPM Configuration for RLVR ICRPM replaces the prompt-difficulty predictor (RPM). At each training step, rollouts are generated by vLLM using the current actor. After reward computation, prompt embeddings ∈ℝ1024 τ ^1024 are extracted from the current actor’s last hidden states and paired with observed solve rates to form the context HtH_t, , where the solve rate is the fraction of generated responses that exactly solve the Countdown instance. These task representations are not frozen; their policy-induced drift is treated as part of the non-stationary task-risk context modeled by ICRPM. The ICRPM is trained with the objective in Eq. (35), using W=3W=3, γ=0.9γ=0.9, and 20 gradient steps per training step, matching the vanilla RPM. In RLVR, ICRPM is scaled to match the vanilla RPM parameter count (≈ 42.8 M). The task-sampling ratio is 80%. E.3 Domain Randomization Experiments For zero-shot DR we follow the MPTS protocol with TD3 as the policy backbone on Ergo-Reacher and Lunar-Lander [36]. Each task τ is an environment instance induced by a continuous domain-randomisation identifier; the risk ℓ()=L(t,) ( τ)=L( θ_t, τ) is the negative episodic return without inner-loop adaptation. The same task-batch size B, candidate-pool size B B, and baseline protocol as Meta-RL are used. ICRPM uses a parameter-matched Transformer architecture and the same temporal augmentation settings (W=3,γ=0.9W=3,γ=0.9). E.4 Few-Shot Sinusoid Regression Each task is specified by =(a,b) τ=(a,b) with f(x)=asin(x−b)f_ τ(x)=a (x-b), a∈[0.1,5.0]a∈[0.1,5.0], b∈[0,π]b∈[0,π]. The meta-learner is MAML and task risk is query-set MSE after adaptation. We report average MSE and CVaRα (Eq. (33)) with α∈0.5,0.7,0.9α∈\0.5,0.7,0.9\, along with SRC (Eq. (34)) throughout training. ICRPM is parameter-matched to the vanilla RPM and uses the same temporal augmentation settings (W=3W=3, γ=0.9γ=0.9). E.5 Ablation Studies Ablations are conducted on Lunar-Lander DR with 10 independent runs each. The ICRPM parameter count is equal to the vanilla RPM across all ablation studies for fair comparison. We vary four design dimensions: acquisition principle, RPM backbone, history augmentation, and the window/decay hyperparameters. 1. Acquisition principle. We compare worst-case selection by predicted risk μ(ℓ)μ( ) with a diversity-aware variant that combines μ(ℓ)μ( ) with a set-level diversity bonus. Both variants are evaluated under identical optimisation budgets and compared with PDTS [23]. This ablation tests whether the improvement comes from the risk predictor itself rather than a specific acquisition heuristic. 2. RPM backbone and history augmentation. We compare the default Transformer backbone with a parameter-matched MLP that retains the same context-conditioned look-ahead objective in Eq. (35). To isolate the role of history reuse from look-ahead prediction, we also evaluate MPTS-augment, which augments vanilla MPTS [22] with the same three-step history window and temporal decay γ=0.9γ=0.9 used by ICRPM, but without the context-conditioned one-step extrapolative objective. This comparison separates the effect of the backbone, the sliding-window augmentation, and the proposed look-ahead formulation. 3. Look-back window W. We vary W∈1,2,3,8W∈\1,2,3,8\ with γ=0.9γ=0.9 fixed. Smaller windows use fresher histories but provide less task coverage, while larger windows improve sample reuse but can introduce stale context from earlier policy states. We set W=3W=3 as the default trade-off between coverage and freshness. 4. Temporal decay factor γ. We vary γ∈0.5,0.7,0.9,1.0γ∈\0.5,0.7,0.9,1.0\ with the default window size fixed. Smaller values emphasize only the most recent batches, whereas γ=1.0γ=1.0 weights stale and fresh histories equally. Among the tested values, γ=0.9γ=0.9 provides a stable balance by mildly discounting older histories, and we therefore use it as the default. All ablation settings report CVaR0.9 return, average return, and SRC (Eq. (34)) on the task-distribution validation set. The results show that ICRPM remains effective across acquisition rules and backbone choices, while MPTS-augment does not match ICRPM under the same history-window and decay configuration. This indicates that the gains are not explained by history reuse alone, but mainly come from context-conditioned one-step look-ahead risk prediction. E.6 Compute resources. Table 5 reports approximate single-run wall-clock times, where each entry corresponds to one method on one benchmark/environment with one random seed, including training and periodic evaluation but excluding preliminary debugging runs. For Sinusoid and continuous-control experiments, ERM [30], MPTS [22], GDRM [31], OHTM [33], and ICRPM have comparable runtime, while DRM is about 30–40% slower due to its additional importance-weight computation. For RLVR Countdown, Uniform, PCL [14], MoPPS [13], and ICRPM have comparable runtime, while GRESO [34] is about 60% slower due to additional rollout-selection overhead. Table 5: Approximate single-run compute resources. Experiment Hardware Common Slower Sinusoid RTX 4090 24GB ∼ 12.2 min DRM: ∼ 16.4 min DR RTX 4090 24GB ∼ 3.0 h DRM: ∼ 4.2 h Meta-RL A10 24GB ∼ 8.7 h DRM: ∼ 11.6 h RLVR 4×A100 80GB ∼ 11.1 h GRESO: ∼ 17.8 h E.7 Assets, Licenses, and Terms of Use. All experiments use publicly available research assets. For LLM post-training, we use Qwen3-4B-Base as the policy backbone, and implement training and inference with veRL and vLLM. These model and software assets are released under the Apache License 2.0. For continuous-control experiments, we use OpenAI Gym with MuJoCo-based environments through the legacy mujoco-py interface. OpenAI Gym and mujoco-py are released under permissive open-source licenses, while the MuJoCo simulator is used subject to its applicable license terms; we do not redistribute MuJoCo binaries, license keys, or third-party assets. Our TD3 and MAML implementations are adapted from or benchmarked against publicly available research codebases released under permissive licenses. For Countdown, we use public arithmetic-reasoning datasets from Hugging Face for research evaluation only, do not redistribute raw dataset files, and provide sufficient preprocessing and evaluation details to support reproducibility. We comply with the terms of use of all upstream assets and retain the corresponding license and copyright notices where applicable. Appendix F Limitations ICRPM is proposed as a context-conditioned extrapolative RPM for improving risk approximation in MPTS-style task sampling. Several aspects remain open for future investigation. 1. History-based extrapolation. ICRPM relies on recent optimization histories to infer next-step task risk. This is effective when recent histories capture local risk drift, but may be less reliable under abrupt policy updates, highly noisy rewards, or task regions poorly covered by the history buffer. 2. Scope of validation and analysis. Our theory uses standard boundedness and smoothness assumptions to characterize temporal mismatch, and our experiments cover regression, robust continuous control, and RLVR-based LLM post-training. Extending the analysis and evaluation to larger models, longer horizons, and broader task distributions remains an important direction for future work. Appendix G Broader Impacts ICRPM improves sample-efficient task selection by reducing temporal-lag bias in online risk prediction. This can lower interaction and compute costs and improve robustness on difficult or tail tasks in simulated control and verifiable-reward reasoning settings. However, more efficient task selection may also reduce the cost of RL-based post-training for capable models, including applications with potential misuse. Moreover, learned risk predictors may amplify blind spots if the task distribution is biased or incomplete. We therefore recommend using such methods with broad task coverage, safety-oriented stress tests, and task-distribution audits before any real-world deployment. Our experiments use simulated environments and algorithmic reasoning benchmarks only, and do not involve human subjects, private user data, or real-world deployment.