PaperlayerView exact revision on arXiv

Source-linked AI summary

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

arXiv:2608.17180v1cs.LGcs.AI

TL;DR

CRL lacks a principled answer to how much constrained fine-tuning budget each task region should receive after pretraining. TSFT predicts specialization performance, allocates budget via integer linear programming, and improves task coverage across diverse domains while approaching oracle performance.

  • Problem

    Existing CRL work largely fixes per-policy budgets, leaving how much fine-tuning each task region should receive under a constrained budget open.

  • Method

    TSFT uses an online parametric performance model with periodic re-estimation and solves each model-based budget allocation through integer linear programming.

  • Results

    Across combinatorial optimization, continuous control, and LLM fine-tuning, TSFT improves task coverage, achieves up to 2– 3× gains over simple strategies, and performs comparably to an oracle.

  • Takeaways & Limitations

    Budget-aware specialization offers a unified pretrain-finetune direction for improving task coverage in contextual reinforcement learning.

  • Takeaways & Limitations

    TSFT assumes fixed task grouping, and its surrogate objective, exhaustive evaluation, and exponential performance model may limit alignment, scalability, or applicability.

Abstract

from arXiv · show

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.

1 Introduction

The paper reframes contextual reinforcement learning as budget-constrained specialization of a pretrained policy, addressing heterogeneous fine-tuning returns across task regions. It introduces TSFT, which predicts performance gains and allocates fine-tuning budget exactly via integer linear programming, with validation across multiple decision domains.

  • Background: CRL targets broad coverage of related tasks represented as vectors in a unified context space, but independent, multi-task, and multi-policy paradigms each face coverage, capacity, transfer, or computational limitations.The introduction contrasts these existing approaches as the background for the proposed unified alternative.
  • Problem Setting: The pretrain-finetune paradigm first trains one policy across the context space, then produces multiple specialized variants under a constrained budget.This unified alternative combines aspects of multi-task and multi-policy training while retaining a shared pretrained initialization.
  • Motivation: Fine-tuning returns vary across task regions, with some improving rapidly while others saturate or degrade, making uniform compute allocation sample-inefficient.This motivates deciding how much fine-tuning each region should receive rather than only deciding where to train.
  • Method: TSFT formulates fine-tuning allocation as an online, model-based budget-allocation problem and solves each resulting maximum-coverage variant exactly with integer linear programming.A parametric model predicts how task performance and the induced coverage set evolve as additional budget is assigned.
  • Contributions: The paper validates TSFT across combinatorial optimization, continuous control, and LLM fine-tuning, alongside theoretical error analysis.These experiments assess the proposed budget-allocation framework across diverse decision domains.

2 Related Work

Prior work frames contextual reinforcement learning as structured multi-task learning challenged by negative transfer, diverse dynamics, sparse feedback, and unstable optimization. Alternatives use transfer, policy composition, multi-policy source selection, or pretrained representations to enable task specialization and adaptation.

  • Multi-Task Learning and Contextual RL: Multi-task learning improves generalization across related tasks but can suffer negative transfer when tasks are insufficiently related.In reinforcement learning, diverse task dynamics, sparse feedback, and unstable optimization exacerbate these challenges.
  • Multi-Task Learning and Contextual RL: Single-policy contextual RL approaches remain limited by model capacity and may experience greater negative transfer as task diversity increases, despite improved generalization from larger, better-regularized value functions [28].
  • Multi-Task Learning and Contextual RL: Transfer learning [14], policy composition [30], and multi-policy CRL methods mitigate task interference by enabling component specialization or selecting source tasks for zero-shot coverage of the remaining context space [16] [19].MBTL models source-task performance and transfer gaps to guide source-task selection.
  • RL Fine-Tuning: RL fine-tuning research includes learning adaptable initializations or task representations with limited new data, reusable behaviors through unsupervised skill discovery, and more direct pretraining followed by task-specific fine-tuning [31] [32] [33].Representative approaches include MAML and PEARL for adaptation and unsupervised skill discovery for reusable behaviors.

3 Problem Statement

The paper models contextual reinforcement learning as a family of context-specific MDPs and defines task coverage through a performance threshold. It formulates task specialization as allocating a constrained fine-tuning budget across source task sets to maximize the union of specialized policies’ coverage.

  • Contextual MDP: A contextual MDP is a finite, bounded family of MDPs indexed by context vectors that may alter transitions, rewards, and initial-state distributions [6] [8] [9].
  • Task Specialization Fine-Tuning: Given a pretrained policy and source task set, fine-tuning with k budget units produces a specialized policy whose performance is evaluated on each context task.
  • Task Specialization Fine-Tuning: A task is covered when the specialized policy meets a predefined performance threshold, with the inequality reversed for maximization objectives; the policy’s coverage set contains all such tasks.
  • Budget Allocation: With N source task sets and K total budget units, the objective allocates budget across specialized policies to maximize global coverage, defined as the union of their potentially overlapping coverage sets.
  • Budget Allocation: A budget unit counts consumed data samples over training epochs or gradient steps, making the allocation problem combinatorial and non-convex.

