PaperlayerView exact revision on arXiv

Source-linked AI summary

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

arXiv:2610.14913v1cs.AIcs.LG

TL;DR

Expensive task evaluation limits task selection in robust RL and LLM post-training, while vanilla risk predictors are challenged by temporal lag and sparse histories. The paper introduces in-context risk prediction with one-step look-ahead and a temporal-weighted history window. It reports improved risk ranking, sample efficiency, and tail-risk performance across several settings.

  • Problem

    Policy evaluation requires extensive environment trajectories in robust RL and long Chain-of-Thought rollouts in LLM post-training, while vanilla RPMs face temporal-lag bias and sparse-history generalization limits.

    Source: §1 · p. 1

  • Method

    ICRPM predicts evolving task risk with one-step look-ahead and a temporal-weighted sliding window that reuses recent history while downweighting stale data.

    Source: §1 · p. 2

  • Results

    ICRPM improves risk ranking, sample efficiency, and tail-risk performance across regression, robust Meta-RL, domain randomization, and RLVR-based LLM post-training.

    Source: §5 · p. 10

  • Takeaways & Limitations

    Task samplers should model policy-induced risk dynamics rather than rely on static or lagged difficulty estimates.

    Source: §5 · p. 10

    The formal lagged-risk bound assumes bounded Lipschitz risks, bounded gradients, and i.i.d. update samples.

    Source: §3.1 · p. 4

Abstract

from arXiv · show

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.

1 Introduction

Expensive task evaluation motivates MPTS, which uses optimization history to prioritize tasks, but its risk predictors face lag and sparse-history limitations. The paper proposes in-context prediction and reports gains across robust RL and LLM post-training.

Conditional Agent-Environment Interactions via Model Predictive Strategies

  • MPTS reduces unnecessary interactions by training risk predictive models on optimization histories, but vanilla models face temporal-lag bias and sparse-history generalization limits.These issues also affect prompt curriculum learning, where lag and sparse evaluations can distort curricula and reduce sample efficiency.Source: §1 · p. 2

Enhancing Risk Predictive Models via In-Context Sequence Modeling

  • ICRPM models evolving task risk through one-step look-ahead prediction and a temporal-weighted sliding window, addressing extrapolation bias and data scarcity without changing policy optimization.The window reuses recent trajectories while downweighting stale data, with negligible additional rollout cost.Source: §1 · p. 2

Primary Contributions and Empirical Findings

  • A lagged-risk decomposition separates finite-sample reconstruction error from an O(η) policy-update mismatch that can destabilize task-risk rankings.The analysis motivates modeling policy-induced changes rather than relying on lagged risk estimates.Source: §1 · p. 2
  • ICRPM improves sample efficiency and tail-risk robustness on robust DR and Meta-RL benchmarks, and accelerates curriculum-based RL post-training for LLMs.The reported findings span continuous-control and LLM reasoning tasks.Source: §1 · p. 2

2 Research Background

The background formalizes tasks and their policy-dependent risks, then describes MPTS as a predict-then-select process driven by a risk predictive model trained on optimization history. Its estimates guide task choice, but vanilla models use risks estimated under the current policy after a policy update.

Notation and Setup

  • Tasks include continuous descriptors for DR and Meta-RL and discrete reasoning prompts for RLVR; risk is defined for a policy on each task at each optimization step.The policy may be a control policy or an autoregressive model.Source: §2 · p. 2

General MPTS Pipeline

  • MPTS predicts task risk or prompt success, then selects tasks near a target threshold for curriculum learning or near-worst-case risk for robust sampling.Prompt selection typically targets a success rate around 0.5.Source: §2 · p. 3
  • MPTS updates the RPM from recent history, evaluates a candidate pool in batches, and uses the selected tasks for the next policy update.Repeated rounds support prompt curriculum learning, robust Meta-RL, and DR.Source: §2 · p. 3

RPM and Variational Inference

  • Existing RPMs commonly use a VAE-based latent predictor trained with streaming variational inference, then estimate task outcomes through Monte Carlo forward passes.The resulting predictions feed the acquisition rule for robust or curriculum optimization.Source: §2 · p. 3
  • Vanilla RPMs estimate risk under policy θ_t but use it after updating to θ_t+1, creating a temporal mismatch in task difficulty estimates.Risk rankings may shift as the policy changes.Source: §2 · pp. 2–3

3 From Theoretical Investigation to Practical Enhancement

This section diagnoses why vanilla risk prediction can mis-rank tasks as policies change, then introduces ICRPM to predict next-step risk using optimization context and reused recent histories.

