Task Specialization Fine-Tuning for Contextual Reinforcement Learning

Jianan Zhou; Jung-Hoon Cho; Tianyue Zhou; Han Zheng; Jie Zhang; Roy Dong; Yining Ma; Cathy Wu

cs.LGAug 17, 2026

Task Specialization Fine-Tuning for Contextual Reinforcement Learning

Authors: Jianan Zhou, Jung-Hoon Cho, Tianyue Zhou, Han Zheng, Jie Zhang, Roy Dong, Yining Ma, Cathy Wu

Organizations: Nanyang Technological University · MIT · UIUC

Abstract

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 $\mathcal{N}$	Single-task	From scratch
Multi-Task Training	Single	Multi-task	From scratch
Multi-Policy Training	Multiple $N\ll\mathcal{N}$	Single-task	From scratch
TSFT (Ours)	Multiple $N\ll\mathcal{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

\sigma

of a GMM.

Panel of Figure 5 — 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).

Panel of Figure 7 — 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

\epsilon

. 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)
$(X_{S_{1}},X_{S_{3}},X_{S_{5}})$	10.4%	12.6%	18.7%
$(X_{S_{1}},X_{S_{2}},X_{S_{3}})$	7.1%	8.1%	11.4%
$(X_{S_{2}},X_{S_{3}},X_{S_{4}})$	6.4%	6.1%	12.3%
$(X_{S_{3}},X_{S_{4}},X_{S_{5}})$	9.0%	5.7%	16.9%
$(X_{S_{1}},X_{S_{2}},X_{S_{5}})$	10.5%	11.7%	21.8%
$(X_{S_{1}},X_{S_{4}},X_{S_{5}})$	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.

Panel of Figure 10 — 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.

Type	Asset	License	Usage
Code	POMO [ 41 ]	MIT License	Policy network (CVRP/CVRPTW)
	Stable Baselines3 [ 61 ]	MIT License	PPO algorithm (CartPole/Ant)
	MOORE [ 45 ]	MIT License	Meta-World multi-task RL
	VeRL	Apache 2.0	LLM RFT framework
Model	Qwen3-4B-Base [ 54 ]	Apache 2.0	LLM pretrained policy
Datasets	DAPO-17K [ 46 ]	Apache 2.0	LLM training / evaluation

Appendix

Table 12: Licenses and usage for existing assets.

Contextual RL

See all 19 papers

Reward-Conditioned Reinforcement Learning
Mar 5, 2026Michal Nauman, Marek Cygan, Pieter AbbeelReinforcement Learning Off-Policy RL
University of Warsaw · University of Warsaw, Nomagic · UC Berkeley, Amazon FAR
Critic-Free Pretraining for Efficient Online Reinforcement Learning Fine-Tuning
Aug 11, 2026Daoyi Li, Yixian Zhang, Wenbo Ding +2Reinforcement Learning Critic-Free RL
Tsinghua University
Context Bootstrapped Reinforcement Learning
Mar 19, 2026Saaket Agashe, Jayanth Srinivasa, Gaowen Liu +2RL for Language Model Reasoning Reinforcement Learning with Verifiable Rewards
University of California, Santa Barbara · Cisco Research

Mar 5, 2026J/K move · Enter open · S save

Reward-Conditioned Reinforcement Learning

Michal Nauman, Marek Cygan, Pieter Abbeel

University of Warsaw · University of Warsaw, Nomagic · UC Berkeley, Amazon FAR

Reinforcement Learning Off-Policy RL Multi-Objective Reinforcement Learning

Open paper AskPDF arXiv

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.

Abstract

Contextual RL

Explore similar work

Reward-Conditioned Reinforcement Learning

Critic-Free Pretraining for Efficient Online Reinforcement Learning Fine-Tuning

Context Bootstrapped Reinforcement Learning