4 Methodology

TSFT formulates task-specialization budget allocation as a model-based multiple-choice coverage problem, solving it exactly via ILP while estimating unknown coverage through parametric performance models. It embeds this allocation procedure in an online framework with warmup, iterative re-estimation, and execution-budget control, and decomposes its optimality gap into warmup, modeling, and planning errors.

  • Online TSFT Framework: The framework solves the model-based allocation online at each decision step and periodically re-estimates performance models to mitigate error accumulation, enabling adaptive and sample-efficient task specialization.This procedure embeds exact ILP allocation into contextual reinforcement learning rather than planning only once over the full horizon.
  • ILP Formulation: TSFT formulates allocation as a multiple-choice budgeted maximum coverage problem and solves it exactly with an ILP over policy allocation levels and task coverage variables.The formulation selects exactly one allocation level per policy, respects the total budget, and uses coverage variables activated only when a task is covered.
  • Surrogate Performance Model: Because task coverage is unknown before training, TSFT fits a parametric performance model for each policy-task pair using nonlinear least squares to predict performance across budget levels.The fitting process can be performed in parallel over thousands of tasks within seconds, but modeling error can make the surrogate-optimal allocation differ from the true optimum.
  • Online TSFT Framework: TSFT begins with uniform warmup across policies, evaluates saved checkpoints to build training data, and repeatedly fits models, solves an ILP, allocates a bounded execution budget, and collects new evaluations.The execution budget is distributed proportionally to the ILP solution rather than executed to completion, limiting over-commitment to imperfect models while supporting balanced exploration.
  • Theoretical Analysis: Theorem 1 decomposes TSFT’s optimality gap from the Oracle into warmup error, surrogate-model error 2δm, and algorithmic planning error ηalg.Warmup error vanishes when the post-warmup state lies on an optimal trajectory, while ηalg = 0 for a one-shot exact full-horizon planner optimizing and executing a fixed surrogate exactly.

5 Experiment

Across combinatorial optimization, continuous control, and LLM fine-tuning, TSFT consistently improves task coverage over non-oracle baselines and approaches Oracle-Warmup. Additional analyses show strong sample efficiency, effective online modeling, modest computational overhead, and robustness across budgets.

  • Combinatorial Optimization: In combinatorial optimization, TSFT’s gains over the strongest baseline increase with budget or policy count on CVRP, while remaining more moderate on CVRPTW.The comparison is reported in Table 2, with coverage rate defined as global coverage divided by the total number of context-space tasks.
  • Results: TSFT achieves the best non-oracle coverage across combinatorial optimization, continuous control, and LLM fine-tuning, consistently approaching Oracle-Warmup.It performs best across all six CVRP and CVRPTW settings, improves coverage in CartPole, Ant, and Meta-World, and outperforms non-oracle LLM baselines.
  • Continuous Control: In continuous control, TSFT consistently improves coverage and narrows the Oracle-Warmup gap as budget grows, while MTL can be unstable on Ant but remains strong on Meta-World.On Meta-World, MOORE reaches an average success rate of 0.68, yet TSFT still achieves better coverage.
  • LLM Fine-Tuning: TSFT consistently outperforms all non-oracle LLM baselines and approaches Oracle-Warmup, whereas heuristic allocation recovers only part of the specialization benefit.The results indicate that model-based budget allocation scales to LLM reinforcement-learning fine-tuning with billions of parameters.
  • Analysis: TSFT’s surrogate modeling and ILP allocation add modest overhead: thousands of models fit in parallel within seconds, and an N = 10, K = 150 ILP solves optimally in under one second.Online model re-estimation becomes effective as additional data arrive, although coverage can initially shrink after pretraining because of distribution shift.

6 Conclusion · Appendix · A Discussion

