面向上下文强化学习的任务特化微调

# 任务特化微调在上下文强化学习中的应用  
来源：https://arxiv.org/html/2608.17180  
作者：Jung-Hoon Cho（MIT）、Tianyue Zhou（MIT）、Han Zheng（MIT）、Jie Zhang（南洋理工大学）、Roy Dong（伊利诺伊大学厄巴纳-香槟分校）  
邮箱：[email protected], [email protected], [email protected], {jhooncho, tianyuez, hanzheng, yiningma, cathywu}@mit.edu  
Yining Ma（通讯作者，MIT）、Cathy Wu（MIT）  

###### 摘要  
上下文强化学习旨在通过最大化任务在上下文空间中的覆盖度，将经典强化学习泛化至相关任务族。现有方法通常从头训练，依赖单一策略的多任务学习或策略性地训练多个策略。我们提出一种统一替代方案：先预训练一个具有良好初始性能的单一策略，再针对任务特化微调多个策略。然而，这一新范式引入了异质边际回报与样本效率低下等挑战。由此产生关键研究问题：**给定预训练策略与受限预算，每个任务区域应分配多少微调资源以实现样本高效的上下文强化学习？**  
为此，我们提出**任务特化微调**——一种在线框架，通过简单参数化模型预测微调性能，并借助整数线性规划精确求解离散预算分配问题。在组合优化、连续控制和大语言模型微调等多领域的实验表明，任务特化微调在任务覆盖度上显著优于基线方法，并接近理论最优性能。本工作为基于模型的上下文强化学习开辟了新方向，契合当代预训练-微调范式。  

## 1 引言  
强化学习在多领域已取得显著成功，但在处理仅关键环境参数不同的相关任务族时仍显脆弱。上下文强化学习将此类任务族形式化为上下文马尔可夫决策过程，其中各任务由统一上下文空间中的向量参数化，目标是在该空间中实现广泛覆盖。现有上下文马尔可夫决策过程求解范式可分为三类：  
1) **独立训练**：为每个任务学习独立策略，虽直接但对高维上下文空间覆盖成本高昂；  
2) **多任务训练**：学习跨任务的统一策略，但受限于模型容量且易受负迁移影响；  
3) **多策略训练**：针对单任务训练多个策略，依赖其泛化能力覆盖未见任务。  
尽管存在差异，这些范式均假设每个策略从头训练至收敛。这与现代机器学习的预训练-微调惯例相悖，限制了上下文强化学习在高维扩展空间中的可扩展性。  

本文主张统一替代方案：预训练一个在上下文空间中具有优异初始性能的单一策略，再在受限预算下微调出多个特化策略（图1）。该范式统一了多任务与多策略训练，具有三方面优势：预训练策略提供良好初始化；微调实现多策略特化，避免负迁移与从头训练的计算开销；将上下文强化学习融入现代机器学习的可扩展框架。  

**表1**：上下文强化学习范式对比  
| 范式          | 策略数量 | 训练方法     | 训练模式       |  
|---------------|----------|--------------|----------------|  
| 独立训练      | 多       | 单任务       | 从头训练       |  
| 多任务训练    | 单       | 多任务       | 从头训练       |  
| 多策略训练    | 多       | 单任务       | 从头训练       |  
| 任务特化微调  | 多       | 多任务       | 预训练+微调    |  

然而，该范式引入独特挑战：微调增益具有异质性（部分区域提升迅速，部分区域快速饱和甚至性能下降），均匀分配微调计算资源通常导致样本效率低下。由此产生关键问题：**给定预训练策略与受限预算，每个任务区域应分配多少微调资源以实现样本高效的上下文强化学习？**  
现有文献主要研究**训练位置选择**（在固定预算假设下选择源任务以最大化覆盖），而正交的**资源分配量问题**仍未解决。本文引入预算感知视角，提出**任务特化微调**框架，用于上下文强化学习的基于模型的预算分配。该框架通过参数化模型预测任务性能随资源增加的变化趋势，并将每个规划步骤转化为基于分配向量的最大覆盖问题变体，可通过整数线性规划精确求解。为缓解模型误差影响，我们将整数线性规划嵌入在线框架，周期性重估模型：每轮求解当前规划问题，执行有限步长的分配策略，收集新数据并更新模型。  