3.1 Vulnerabilities of Vanilla RPM

  • Vanilla RPM error combines finite-sample reconstruction and approximation error from sparse histories with O(η) policy-update mismatch, which can distort task-risk rankings.Source: §3.1 · p. 4

3.2 Context-based Extrapolative Risk Prediction

  • ICRPM conditions look-ahead prediction on recent optimization context to target next-step risk rather than reconstructing lagged risk, alleviating O(η) temporal bias.Source: §3.2 · pp. 4–5
  • A temporal-weighted sliding window reuses recent context-target history pairs, increasing training pairs while exponentially downweighting outdated histories.Source: §3.2 · p. 5

3.3 Implementation and Neural Architecture

  • Each iteration uses prior-policy task-risk pairs as context and latest-policy task-risk pairs as targets; Transformer tokenization captures correlations between target and past risk sequences.Source: §3.3 · p. 6

4 Results and Analysis

The experiments test whether ICRPM improves task-risk rankings and downstream performance across regression, robust Meta-RL, domain randomization, and LLM post-training, using SRC and task-specific performance metrics. Ablations examine which components account for its gains.

4.1 Robust Few-Shot Sinusoid Regression

  • In sinusoid regression, ICRPM lowers average and CVaR0.9 MSE, converges faster, maintains higher SRC, and achieves the lowest meta-test CVaRα MSE across α∈{0.5, 0.7, 0.9}.Vanilla MPTS has less stable convergence and worse worst-case performance.Source: §4.1 · p. 7, Figure 3 · p. 7

4.2 Robust Meta-RL and Domain Randomization

  • On Walker2dVel, CVaR0.9 is −1.5 vs. −27.9 for MPTS; ICRPM leads validation CVaR0.9 on all three Meta-RL benchmarks with competitive average returns.At meta-test, it achieves the best CVaRα returns across nearly all environments and robustness levels, with substantially reduced variance.Source: §4.2 · pp. 8–9, Table 1 · p. 8, Figure 4 · p. 8
  • On Lunar-Lander, ICRPM raises convergence CVaR0.9 by over 40% vs. ERM; it leads CVaR0.9 on Ergo-Reacher and Lunar-Lander with competitive average returns.Its SRC advantage over MPTS persists in both domain-randomization environments.Source: Figure 4 · p. 8, §4.2 · p. 9

4.3 Sample-Efficient RLVR on Countdown

  • 55.7% CD4 accuracy: ICRPM exceeds GRESO by 1.9 points and MoPPS by 2.8, with similar gains on CD34 and faster convergence.Its SRC curve indicates more reliable online prompt-difficulty rankings as the policy evolves.Source: §4.3 · p. 9, Table 2 · p. 9, Figure 6 · p. 9

4.4 Ablation Studies

  • Ablations show ICRPM’s gains across acquisition rules and backbones mainly come from context-conditioned look-ahead prediction, not the acquisition rule, Transformer, or history reuse alone.Window size affects performance more than decay factor: W=3 gives the best CVaR0.9 performance, and γ=0.9 is used as a robust default.Source: §4.4 · p. 10, Figure 7 · p. 10, Figure 8 · p. 10

5 Conclusion

The conclusion identifies policy-update lag in vanilla risk predictive models as a source of biased risk estimates and unstable rankings, then presents ICRPM as an in-context, look-ahead alternative. The related work situates this approach among robust optimization, task sampling, curriculum learning, and LLM post-training methods.

5 Conclusion

  • Vanilla RPMs estimate risk under the current policy but use it for post-update task selection, inducing first-order lagging bias and unstable task rankings.ICRPM combines context-conditioned one-step extrapolation with temporal-weighted history augmentation; experiments report gains across regression, robust Meta-RL, domain randomization, and RLVR-based LLM post-training.Source: §5 · p. 10

A Related Work

  • Robust objectives including DRM and GDRM address high-risk tasks or subpopulation shifts, but evaluation over large task spaces remains expensive.ERM optimizes average loss; DRM emphasizes high-risk tasks via CVaR [17] [18] [32], while GDRM reweights group-level risks.Source: A Related Work · p. 14
  • RATS and MPTS use learned risk-predictive models to select tasks efficiently, but sparse early histories and pre-update risk estimates challenge prediction.PDTS combines posterior sampling with diversity regularization [23]; this work targets RPM limitations with in-context sequence modeling and one-step look-ahead.Source: A Related Work · p. 14
  • ICRPM is presented as orthogonal to curriculum-design strategies: it improves risk-estimate accuracy and timeliness under a given task distribution.This makes the approach potentially compatible with existing curricula.Source: A Related Work · p. 15
  • The framework maps prompts to task identifiers and uses an RPM to estimate prompt difficulty for LLM post-training curricula without additional generation overhead.Prompt-level difficulty estimation otherwise requires costly rollouts from the current model.Source: A Related Work · p. 15

