Contextual Reinforcement Learning (CRL) seeks to generalize classical RL by maximizing task coverage across a context space of related tasks. While prior works often train from scratch and rely on either multi-task learning for a single policy or strategically training multiple policies, we advocate for a unified alternative: pretraining a single policy with good initial performance, followed by fine-tuning multiple policies for task specialization. This new paradigm, however, introduces unique challenges, such as heterogeneous marginal returns and sample inefficiency. This raises a critical research question: given a pretrained policy and a constrained budget, how much fine-tuning should each task region receive to enable sample-efficient CRL? To this end, we propose Task Specialization Fine-Tuning (TSFT), an online framework that predicts fine-tuning performance with a simple parametric model and exactly solves the resulting discrete budget allocation problem via integer linear programming. Extensive experiments across diverse decision domains, including combinatorial optimization, continuous control, and LLM fine-tuning, demonstrate that TSFT significantly outperforms baselines in task coverage and approaches oracle performance. Our work charts a new direction for model-based CRL, aligning with the modern pretrain-finetune era.
Figures & tables
Paradigm
# Policy
Training Method
Training Mode
Independent Training
Multiple N
Single-task
From scratch
Multi-Task Training
Single
Multi-task
From scratch
Multi-Policy Training
Multiple N≪N
Single-task
From scratch
TSFT (Ours)
Multiple N≪N
Multi-task
Pretrain + Fine-tune
Table 1: Comparison of CRL paradigms along three orthogonal axes.
Figure 1: Conceptual Overview of Task Specialization Fine-Tuning for Contextual RL. We illustrate the framework using a two-dimensional CartPole context space, where a pretrained policy is fine-tuned on three source task sets under different budget allocations to maximize final task coverage.
CVRP
CVRPTW
CartPole
Ant
N/K
3/100
3/150
4/150
3/100
3/150
4/150
3/50
3/100
4/100
3/50
3/100
4/100
Oracle
26.7%
29.3%
32.5%
49.4%
50.0%
50.0%
89.2%
89.7%
90.5%
100.0%
100.0%
100.0%
Oracle-Warmup
21.4%
24.5%
28.0%
46.6%
47.5%
47.5%
86.8%
87.5%
90.4%
100.0%
100.0%
100.0%
Pretrained
6.8%
6.8%
6.8%
40.9%
40.9%
40.9%
66.6%
66.6%
66.6%
31.1%
31.1%
31.1%
MTL
9.5%
11.0%
11.0%
44.8%
46.0%
46.0%
71.9%
74.5%
74.5%
31.1%
31.1%
31.1%
Random
10.1%
11.0%
12.9%
42.9%
43.2%
43.3%
75.4%
75.0%
82.4%
15.8%
42.5%
45.7%
Table 2: Performance Comparison in Combinatorial Optimization and Continuous Control.
Meta-World (10/50)
LLM-FT (4/100)
Oracle
86.0%
77.8%
Oracle-Warmup
84.0%
77.8%
Pretrained
58.0%
0.0%
MOORE
60.0%
/
Random
71.2%
26.7%
Uniform
70.0%
22.2%
Table 3: Results for Meta-World and LLM-FT.
Figure 2: From left to right: Performance comparison under diverse budget settings; Average recall for predicting coverage sets over the future 50 budget units, given varying amounts of data for model fitting; An example run illustrating the allocation policies and the global coverage trajectories over the training process across different methods for the case with N=3 and K=100 .
Appendix figures & tables17 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 3: Illustration of RH-DP.
Hyperparameter
Value
Parameter precision
BF16 (mixed precision)
Group size (rollouts per prompt)
8
KL loss coefficient
0 (KL loss disabled)
KL penalty in reward
Disabled
Advantage normalization by group std
Disabled
Entropy coefficient
0
Appendix
Table 4: GRPO hyperparameters used for LLM fine-tuning.
Figure 4: Examples of node distributions obtained by varying the covariance σ of a GMM.
Figure 5: Visualization of the context spaces and corresponding source task sets for CVRP, CVRPTW, CartPole, Ant, and Meta-World, respectively. The task embeddings in Meta-World are the K -dimensional expert weights associated with each task from the pretrained model. We use PCA to reduce them to two dimensions for visualization.
Figure 6: Coverage rate comparison on continuous control. Bars report mean coverage, and error bars denote standard deviation.
Trial 1
Trial 2
Trial 3
Trial 4
Trial 5
3/50
3/100
4/100
3/50
3/100
4/100
3/50
3/100
4/100
3/50
3/100
4/100
3/50
3/100
4/100
Oracle
86.5%
86.8%
90.0%
89.7%
91.2%
91.5%
89.9%
90.5%
90.6%
90.0%
90.0%
90.2%
89.7%
90.0%
90.1%
Oracle-Warmup
82.1%
82.4%
89.7%
83.7%
85.5%
91.5%
89.9%
90.5%
90.6%
88.5%
89.0%
90.0%
89.7%
89.9%
90.0%
Pretrained
53.5%
53.5%
53.5%
83.1%
83.1%
83.1%
68.2%
68.2%
68.2%
60.5%
60.5%
60.5%
67.6%
67.6%
67.6%
MTL
63.7%
73.0%
73.0%
83.1%
83.1%
83.1%
84.6%
88.4%
88.4%
60.5%
60.5%
60.5%
67.6%
67.6%
67.6%
Random
60.2%
65.7%
79.5%
78.4%
77.5%
90.2%
82.0%
87.6%
77.7%
75.6%
64.8%
76.9%
81.0%
79.4%
87.5%
Appendix
Table 5: Full Results for CartPole.
Ant Trial 1
Ant Trial 2
Ant Trial 3
Meta-World 5/25
3/50
3/100
4/100
3/50
3/100
4/100
3/50
3/100
4/100
Trial 1
Trial 2
Trial 3
Trial 4
Trial 5
Oracle
100.0%
100.0%
100.0%
100.0%
100.0%
100.0%
100.0%
100.0%
100.0%
68.0%
72.0%
70.0%
70.0%
68.0%
Oracle-Warmup
100.0%
100.0%
100.0%
100.0%
100.0%
100.0%
100.0%
100.0%
100.0%
62.0%
70.0%
68.0%
68.0%
68.0%
Pretrained
93.4%
93.4%
93.4%
0.0%
0.0%
0.0%
0.0%
0.0%
0.0%
48.0%
48.0%
52.0%
56.0%
42.0%
MTL
93.4%
93.4%
93.4%
0.0%
0.0%
0.0%
0.0%
0.0%
0.0%
62.0%
60.0%
64.0%
64.0%
54.0%
Random
21.6%
33.1%
60.2%
0.0%
74.2%
73.4%
25.8%
20.3%
3.5%
54.0%
58.0%
60.0%
66.0%
56.0%
Appendix
Table 6: Full Results for Ant and Meta-World.
CVRP
CVRPTW
CartPole
Ant
5 Points
30.2%
30.7%
28.2%
3.3%
10 Points
1.9%
0.8%
8.3%
0.0%
Appendix
Table 7: Approximate Monotonicity Rate.
CVRP 3/100
Ant 3/50
Type
Model
Coverage
Avg. ID MSE
Avg. OOD MSE
Type
Model
Coverage
Avg. ID MSE
Avg. OOD MSE
Monotonic
Exponential
18.7%
0.027682
0.103628
Monotonic
Exponential
88.4%
0.000694
0.001309
Power-Law
19.7%
0.065287
0.192049
Power-Law
88.4%
0.000702
0.001048
Logarithmic
19.0%
0.063002
0.179515
Logarithmic
90.7%
0.000760
0.001076
Non-monotonic
Quadratic
8.1%
0.027072
0.565207
Non-monotonic
Quadratic
85.2%
0.000625
0.012584
Cubic
7.1%
0.015860
0.585428
Cubic
89.7%
0.000628
1.969833
Appendix
Table 8: Comparison of different parametric models. We report the TSFT coverage and the average in-distribution (ID) and out-of-distribution (OOD) mean squared errors (MSEs).
Figure 7: Top left panel: Comparison of recall rates of coverage sets. Top right panels: Comparison of fitted models on a single task. Bottom panels: Predicted coverage sets from three modeling approaches (i.e., Ours, GP, and PFGP) after 50 additional budget units. Blue , red , and dark red points denote false positives (FP), false negatives (FN), and true positives (TP), respectively.
Figure 8: Left panel: Sensitivity test of performance threshold ϵ . Right panels: Sample efficiency comparison between TSFT and MTL on CVRP. We use the best TSFT results reported in Table 2 as thresholds, namely 187 and 264 for budgets of 100 and 150 units, respectively.
Figure 9: Left panel: Performance comparison under different number of policies. Middle panel: Sensitivity test of warmup budget. Right panel: Sensitivity test of execution budget.
CVRP 3/100
Adaptive
LinUCB
TSFT (DP)
(XS1,XS3,XS5)
10.4%
12.6%
18.7%
(XS1,XS2,XS3)
7.1%
8.1%
11.4%
(XS2,XS3,XS4)
6.4%
6.1%
12.3%
(XS3,XS4,XS5)
9.0%
5.7%
16.9%
(XS1,XS2,XS5)
10.5%
11.7%
21.8%
(XS1,XS4,XS5)
7.4%
7.3%
18.4%
Appendix
Table 9: Sensitivity Test of Source Task Set.
Setting
Method
Runtime / Memory
Varying the number of policies N with fixed budget K=100
K=100
N=2
N=3
N=4
N=5
N=6
N=7
TSFT (DP)
0.04s / 2.5MB
2.1s / 47.8MB
47.0s / 771.8MB
642.3s / 9476.2MB
–
–
TSFT (RH-DP)
0.02s / 1.0MB
0.4s / 8.4MB
3.3s / 65.2MB
25.2s / 406.4MB
136.8s / 2186.9MB
717.2s / 10418.9MB
TSFT (ILP)
0.08s / 2.2MB
0.2s / 6.1MB
0.5s / 8.7MB
0.5s / 9.0MB
0.5s / 10.3MB
0.4s / 11.0MB
Varying the total budget K with fixed number of policies N=3
Appendix
Table 10: Computational cost comparison of TSFT under different numbers of policies N and total budgets K . Each entry reports runtime and peak memory usage per decision-making step.
TSFT (DP)
CVRP 3/100
CVRP 3/150
Original
18.7%
23.0%
Reduced Evaluation (10x)
18.9%
22.3%
Appendix
Table 11: Results for Cheap Evaluation.
Figure 10: Top row: Performance heatmaps of policies selected by TSFT. Bottom row: Performance heatmaps and coverage sets (yellow regions) of the pretrained policy and the multi-task learning (MTL) policy under the same budget as TSFT, together with the final coverage set achieved by TSFT.
Single-task RL agents are typically trained under a fixed reward function, which limits their robustness to reward misspecification and their ability to adapt to changing preferences. We introduce Reward-Conditioned Reinforcement Learning (RCRL), an off-policy method that conditions agents on reward parameterizations while collecting experience under a single nominal objective. By recomputing counterfactual rewards from shared replay data, RCRL exposes the agent to multiple reward objectives without additional environment interaction, connecting single-task RL with ideas from multi-objective and multi-task learning. Across single-task, multi-task, and vision-based benchmarks, RCRL improves sample efficiency under the nominal reward parameterization, enables efficient adaptation to new parameterizations, and supports zero-shot behavioral adjustment at deployment. Our results show that RCRL provides a scalable mechanism for learning robust, steerable policies without sacrificing the simplicity of single-task training.
Michal Nauman, Marek Cygan, Pieter Abbeel
University of Warsaw · University of Warsaw, Nomagic · UC Berkeley, Amazon FAR
Offline-to-online (O2O) reinforcement learning aims to leverage policies pretrained on static datasets while improving them through online interaction. However, directly reusing an offline-trained critic can hinder online fine-tuning: as the policy and data distribution change rapidly, value estimates inherited from offline training may become misaligned with the online environment, leading to inaccurate policy improvement and inefficient exploration. To address this problem, we introduce Critic-Free Pretraining: an efficient paradigm that completely abandons the approach of offline critic training, allowing a freshly initialized critic to adapt without inheriting biased estimates. CFP is compatible with various mainstream O2O algorithms and consistently matches or improves upon conventional O2O algorithms across a diverse set of tasks, with particularly pronounced gains on several challenging tasks.
Reinforcement Learning from Verifiable Rewards (RLVR) suffers from exploration inefficiency, where models struggle to generate successful rollouts, resulting in minimal learning signal. This challenge is particularly severe for tasks that require the acquisition of novel reasoning patterns or domain-specific knowledge. To address this, we propose Context Bootstrapped Reinforcement Learning (CBRL), which augments RLVR training by stochastically prepending few-shot demonstrations to training prompts. The injection probability follows a curriculum that starts high to bootstrap early exploration, then anneals to zero so the model must ultimately succeed without assistance. This forces the policy to internalize reasoning patterns from the demonstrations rather than relying on them at test time. We validate CBRL across two model families and five Reasoning Gym tasks. Our results demonstrate that CBRL consistently improves success rate, provides better exploration efficiency, and is algorithm-agnostic. We further demonstrate CBRL's practical applicability on Q, a domain-specific programming language that diverges significantly from mainstream language conventions.
Saaket Agashe, Jayanth Srinivasa, Gaowen Liu +2
University of California, Santa Barbara · Cisco Research