Task Specialization Fine-Tuning for Contextual Reinforcement Learning

2026-08-17

AI Summary

本文提出任务专业化微调(TSFT)方法，通过预训练策略和在线框架解决上下文强化学习中的样本效率问题，提高任务覆盖率。

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.

Problem

• Contextual Reinforcement Learning
• fine-tuning
• sample efficiency
• task specialization
• budget allocation

Innovation

• Task Specialization Fine-Tuning
• Contextual Reinforcement Learning
• Integer Linear Programming