The paper presents TSFT as an online, model-based framework that allocates constrained fine-tuning budgets across task regions using a parametric model and exact integer linear programming. It also identifies scalability, surrogate-objective alignment, pretrained-policy quality, and unstructured context relationships as important limitations and directions for extension.

  • 6 Conclusion: TSFT allocates a constrained fine-tuning budget across task regions through a simple parametric model and solves the resulting allocation problem exactly with integer linear programming.The framework embeds this allocation process into an online procedure.
  • 6 Conclusion: The framework’s limitations include possible misalignment between its surrogate and original objectives, which can make surrogate-optimal allocations perform poorly on the original objective.This is identified as the first of four limitations.
  • 6 Conclusion: Exhaustive evaluation may become computationally prohibitive as the context space or number of specialized policies grows, limiting TSFT’s scalability.This is identified as the second of four limitations.
  • A Discussion: At the formulation level, Eqs. (3)-(7) define a multiple-choice budgeted MCP, where each policy selects one of K + 1 allocation levels with corresponding budget consumption and coverage set.The formulation groups each policy’s candidate allocation levels and requires exactly one level per policy.
  • A Discussion: TSFT assumes a pretrained policy with meaningful capabilities for adaptation; without relevant capabilities, fine-tuning approaches learning target tasks from scratch and may require substantially more effort.The assumption reflects the intended setting of adapting an already capable pretrained model, such as a foundation model.
  • A Discussion: When context relationships are less explicit or smoothly varying, source-task grouping methods can complement TSFT’s focus on how much to train.Task embeddings and gradient-based task-affinity grouping are proposed as orthogonal directions for future work.

B Methodology Detail … B.3 Alternative Model

The methodology formulates budget allocation as a model-based MDP, predicts policy-task performance, and optimizes coverage-aware allocations with dynamic programming or ILP. It also compares alternative formulations and performance models, balancing expressiveness against computational tractability.

  • B.1.1 MDP Formulation: The MDP allocates discrete budget units across policies, with deterministic allocation transitions and early stopping when further training yields no improvement.The underlying training and realized performance remain stochastic despite deterministic transitions in the abstract allocation state.
  • B.1.1 MDP Formulation: A parametric policy-task performance model predicts fine-tuning outcomes from consumed budget, while nonlinear least squares fitting scales to thousands of tasks within seconds.The model allows performance shifts in either direction, capturing both improvement and degradation during fine-tuning.
  • B.1.1 MDP Formulation: Coverage reward is the marginal increase in the union of policies’ task coverage, coupling decisions across policies and penalizing redundant specialization.The global objective therefore incentivizes expanding collective coverage rather than optimizing policies independently.
  • B.1.2 MDP Solving: Dynamic programming optimizes the surrogate coverage objective induced by predicted performance and is globally optimal for the full budget, but its complexity is exponential in policies N.Receding-horizon DP trades global optimality for tractability by using limited lookahead and execution horizons.
  • B.1.3 Dynamic Programming: The DP implementation recursively evaluates allocation states, selects the action with maximum future value, and returns no action when further training cannot improve coverage.Memoization supports repeated state evaluation, while termination implements early stopping.
  • B.2 Alternative Formulation: Alternative ILP encodings reduce variables under monotone coverage, whereas the original ILP and CP formulation are more general; the original ILP benefits from mature mixed-integer optimization techniques.The CP formulation supports arbitrary non-monotone coverage but can become expensive for large context spaces.
  • B.3 Alternative Model: Gaussian processes provide non-parametric posterior-mean performance predictions for each policy-task pair, implicitly estimating how coverage evolves with training.PFGP combines a smooth parametric trend with a GP residual model to capture local deviations and stochastic fluctuations.
  • B.3 Alternative Model: Additional alternatives include power-law, logarithmic, polynomial, and piecewise-linear models, spanning smooth monotonic trends and flexible non-monotonic trajectories.Polynomial degrees evaluated are 2, 3, and 4, while the piecewise-linear breakpoint is selected using BIC.

C Theoretical Analysis · D Experiment Detail · D.1 Baseline