**贡献**：  
1) 概念层面：在上下文强化学习中提出**预算约束微调问题**，将研究重点从训练位置选择转向预训练-微调范式下的资源分配；  
2) 方法层面：提出**任务特化微调框架**，通过整数线性规划优化基于模型的预算分配，并提供理论误差分析；  
3) 实验层面：在组合优化、连续控制和大语言模型微调中验证任务特化微调的有效性，任务覆盖度较简单策略提升2–3倍，样本效率较多任务基线提升2倍，且在不同设置下接近理论最优策略性能。  

## 2 相关工作  
### 多任务学习与上下文强化学习  
多任务学习通过联合学习提升泛化能力，但当任务相关性不足时易产生负迁移。在强化学习中，任务多样性、稀疏反馈和不稳定优化加剧了这一挑战。上下文强化学习可视为结构化的多任务强化学习，其任务族通过影响环境动力学、奖励或初始状态分布的上下文变量参数化。  

常见方法包括通过共享任务表征、上下文条件策略或任务条件策略/价值头训练跨任务泛化的单一策略。尽管研究表明更大、正则化更好的价值函数可提升泛化能力，单一策略方法仍受模型容量限制，且可能随任务多样性增加而产生负迁移。另一类方法通过迁移学习或策略组合缓解任务干扰，允许不同组件特化。  

近期多策略上下文强化学习方法通过战略性选择源任务并依赖零样本转移覆盖剩余上下文空间。例如，MBTL通过建模源任务性能与转移间隙指导选择，SD-MBTL进一步检测上下文马尔可夫决策过程的泛化结构并切换选择策略。这些方法主要解决**训练位置**问题，通常假设每个策略在固定预算下训练。本文延续多策略范式以缓解负迁移，同时解决其预算分配机制的缺失。  

### 强化学习微调  
微调已成为重用强化学习经验的常见方式。早期工作聚焦于学习可适应少量新数据的初始化或任务表征（如MAML、PEARL）。另一类方法在下游训练前学习可复用行为（如通过无监督技能发现）。  

近期工作更直接地研究预训练后微调：多任务预训练配合任务特定微调被证明是元强化学习的强基线替代方案；在Atari变体上的多任务预训练可提升向未见变体的泛化能力；自监督多任务预训练也被探索用于序列决策模型，提升跨任务的下游微调效率。离线到在线强化学习方法利用先验数据加速在线改进；参数高效方法（如L2M）可减少预训练模型适应新任务时的遗忘。  

这些工作主要解决如何获取有效初始化、稳定单次微调或适应单一智能体至新任务。相比之下，任务特化微调将微调进展本身作为规划信号：建模特化带来的异质边际收益，并协调跨微调策略组合的计算资源分配。  

## 3 问题定义  
### 上下文马尔可夫决策过程  
设标准马尔可夫决策过程为（状态空间、动作空间、转移动力学、奖励函数、初始状态分布）。上下文马尔可夫决策过程是一族上下文特定马尔可夫决策过程，由有限有界上下文集中的上下文向量参数化，该向量可影响转移动力学、奖励函数和初始状态分布。后文简化表示为特定任务。  

### 任务特化微调  
给定参数为θ的预训练策略，可在源任务集上分配k单位预算进行微调，实现上下文子空间上的任务特化。策略在任务上的性能表示为J(θ,x)。若性能满足预设阈值（最小化目标下J(θ,x)≤ε），则称任务被该策略覆盖。策略的覆盖集定义为所有被覆盖任务的并集。此约定无损失一般性：对于最大化目标，仅需反转不等式方向。  

传统研究通常假设训练至收敛，导致样本效率低下。本文考虑更实际的预算约束场景，将任务特化形式化为预算分配问题：给定N个不同源任务集与总预算K单位，目标是将受限预算分配至上下文空间以最大化全局覆盖度（即覆盖集的并集）。从预训练策略θ出发，我们在每个源任务集上独立微调，得到N个特化策略，分配目标形式化为：  
max{ K | ... }