Where should the rollout budget be most appropriately allocated? Not on Transformer, but on taking one step further in risk prediction.
After post-training of large model reinforcement learning, robot control and Meta-RL, there is an increasingly practical problem:
Evaluating tasks one by one is too costly, while using historical data to predict risks may fail to keep up with model updates.
How to find the truly training-worthy Prompts with a limited rollout budget?
The research team from National University of Defense Technology and Tsinghua University accordingly proposes ICRPM (In-Context Risk Predictive Model), which uses recent task and risk records to directly predict the next-step risk.
This method does not need to generate additional rollouts or increase environment interactions for risk prediction, and while basically maintaining the original training overhead, it improves the task risk prediction and selection performance in scenarios such as large model RLVR, Meta-RL and robust control.
Research Background: Why Are Tasks Still Selected Incorrectly Even With Accurate Prediction?
Traditional task sampling requires first measuring the difficulty or risk of candidate tasks. In large model RLVR, this means generating multiple responses for the Prompt and then estimating the problem-solving success rate based on verifiable rewards; in Meta-RL and robust control, it often means expensive environment interactions.
Model Predictive Task Sampling (MPTS) reduces such evaluations through the Risk Prediction Model (RPM): it only performs real measurements on a small number of tasks, then predicts the risks of a larger candidate set and samples accordingly. However, task risk is not a static label. After policy update, originally difficult tasks may become easier, and the ranking between tasks may even be reversed. An RPM that is good at predicting "current risk" may not be able to accurately judge "which one to choose in the next round".
△
Technical Solution: Make RPM Shift From Looking Back at Current Status to Predicting the Next Step
Predict Future Risks With Context
ICRPM retains the basic task sampling framework of MPTS, and focuses on replacing the risk prediction method in it. Traditional RPM mainly fits the risk under the current policy; ICRPM takes the "task-risk" observation of the latest step as context, combines it with the task to be evaluated, and predicts the risk after the policy completes the next update.
In this way, historical records are no longer just ordinary training samples, but also clues reflecting the state of the policy. What the model needs to learn is not only the difficulty difference of different tasks, but also how the difficulty changes with the training progress.
△
The error analysis in the paper explains the necessity of this design. When traditional RPM predicts the current risk and then uses it for the next step, there is a first-order time lag error related to the learning rate. If the context predictor can capture local risk changes, ICRPM can control the corresponding smoothness residual at the second-order magnitude. However, this does not mean that the prediction error automatically returns to zero: the model's own context approximation error still exists.
Use Time Decay Window to Compensate for Sparse Supervision
The number of tasks for real evaluation in each round is limited, and relying solely on the latest batch is likely to cause insufficient data; but if old data is accumulated for a long time, outdated policy states will be mixed in.
ICRPM adopts a sliding window with time decay, reuses the context and prediction targets formed by the latest 3 historical batches, and sets the time decay coefficient to 0.9. Newer records have a greater impact on current predictions, while older records can still provide supervision signals but their weights are gradually reduced, thus achieving a balance between supplementing supervision and avoiding interference from outdated data.
Capture Cross-Task Relationships via Transformer
In implementation, ICRPM encodes tasks and corresponding risks into input representations, uses Transformer to model the relationships between historical tasks, then reads information related to candidate tasks through cross-attention, and finally outputs the risk prediction for the next step.
It is worth noting that the experiment matched the parameter quantities of ICRPM and traditional RPM: about 0.65M parameters in Meta-RL, and about 42.8M parameters in RLVR. Therefore, the performance difference cannot be simply explained by the larger prediction model. The training target still comes from the existing optimization process, and no additional environment interactions or rollouts are required for the risk model.
Ablation Verification: The Key Is Not Only to Look at More History
The study compared different model structures, task selection rules and historical window settings on LunarLander. When Transformer is replaced with a parameter-matched MLP, ICRPM still outperforms traditional MPTS; when MPTS is added with the same historical window and time decay but does not predict the next step, it still cannot catch up with ICRPM.
△
This indicates that the truly critical factor is forward-looking prediction under contextual conditions, not just using Transformer or adding more historical data. The window sensitivity experiment also shows that longer history is not always better: the paper adopts a 3-step window and a decay coefficient of 0.9 to achieve a balance between supplementing supervision and avoiding interference from old data.
Experimental Results: From Robot Control to Large Model RLVR
The paper covers four scenarios: few-shot sinusoid regression, Meta-RL, Domain Randomization (DR) and LLM RLVR. The evaluation focuses not only on average performance, but also on the consistency of task risk ranking and performance on difficult tasks.
Few-Shot Regression: Can Improved Risk Ranking Be Translated Into Benefits?
In the Few-shot Sinusoid Regression experiment, ICRPM not only reduces the average mean squared error (MSE) and tail risk error, but also its Spearman rank correlation coefficient (SRC) between predicted risk and real risk is continuously better than that of MPTS. In the Meta-test stage, ICRPM achieved the lowest MSE among all compared methods at the three evaluation levels of CVaR₀.₅, CVaR₀.₇ and CVaR₀.₉, while maintaining higher risk ranking correlation, indicating that the improvement of task difficulty ranking can be transmitted to downstream performance.
△
Continuous Control: Improve Performance on Difficult Tasks
In the five environments involved in Meta-RL and DR, ICRPM as a whole shows stronger tail risk robustness. On Walker2dVel, the CVaR₀.₉ reward of MPTS is -27.9, and ICRPM increases it to -1.5; on LunarLander, the paper reports that the CVaR₀.₉ of ICRPM at convergence is over 40% higher than that of ERM.
△
△
Specifically, the CVaR₀.₉ of Walker2dVel increases from -27.9±19.1 of MPTS to -1.5±7.2 of ICRPM; HalfCheetahVel increases from -184.0 to -180.2, and ErgoReacher increases from 75.7 to 78.3. However, not all individual indicators are leading: the ICRPM on ReacherPos is -27.8, slightly lower than -26.5 of DRM. Therefore, the conclusion should be understood as an overall trend improvement, not an absolute advantage over all tasks.
Large Model RLVR: Countdown Accuracy Improvement
In the Countdown experiment based on Qwen3-4B-Base, the training set contains about 20,000 Countdown-34 samples, and 8 responses are generated for each Prompt. ICRPM's ranking correlation of Prompt difficulty is significantly better than that of the comparison methods, and brings better final accuracy.
△
△
On the more difficult CD4 test set, the pass@1 of ICRPM reaches 55.7%, 1.9 percentage points higher than 53.8% of GRESO and 2.8 percentage points higher than 52.9% of MoPPS; it reaches 75.1% on CD34, with an average of 65.4%. Compared with the average accuracy of 62.5% of Uniform, 63.8% of GRESO and 63.4% of MoPPS, ICRPM is 2.9, 1.6 and 2.0 percentage points higher respectively.
Computational Overhead: No Additional Rollout Required, Running Cost Is Close to the Original Method
In terms of computational cost, in one RLVR experiment reported in the appendix of the paper, the running time of ICRPM, Uniform, PCL and MoPPS is all about 11.1 hours, and that of GRESO is about 17.8 hours. On Walker2dVel, the relative running time of ICRPM is 1.08 times that of ERM, which is close to 1.07 times of MPTS.
△
Project Value: Task Sampling Needs to Keep Up With Model Changes
The idea proposed by ICRPM is that task difficulty is not fixed, a good task selector should not only know what tasks are difficult now, but also judge what tasks are worth training after policy update. Through forward-looking risk prediction and reuse of recent history, the study improves the reliability of task risk ranking while basically maintaining the sampling and running overhead, and verifies the effect in continuous control and large model post-training.
For training processes that need to carefully use the rollout budget, this provides an idea to improve training efficiency from the task selection link without replacing the downstream optimizer.
Paper Title: Bridging Risk Approximation Gaps in Model Predictive Task Sampling via In-Context Modeling
Paper Authors: Jiarong Wen, Qi Tao, Kaiyu Zhang, Yun Qu, Lizhou Cai, Heming Zou, Yuhang Jiang, Cheems Wang
Paper Link: https://arxiv.org/pdf/2610.14913
This article is from the WeChat Official Account "QbitAI", author: ICRPM Team, published with authorization from 36Kr.