B Expanded Proofs for Theoretical Analysis

The appendix proves bounds for vanilla RPM lagged prediction and for ICRPM look-ahead. The analyses separate reconstruction error from temporal mismatch induced by policy updates.

B Expanded Proofs for Theoretical Analysis

  • The appendix states that its proofs cover Theorem 1 and Proposition 2, with notation defining task risk, the cumulative trajectory history, and RPM hypothesis complexity.ICRPM uses W consecutive history batches, totaling N = WB, whereas vanilla RPM is trained on the latest batch of size B.Source: B Expanded Proofs for Theoretical… · p. 15
  • The vanilla RPM prediction bound separates reconstruction approximation error from generalization error under bounded, Lipschitz risk and sampling assumptions.Its proof controls reconstruction with a uniform-convergence bound, then accounts for the shift from θ_t to θ_t+1.Source: B Expanded Proofs for Theoretical… · pp. 15–16

B.1 Proof of Proposition 2

  • Proposition 2 analyzes first-order look-ahead risk under a bounded policy update and β-smooth task risk.The proof compares the first-order expansion at θ_t with risk at θ_t+1, using the smoothness remainder and triangle inequality.Source: §B.1 · p. 17

C Additional Preliminaries and Training Protocols

The appendix specifies task formulations and training protocols across control, Meta-RL, and LLM post-training. It also describes online prompt selection, a shared predictive-selection view, and benchmark-specific risk definitions.

C Additional Preliminaries and Training Protocols

  • The appendix frames the settings as task selection and evaluation followed by adding observed outcomes to optimization history for later predictive sampling.Its task family allows identifiers to specify MDPs with shared state and action spaces but task-varying transitions or rewards.Source: C Additional Preliminaries and… · p. 17, §C.1 · p. 17

C.1 Task Distributions in Randomized Environments

  • Randomized-environment tasks use identifiers that can encode physical parameters or goals; domain randomization evaluates unseen environments without adaptation, while Meta-RL adapts from support trajectories.The described DR protocol uses TD3 for physical domain randomization.Source: §C.1 · p. 17, §C.2 · p. 18

C.3 RL Post-Training with Verifiable Rewards

  • For RLVR, prompts are task identifiers, solve rates enter the optimization history, and GRPO uses group-normalized advantages with a clipped policy objective regularized by KL divergence.The appendix also describes online prompt selection, noting that exact rollout evaluation is expensive and prompt-wise predictors may fail to share information or lag behind policy changes.Source: §C.3 · p. 18

C.4 Unified Predictive Selection View

  • The unified view treats robust task sampling and active prompt selection as choosing candidates according to predicted outcomes, with different targets for risk and prompt success.The passage identifies predicted success rate for prompt selection and predicted risk for robust sampling.Source: §C.4 · p. 19

D Risk-Aware Evaluation and Predictive Ranking Metrics

D.1 CVaR𝛼for Tail-Task Robustness

The section defines tail-task robustness with CVaR and evaluates predictive task ordering with Spearman Rank Correlation. It then gives the Meta-RL experiment protocol and ICRPM configuration.

D.1 CVaR𝛼for Tail-Task Robustness

  • CVaR_α measures expected risk in the upper tail; under the stated convention, CVaR_0.9 averages risk over the hardest 10% of tasks.For control, −CVaR_α is the corresponding CVaR return, where higher values are better.Source: §D.1 · p. 19

D.2 Spearman Rank Correlation for Risk Prediction

  • Spearman Rank Correlation measures agreement between predicted and observed task-risk rankings, complementing CVaR, which evaluates policy robustness rather than predictor ranking.Values near 1 indicate matching difficulty order; low SRC indicates rank-flipping.Source: §D.2 · pp. 19–20

E Experiment Details

  • Meta-RL experiments use three MuJoCo benchmarks, MAML with TRPO as inner optimizer, and a held-out validation set of 20 tasks excluded from interaction counts.Baseline comparisons follow original configurations where available and align key training budgets and evaluation settings under the shared protocol.Source: §E.1 · p. 20, E Experiment Details · p. 20

