ICRPM team posted
QbitAI is public.
There is an increasingly practical problem in Big Model Reinforcement Training, Robot Control and Meta-RL:
Individual assessment tasks are too expensive to predict risks with historical data and may not keep up with model updates.
How to find a really worthy Prompt with a limited rollout budget?
FromNational Defence and Technology UniversityandChinese UniversityAccording to the research team,ICRPM（In-Context Risk Predictive Model）Using recent missions and risk records to directly predict the next risk.
The method does not require additional rollout generation or increased environmental interaction for risk prediction, and improves mission risk prediction and selection in scenarios such as large model RLVR, Meta-RL and rod control, while basically maintaining the original training cost.
Background: Why is prediction accurate and why is it still wrong to choose a task?
Traditional task sampling requires first measuring the difficulty or risk of the candidate task. In the large model RLVR, this means generating multiple responses for Prompt and then estimating the success rate of solving the problem based on verifiable rewards; in Meta-RL and robot control, this often means expensive environmental interaction.
Model Predictive Task Sampling（MPTS）Using risk prediction models（RPM）Reducing this type of assessment: only real measurements of a small number of tasks, then predicting the risk of a larger set of candidates and sampling accordingly.Mission risk is not a static labelOnce the strategy is updated, the task that was previously difficult may become easier, and the order between tasks may even reverse. A RPM that is good at predicting the current risk threshold may not be able to accurately judge who will be selected for the next round.
△Following policy updates, the risk distribution of tasks is drifting, potentially causing a reversal in the order of tasks.Techniques: Let RPM move from reviewing the present to predicting the next step
Use context to predict future risks
ICRPM retains the basic task sampling framework of MPTS, focusing on replacing its risk forecasting methods.In the context of the mission mission, the observation of risk mites as a contextCombined with the task of evaluating, the forecasting strategy completes the next updated risk.
In this way, the historical record is no longer just an ordinary training sample, but also a clue to the state of the strategy. The model needs to be learned not only about the difficulty differences between different tasks, but also how the difficulty changes as training progresses.
△ICRPM uses recent mission risk records and candidate missions to predict the risk distribution of next steps.The paper's error analysis explains the necessity of this design. Traditional RPM predictions of current risk are used for the next step, with the existence of learning-related risks.One-step time lag errorIf the context predictor is able to capture local risk variations, the ICRPM can control the corresponding smoothness residues at a secondary measurement level. However, this does not mean that the prediction error is automatically reduced to zero: approximate error in the context of the model itself still exists.
Smaller supervision by the time decline window
The number of tasks that are actually evaluated in each round is limited, and a recent batch alone is likely to cause insufficient data; however, long-term accumulation of old data will blend in with outdated strategic status.
Adopted by ICRPMSliding windows with a decline in timeRepeating the context and forecast objectives of the three most recent historical batches and setting the time decline coefficient at 0.9. The greater the impact of the current forecast on the new records, the older records can still provide a surveillance signal, but the weight gradually decreases, thus striking a balance between supplementing surveillance and avoiding outdated data interference.
Transformers to capture cross-tasking relationships
In implementation, ICRPM encodes tasks and corresponding risks into inputs, modeling the relationships between historical tasks by Transformer, and then reading information related to candidate tasks through cross-attention, ultimately producing the next step of risk prediction.
It is worth noting that the experiment matched parameters for ICRPM with the traditional RPM: about 0.65M parameters in Meta-RL and about 42.8M parameters in RVR. Therefore, the performance difference cannot simply be explained as an enlarged predictive model.The training objectives are still based on existing optimization processes, with no need for additional interaction or rollout for risk models.
Dissolution verification: the key is not just to look at history
Research on LunarLander compared different model structures, task selection rules, and history window settings. By replacing Transformer with parametrically matched MLP, ICRPM was still superior to traditional MPTS; MPTS added the same history window and time decline, but could not predict the next step and still not keep up with ICRPM.
△Different sampling rules, predictor structure, and historical enhancement methods of melting contrast.This suggestsThe real key is predicting the outcome in the context above.Windows sensitivity experiments also show that history isn't getting any better: the paper uses 3-step windows with a decline coefficient of 0.9 to strike a balance between supplementing surveillance and avoiding obsolete data interference.
Experimental results: From robot control to large model RLVR
The paper covers a small sample of string regression, Meta-RL, Domain Randomization（DR）The evaluation focuses not only on average performance, but also on consistency in task risk ranking and performance on difficult tasks.
Small sample returns: Can improvements in risk ranking turn into gains?
In the few-shot Sinusoid Regression experiment, the ICRPM not only reduced the average median error（MSE）Spearman ranking coefficients between tail risk error, predicted risk and actual risk（SRC）In the meta-test phase, ICRPM achieved the lowest MSE of all comparison methods at all three assessment levels of VaR0.5, VaR0.7 and VaR0.9, while maintaining a higher risk ranking relevance, indicating that improvements in task difficulty ranking can be transmitted to downstream performance.
△Average errors, tail errors, risk ranking correlations and test performance of a string regression experiment.Continuous Control: Improving Hard Task Performance
In the five environments involved in Meta-RL and DR, ICRPM overall demonstrated stronger tail risk ruggedness. In Walker2dVel, MPTS's CVaR0.9 returns were 27.9, and the ICRPM was upgraded to 27.9.-1.5On LunarLander, the paper reported a CVaR 0.9 improvement over ERM at ICRPM convergence.More than 40 percent。
△The training results of two DR environments and three Meta-RL environments subsequently showed CVaR 0.9 returns, average returns and risk ranking correlations (SRC).
△The final test CVaR results for five environments, with different algorithms and error ranges.Specifically, Walker2dVel's CVaR0.9 was upgraded from MPTS's 27.9±19.1 to ICRPM.-1.5±7.2HalfCheetahVel was upgraded from 184.0 to-180.2ErgoReacher is up from 75.778.3But not all individual indicators lead: the ICRPM on the ReacherPos is 27.8, slightly below 26.5 on the DRM. Consequently, the conclusion should be understood as an overall trend improvement rather than an absolute advantage over all tasks.
Large model RLVR: Increased countdown accuracy
In a Countdown experiment based on Qwen3-4B-Base, the training set contained about 20,000 Countdown-34 samples, each of which generated eight answers for Prompt. ICRPM's ranking correlation with Prompt difficulty was significantly better than the comparison method and yielded better final accuracy.
△A training curve related to the accuracy of the countdown test and the predictive ordering of the Prompt difficulty.
△In contrast to the accuracy of the countdown test, the CD4 and CD34 respectively represent two sets of ratings, with Avg. having two averages.In the more difficult CD4 test set, the ICRPM's pass@1 reached55.7%It was 5.8 percentage points higher than GRESO at 53.8% and 2.8 percentage points higher than MoPPS at 52.9%.75.1%The average of the two65.4%The average accuracy of ICRPM was 2.9, 1.6 and 2.0 percentage points higher than Uniform's 62.5%, GRESO's 63.8% and MoPPS's 63.4%.
Calculation costs: no additional rollout, running costs close to the original method
In terms of computational costs, the paper's appendix reported that in an experiment with RLVR, ICRPM and Uniform, PCL, and MoPPS were running approximately the same amount of time.11.1 hoursOn Walker2dVel, the relative running time of ICRPM is approximately 17.8 hours.1.08 times the ERMThis is close to 1.07 times the MPTS.
△Walker2dVel's running time versus tail return. ICRPM achieved better CVR0.9 returns at close to the running cost of MPTS.Project value: task sampling needs to keep pace with model changes
The idea of ICRPM is that the difficulty of the task is not constant.A good task selector not only knows what tasks are difficult at the moment, but also decides what tasks are worth training after updating the strategy.Through forward-looking risk prediction and recent historical replication, research has improved the reliability of task risk sequencing while maintaining basic sampling and running costs, and validated results through continuous control and post-large model training.
For the training process that needs to be refined using the rollout budget, this provides a way to improve training efficiency from the task selection link without having to replace the downstream optimizer.
Title: Bridging Risk Approximation Gaps in Model Predictive Task Sampling via In-Context Modeling
The paper was written by Jiarong Wen, Qi Tao, Kaiyu Zhang, Yun Qu, Lizhou Cai, Heming Zou, Yuhang Jiang, Cheems Wang.
Link to the paper:https://arxiv.org/pdf/2610.14913