The theoretical analysis bounds TSFT’s post-warmup gap through surrogate error and planning accuracy, then decomposes its full-Oracle gap into warmup and post-warmup losses. Experiment details specify hardware across three domains, while baseline implementations define Adaptive and LinUCB allocation and reward procedures.

  • C Theoretical Analysis: TSFT’s post-warmup optimality gap relative to Oracle-Warmup is bounded by 2δm + ηalg under bounded surrogate error and surrogate planning accuracy.For exact full-horizon planning, ηalg = 0, reducing the bound to 2δm.
  • C Theoretical Analysis: Because coverage is thresholded, δm measures induced coverage-level surrogate error, and online model re-estimation may reduce it without guaranteeing monotonic improvement.Small performance-prediction errors can change coverage near the threshold.
  • C Theoretical Analysis: The full-Oracle gap decomposes into warmup error plus the post-warmup term bounded by 2δm + ηalg.Warmup error can be zero when the post-warmup state lies on an optimal trajectory.
  • C Theoretical Analysis: Surrogate prediction error need not cause allocation error when high-quality feasible allocations retain their ranking, but near-ties or threshold-adjacent tasks can change selection.Thus, allocation quality depends on ranking preservation rather than exact equality between surrogate and true coverage.
  • D Experiment Detail: Experiments use RTX 4090 and Threadripper PRO 7975WX hardware for combinatorial optimization, V100 and Xeon E5-2670 hardware for control, and 16 H200 GPUs for LLM fine-tuning.The hardware configurations differ by experimental domain.
  • D.1 Baseline: Adaptive maintains uniformly initialized policy weights, converts them through temperature-scaled softmax probabilities, and samples execution-budget allocations from the resulting multinomial distribution.Updated checkpoints are evaluated after training, as described in the baseline implementation.
  • D.1 Baseline: LinUCB uses α = 1.0 for exploration, selects the policy with the largest upper-confidence score, and allocates one budget unit at each decision step.Its reward is normalized monotonic coverage improvement over the selected policy’s historical best, avoiding penalties for temporary fine-tuning drops.

D.2 Policy Training · D.3 Context Space · E Additional Result

The paper instantiates TSFT across combinatorial optimization, continuous control, and LLM fine-tuning, using domain-specific policy-training protocols and context spaces defined by task parameters or learned representations.

  • D.2 Policy Training: POMO policies are pretrained for 5,000 epochs across the full combinatorial-optimization context space using 10,000 instances per epoch.The setup uses Adam with learning rate 1 × 10−4, weight decay 1 × 10−6, batch size 64, and problem and POMO sizes of 100.
  • D.2 Policy Training: Continuous-control experiments use PPO through Stable Baselines3 with shared MLP policy-value backbones sized for CartPole and [256, 256] for Ant.The experiments use default PPO hyperparameters, including learning rate 3 × 10−4, n_steps=2048, batch size 64, discount factor 0.99, and GAE parameter 0.95.
  • D.2 Policy Training: LLM specialization fine-tunes Qwen3-4B-Base with GRPO in VeRL and vLLM, keeping the configuration fixed across four policies while changing only the training corpus.After every GRPO update, each checkpoint is evaluated on all context-space benchmarks, with evaluation capped at 100 problems per benchmark.
  • D.3 Context Space: CVRP contexts vary vehicle capacity from 10 to 400 in steps of 10 and control node distributions through covariance changes in a Gaussian mixture model.Customer demands are sampled uniformly from {1, 2, . . . , 9} and normalized by vehicle capacity before network input.
  • D.3 Context Space: CVRPTW extends CVRP with customer time windows and service times, varying capacity from 10 to 400 and service time from 0.04 to 1.00 in steps of 0.04.Larger service times correspond to tighter time windows, and five source task sets each contain 49 tasks centered around reference tasks.
  • D.3 Context Space: Continuous-control context spaces discretize CartPole across pole length, cart mass, and pole mass, and Ant across gravity and friction.Ant uses a 25×40 = 1,000-context grid, with five 20-context source rectangles spanning 100 grid points; CartPole uses five axis-aligned 27-context regions.
  • D.3 Context Space: Meta-World represents 50 manipulation tasks using pretrained expert weights rather than one-hot encodings, projecting those learned task representations into two dimensions with PCA.The LLM context space contains nine benchmarks spanning math reasoning, code generation, and general reasoning, each with a benchmark-specific performance threshold.

E.1 Full Result

Full results are reported for CartPole, Ant, and Meta-World, with continuous-control coverage compared in Figure 6. Experiments use multiple seeds for CartPole and Ant, while other domains generally use single-seed evaluation due to stability or computational cost.

  • E.1 Full Result: Experiments use 5 seeds for CartPole and 3 for Ant, while combinatorial optimization and LLM fine-tuning use one seed, and Meta-World uses one main-experiment seed plus a smaller five-seed robustness evaluation.The Meta-World robustness setting uses N = 5 and K = 25.
  • E.1 Full Result: Coverage rates on continuous control are compared using mean bars and standard-deviation error bars in Figure 6.
  • E.1 Full Result: Full results cover CartPole, Ant, and Meta-World, with CartPole reported in Table 5 and Ant and Meta-World in Table 6.

E.2 A Comprehensive Study of Performance Model