E.1 Meta-Reinforcement Learning Experiments

  • The Meta-RL ICRPM uses a parameter-matched Transformer and trains with a temporally weighted sliding window; selection chooses the highest predicted-risk candidates.The configuration uses W = 3 and γ = 0.9.Source: §E.1 · pp. 20–21

E.2 RLVR Countdown Experiment

The appendix specifies the Countdown RLVR setup and ICRPM configuration, alongside domain-randomization and sinusoid-regression experiment protocols.

E.2 RLVR Countdown Experiment

  • Countdown RLVR uses Qwen3-4B-Base with GRPO+ and evaluates pass@1 on held-out CD34 and the harder CD4 benchmark.The reward is 1 for correct answers, 0.1 for incorrect but well-formatted answers, and 0 otherwise; ICRPM uses 80% task sampling and is parameter-matched to the vanilla RPM.Source: §E.2 · p. 21

E.3 Domain Randomization Experiments

  • Zero-shot DR evaluates TD3 on Ergo-Reacher and LunarLander, using negative episodic return as task risk and no inner-loop adaptation.ICRPM is parameter-matched and uses W = 3 and γ = 0.9.Source: §E.3 · p. 21

E.4 Few-Shot Sinusoid Regression

  • The few-shot sinusoid-regression tasks vary amplitude and phase, use MAML, and measure post-adaptation query-set MSE and CVaR alongside SRC.ICRPM is parameter-matched to vanilla RPM and uses W = 3 and γ = 0.9.Source: §E.4 · p. 21

E.5 Ablation Studies

The ablations attribute ICRPM’s gains chiefly to context-conditioned look-ahead, while compute comparisons report similar runtimes to several baselines. The section also documents the assets and licenses used.

E.5 Ablation Studies

  • Ablations vary acquisition, backbone, history reuse, window, and decay, with 10 Lunar-Lander DR runs and parameter-matched ICRPM and vanilla RPM.The default window is W=3 and the default decay is γ=0.9, chosen as a balance between history coverage and freshness.Source: §E.5 · p. 22
  • ICRPM remains effective across acquisition rules and backbones, while MPTS-augment does not match it, indicating gains mainly from context-conditioned one-step look-ahead.The ablations report CVaR0.9 return, average return, and SRC on the task-distribution validation set.Source: §E.5 · p. 22

E.6 Compute resources.

  • DRM is about 30–40% slower than the compared Sinusoid and continuous-control methods, and GRESO is about 60% slower than the compared Countdown methods.The passage attributes the additional time to importance-weight computation for DRM and rollout-selection overhead for GRESO.Source: §E.6 · p. 22
  • For Sinusoid and continuous-control, ICRPM has comparable runtime to ERM, MPTS, GDRM, and OHTM; for Countdown, it is comparable to Uniform, PCL, and MoPPS.Table 5 entries are approximate single-run times for one method, benchmark/environment, and seed, including training and periodic evaluation.Source: Table 5 · p. 22, §E.6 · p. 22

E.7 Assets, Licenses, and Terms of Use.

  • The experiments use Qwen3-4B-Base with veRL and vLLM for LLM post-training, MuJoCo-based control environments, and public arithmetic-reasoning datasets for Countdown.The authors state that they comply with upstream asset terms and do not redistribute MuJoCo binaries, license keys, or raw Countdown datasets.Source: §E.7 · pp. 22–23

F Limitations

ICRPM’s history-based extrapolation may be less reliable when recent histories poorly reflect current risk, and evaluation remains bounded by the stated assumptions and tested settings. The broader-impact discussion also identifies misuse and blind-spot risks and limits the evidence to simulated environments and algorithmic reasoning benchmarks.

F Limitations

  • History-based extrapolation may be less reliable under abrupt policy updates, highly noisy rewards, or task regions poorly covered by the history buffer.The analysis uses standard boundedness and smoothness assumptions; validation covers regression, robust continuous control, and RLVR-based LLM post-training.Source: F Limitations · p. 23
  • The analysis and evaluation have not been extended to larger models, longer horizons, or broader task distributions.The reported experiments use simulated environments and algorithmic reasoning benchmarks only, without real-world deployment.Source: F Limitations · p. 23, G Broader Impacts · p. 23

G Broader Impacts

  • Lower post-training costs may also reduce barriers for capable models, while biased or incomplete task distributions may let learned risk predictors amplify blind spots.The authors recommend broad task coverage, safety-oriented stress tests, and task-distribution audits before real-world deployment.Source: G Broader Impacts · p. 23
Loading 2610.14913v1…