TSFT uses a simple exponential performance model because effective budget allocation requires only coarse learning-curve trends, not highly accurate trajectory prediction. Evidence supports this monotonic inductive bias despite noisy training curves, while acknowledging limitations for strongly non-monotonic behavior.

  • Model Motivation: TSFT assumes approximately monotonic fine-tuning trajectories and uses an exponential model to capture coarse learning-curve trends for budget allocation.Accurate prediction of the entire stochastic training trajectory is unnecessary for effective allocation.
  • Robustness to Noise: Nearly all evaluated trajectories show local non-monotonic behavior when 10 points are sampled, yet the monotonic model remains effective on noisy training curves.The study quantifies trajectory monotonicity by uniformly sampling 5 or 10 points from each curve (Table 7).
  • Model Comparison: Monotonic Exponential, Power-Law, and Logarithmic models are compared with Quadratic, Cubic, Quartic, and Piecewise models on CVRP and Ant using fitting, prediction, and coverage metrics.The comparison tests whether monotonicity reduces overfitting under limited observations across relatively monotonic and more challenging non-monotonic domains.
  • Limitations: The exponential model is not universally applicable and may allocate budgets suboptimally when fine-tuning trajectories exhibit strongly non-monotonic global behavior.Future directions include uncertainty-aware exploration and more expressive surrogate models such as neural networks.

E.3 Sensitivity Analysis · E.4 Computational Cost

TSFT remains robust across performance thresholds, execution budgets, and warmup settings, while scaling better than baselines as the policy count increases. Its ILP-based allocation is computationally efficient and exact, unlike DP approaches that become impractical for larger policy sets.

  • E.3 Sensitivity Analysis: The sensitivity study tunes ϵ, N, W, and E through grid search over predefined ranges selected for computational feasibility and empirical robustness.These hyperparameters respectively control the performance threshold, number of policies, warmup budget, and execution budget.
  • E.3 Sensitivity Analysis: TSFT remains competitive and stable across performance thresholds ϵ from 1.15% to 1.4%, consistently outperforming standard baselines as coverage rises with looser thresholds.The analysis uses the 3/100 CVRP setting and increments of 0.05%.
  • E.3 Sensitivity Analysis: TSFT scales substantially better than baselines as the number of policies increases, whereas most baselines degrade under the fixed budget K = 150.LinUCB is mentioned as an exception, but the supplied passage is truncated before its full comparison.
  • E.3 Sensitivity Analysis: TSFT generally performs best with relatively small warmup budgets, while excessive warmup can reduce adaptive-specialization budget and degrade global coverage.The warmup budget varies from 3 to 15 per policy in the 3/100 and 3/150 CVRP settings.
  • E.3 Sensitivity Analysis: TSFT remains relatively stable as the execution budget unit varies from 5 to 30, indicating limited sensitivity once a reasonable execution budget is used.The result suggests adaptive allocation can benefit robustly without extensive tuning of E.
  • E.4 Computational Cost: ILP solves the allocation problem exactly while remaining computationally efficient, whereas exact DP scales poorly with policy count and RH-DP remains prohibitive for large policy sets.ILP benefits from compact formulation, LP relaxations, and branch-and-bound pruning, unlike DP’s near-complete enumeration and caching of reachable allocation states.

E.5 Cheap Evaluation … G License

TSFT enables complementary policy specialization across task-space subregions, while keeping evaluation cheaper than training and supporting greener adaptation practices. The paper also documents asset licensing and plans MIT-licensed release of its code and datasets.

  • E.5 Cheap Evaluation: TSFT periodically evaluates policies, with total evaluation taking hours rather than the days required for policy training.Exhaustive evaluation may still be impractical when evaluations are expensive.
  • E.5 Cheap Evaluation: Online evaluation cost can be reduced by using fewer validation samples or evaluating less frequently with a larger execution interval.These adjustments reduce the evaluation burden for fitting TSFT’s performance model.
  • E.6 Visualization: TSFT allocates different fine-tuning budgets across policies, leaving some at early stages while others progress to later stages.The visualization uses final policies from the 5/150 CVRP setting, where TSFT guides specialization.
  • E.6 Visualization: TSFT policies achieve low optimality gaps near their source-task regions, degrade farther away, and form complementary specialized coverage across the context space.Figure 10 compares TSFT-selected policies with pretrained and multi-task policies under the same budget, alongside TSFT’s final coverage set.
  • F Broader Impact: By improving sample efficiency, task-specialization fine-tuning may reduce the computational and energy costs of adapting RL agents across diverse task conditions.The paper identifies greener AI practices and improved accessibility for researchers and practitioners with limited resources as positive societal impacts.
  • G License: Existing assets’ licenses and usage are listed in Table 12, while the source code and datasets will be publicly released under the MIT License upon publication.Table 12 documents the licenses and usage of existing assets.
Loading 2608.17180v1…