随机世界元训练的大规模情境强化学习

Towards Large-Scale In-Context Reinforcement Learning by Meta-Training in Randomized Worlds 基于随机世界元训练的大规模情境强化学习

https://arxiv.org/pdf/2502.02869

摘要

上下文内强化学习（In-Context Reinforcement Learning, ICRL）使智能体能够从其交互经验中自动、实时地进行学习。然而，扩展ICRL面临的一个主要挑战是缺乏可扩展的任务集合。为解决这一问题，我们提出了程序化生成的表格型马尔可夫决策过程（Markov Decision Processes），命名为AnyMDP。通过精心设计的随机化过程，AnyMDP能够大规模生成高质量任务，同时保持相对较低的结构偏差。为了实现高效的大规模元训练，我们进一步在ICRL框架中引入了逐步骤监督和先验信息引导。实验结果表明，在AnyMDP任务达到足够大规模的情况下，所提出的模型能够泛化到训练集之外未见过的任务。AnyMDP提供的可扩展任务集还使得我们能够更深入地实证研究数据分布与ICRL性能之间的关系。我们进一步发现，ICRL的泛化能力可能需要以更高的任务多样性和更长的适应周期为代价。这一发现对扩展鲁棒的ICRL能力具有重要意义，强调了任务设计应注重多样性与广泛性，并优先考虑渐近性能而非少样本适应能力。

1 引言

大规模预训练已在自然语言、图像和视频处理等领域取得了巨大成功 [1–4]。这些模型已展现出通过上下文内学习（In-Context Learning, ICL）解决未见任务的能力，这是一种利用上下文信息提升性能的学习范式。与依赖基于梯度的模型权重更新的“权重内学习”（In-Weight Learning, IWL）不同，ICL使模型能够以少样本方式获取新技能，从而增强其在新环境中的适应能力。ICL与基于模型的元学习方法 [6, 7] 具有共通之处，其框架可容纳多种传统学习范式，包括监督学习 [7, 8]、模仿学习 [9–11] 和强化学习 [12–15]。这显著减少了对繁琐的人工设计目标函数和优化策略的依赖，而这些在IWL中通常是必需的。此外，基于梯度的IWL因其在持续适应新任务方面的低效性而受到批评 [16]。相比之下，ICL展现出类似人类大脑适应性的可塑性 [17]。

一方面，基于海量未筛选数据的预训练在ICL的可控性方面存在局限。ICL能力的涌现机制及其局限性尚未被充分理解 [18]。此外，预训练中ICL的鲁棒性也受到质疑 [19]。另一方面，元学习的目标是发展上下文内学习功能，而非掌握特定技能。然而，这些方法的有效性常常受限于大规模、多样化数据集的缺乏以及可用任务范围的狭窄——这一限制在元强化学习（MetaRL）和上下文内强化学习（ICRL）中尤为突出。因此，大多数现有研究仅在定义狭窄的领域内评估模型的适应能力，从而限制了这些方法的广泛泛化潜力。

为进一歩提升ICRL的可扩展性和通用性，我们将广泛使用的多臂老虎机基准 [20] 扩展至表格型马尔可夫决策过程（MDPs）。我们提出AnyMDP，一种可扩展的任务生成框架，其中MDP的转移函数和奖励函数均被完全随机化设计。为了生成非平凡且结构偏差最小的环境，我们提出一种程序化生成方法，合成多样化任务，为学习者带来显著挑战。为高效扩展ICRL，我们进一步在建模与训练框架中引入两项关键创新：逐步骤监督和先验知识引导。这使我们能够在前所未有的规模上进行元训练——超过60亿个时间步，上下文长度超过512K。该模型被称为OmniRL，完全在AnyMDP任务上训练，并展现出对完全未见过的Gymnasium环境的强泛化能力。这验证了该框架的有效性。利用这一可扩展任务集，我们进一步研究了数据分布如何在大规模下影响ICRL性能。实验揭示了三个关键洞见：（1）数据覆盖范围和广泛性对有效元学习至关重要，这些因素划定了以IWL为主导的任务识别范式与以ICL为主导的通用学习范式之间的界限；（2）渐近性能（即长期适应能力）可能是比零样本或少样本性能指标更可靠的ICRL能力衡量标准；（3）优先考虑任务多样性与广泛性的训练策略具有巨大潜力，其重要性可能超过对过往经验的保真度。

总结而言，我们的贡献包括：（1）提出AnyMDP，一种可扩展的表格型MDP生成框架，旨在支持在超过10万个任务、数十亿次模拟步长上训练ICRL。此外，我们提出了一种高效的ICRL训练与模型框架，确保其可扩展性。（2）我们研究了数据分布如何影响ICRL，并通过实验证实了多样化和广泛性任务设计的必要性，以及评估渐近性能的重要性。

2 相关工作

2.1 上下文学习的兴起

元学习（meta-learning），也被称为“学会学习”[21]，指的是一类方法，其重点在于获取可在多种任务之间泛化和适应的能力。它涵盖广泛的方法论，包括基于梯度的优化方法[22]和基于模型的元学习方法[6]。元学习的概念与“系统3”的形成相吻合[23]，即通过突触重校准来适应环境，从而补充大语言模型（LLMs）的“系统1”（直觉思维）和“系统2”（上下文推理）。使用大规模、未经筛选的数据集进行预训练的大型模型，也以类似于元学习的方式促进了上下文学习（ICL）的出现[5, 9, 24–28]。为了明确起见，“预训练”指的是以技能获取为目标的训练过程，通常随后需要基于梯度的微调；而“元训练”则指以获得学习能力为目标的训练，无需后续基于梯度的调整。已有大量理论研究[19, 29–31]致力于严谨地揭示ICL能力与预训练数据分布之间的关联。实证证据表明，突发性（burstiness）显著提升了ICL在语言建模任务[30]和决策问题[26]中的表现。当前的ICL实现主要作为少样本学习器运作[32, 33]，尤其是在任务多样性不足的情况下。已有分析和实验表明，基于计算的ICL可能展现出比基于梯度的IWL更丰富的行为模式[34–36]，特别是在可塑性和持续学习方面[17]，这使得ICL有可能成为一种广义的学习机制。然而，最近的研究表明，模型在可行的情况下倾向于优先依赖IWL而非ICL[31]。我们的研究与日益增长的通用上下文学习研究[32, 37, 38]保持一致，这些研究强调将ICL作为技能获取的主要机制，而非IWL。

2.2 上下文强化学习

上下文强化学习（ICRL）包含一系列能够通过综合自生成的轨迹并结合外部反馈，动态适应上下文信息的算法[6, 12]。它通常采用循环神经网络结构[6]和注意力结构[12, 20]，这些结构能够在内循环中编码历史交互信息。用于优化内循环学习机制参数的训练过程称为外循环（outer loop）。ICRL中外循环优化器的常见选择包括强化学习[6, 13, 20]、进化策略[39, 40]和监督学习[12, 15]。尽管监督学习通常比强化学习和进化策略具有更高的样本效率，但它往往具有较弱的遗憾界（regret bounds）[41, 42]。此外，缺乏一个“oracle策略”（监督学习的关键前提）可能构成一个关键瓶颈。在这种情况下，诸如RL Coaches[9, 43]之类的框架提供了可行的替代方案，弥合了监督学习范式与强化学习之间的差距。上下文强化学习（ICRL）研究中的另一个重大挑战源于大规模任务集的稀缺性，因为任务的多样性和规模对ICL能力至关重要[44]。现有研究主要采用领域随机化（domain randomization）技术，将标准强化学习环境扩展为更广泛的关联环境类别。近年来，构建能在开放世界中通用适应的智能体引起了广泛关注[37, 45–48]。然而，现有环境存在显著的结构偏差，这阻碍了在某一任务类别上训练的ICRL模型直接适用于其他任务类别。

2.3 领域随机化

领域随机化（Domain Randomization, DR）最初被提出是为了改善从仿真到现实的迁移效果[49]。在元学习中，也常通过对部分领域参数进行随机化，从而生成各种变体任务。这些基准通常可分为以下几类：在保持状态转移不变的情况下随机化奖励或目标[12, 15, 20, 22, 50, 51]；在保持目标不变的情况下随机化环境动力学[39, 40, 52–54]；以及在不改变底层状态转移和奖励的情况下随机化观测值和标签[32, 55, 56]。除了静态的基于规则的随机化方法外，自动领域随机化[46, 57]引入了一种动态随机化过程，通过自动搜索在目标领域中表现最优的随机化策略。尽管DR能够生成具有一定多样性的任务集合，但它受限于原始任务的设定。最近，研究人员提出通过随机化生成全新的任务，这种方法旨在最小化结构偏差[38, 58, 59]。这一方向为推进具有最小归纳偏见的可扩展ICRL提供了有前景的路径。

3 方法论 3.1 通过世界随机化生成AnyMDP任务

动机。为了促进通用的上下文学习（ICL），我们认为一个关键步骤是减少训练集中各任务之间的共享结构（共同基础）。通过降低对内循环学习（IWL）的依赖，训练出的模型将主要依靠上下文学习（ICL）来解决具体任务。这表明，应优先考虑任务的多样性与覆盖范围，而非对真实场景的保真度，以确保模型能够泛化到更广泛的模式之外。受这些进展的启发，我们提出了“世界随机化”（world randomization），它在传统领域随机化的基础上进一步扩展，以更大程度地增加任务集的多样性。该方法旨在最小化任务构建过程中的归纳偏见，从而进一步提升任务的多样性和覆盖范围。

3.2建模和培训框架

逐步监督（Step-wise Supervision, SS）：使用强化学习（RL）或进化策略进行上下文强化学习（ICRL）的元训练面临显著挑战，包括繁琐的基础设施需求和高昂的计算成本。近年来，基于监督学习的元训练方法取得了进展，例如算法蒸馏（Algorithm Distillation, AD）[12]、ADε[14]以及决策预训练变换器（Decision Pre-Training Transformers, DPT）[15]，这些方法在可扩展的ICRL元训练方面展现出潜力。然而，在推理阶段出现了一个关键挑战：上下文轨迹由模型自身生成，导致训练时与推理时的轨迹之间不可避免地存在差距。这种差异可能导致部署过程中的灾难性失败。尽管已有研究表明，利用多样化的行为轨迹可以缓解这一问题[15]，但这又带来了维持训练效率的新挑战。

为解决这些局限性，我们提出了逐步监督（Step-wise Supervision, SS）框架，该框架受到DPT以及模仿学习中的数据聚合技术[61, 62]的启发。我们的方法依赖于两个关键策略：行为策略（behavior policy）指在训练过程中用于生成轨迹的策略；参考策略（reference policy）则是被模仿的目标策略，且不直接参与执行。将行为策略与参考策略解耦，使得我们可以在行为策略中引入多样性，从而减小训练与推理轨迹之间的差异，同时保持参考策略的最优性，如图1所示。

与DPT仅模仿基于完整轨迹的单步动作不同，我们的逐步监督框架本质上被设计为与面向序列模型（如Transformer[63]及其优化变体[64, 65]）的高效率分块训练（chunk-wise training）流程相匹配。

这一点也在图2中有所展示。

前向传播在各个片段之间递归地进行计算，而后向传播则在每个片段内部执行。线性注意力层的记忆状态 φt的梯度在片段之间被阻断。因此，公式（3）被替换为：

其中SG表示停止梯度（Stop Gradient）。在元训练过程中，梯度在每个片段内部计算，并首先累积在缓存中，仅在轨迹结束时才应用到模型参数上。

4 实验

4.1 泛化性与可扩展性的验证

元训练过程主要在8块Nvidia Tesla A800 GPU上进行。我们每块GPU使用5的批量大小，并将序列划分为每段2000步的片段（chunks）。优化采用AdamW算法，学习率从峰值2×10⁻⁴开始衰减。对于长度T = 12K的轨迹，每次迭代的平均耗时为8秒，且该耗时随序列长度线性增加。更多细节请参见附录C.2。对于因果序列模型，我们评估了四种架构：RWKV-7 [67]、门控Delta网络（Gated Delta-Net, GDN）[69]、门控自注意力（Gated Self-Attention, GSA）[65] 和 Mamba2 [70]。在先前的测试中，线性注意力类型的模型架构在训练效率和上下文记忆长度方面均优于Transformer。因此，图4中的实验专注于线性注意力架构。测试结果与语言处理领域的结论基本一致，证明了AnyMDP作为长期序列建模基准的有效性。在比较了计算效率和性能之后，我们选择RWKV-7用于后续实验。

在不进行任何额外参数调优的情况下，我们在图3中对未见过的AnyMDP任务、Gymnasium任务以及图16中的DarkRoom [12] 上评估了我们的模型OmniRL，其性能结果如表1所示。值得注意的是，与以往的ICRL研究不同，我们的训练集中并未包含任何DarkRoom任务实例。在实验中，所选任务被限制在观测空间维度ns ≤ 128、动作空间维度na ≤ 5的环境中。对于具有连续观测空间的环境（例如Pendulum-v1），我们使用基于网格的离散化方法将观测空间手动划分为60个离散类别。为了使在na = 5条件下训练的OmniRL适用于动作数更少（na < 5）的环境，我们将未使用的动作重新映射到有效动作上。这进一步证明了OmniRL在不同动作空间维度环境中的兼容性。我们还发现，适当的奖励塑形对OmniRL的有效运行至关重要，如图14所示；详细信息见附录C.3。

在表1中，我们比较了OmniRL、经典表格型Q学习（TQL）[71]结合置信上界（UCB）[72]（简称TQL-UCB）以及近端策略优化（PPO）[73]的归一化性能、回合成本和步成本。结果表明，当OmniRL仅在所提出的AnyMDP任务上进行元训练时，能够有效适应大多数Gymnasium任务，验证了AnyMDP任务的表征能力。这一结果也展示了OmniRL卓越的样本效率，这与先前ICRL的研究发现一致。值得注意的是，尽管OmniRL仅在单智能体任务上进行了训练，它仍能通过配置观测空间来适应Switch2等多智能体任务，从而实现智能体之间的涌现式协作，而无需在训练过程中进行显式的多智能体交互，因此将协作行为的出现与集中式机制解耦。此外，如预期一致，随着状态空间（ns）或动作空间（na）规模的增加，AnyMDP任务的求解难度也随之上升，其中PPO对动作空间扩展更为敏感，而TQL-UCB则对状态空间增长更为脆弱，如图9所示。

ICRL可视化。采用线性注意力机制的ICRL能够在其记忆（φt）中捕获解决环境所需的所有信息。我们进行了全面的t-SNE分析，以检验在在线强化学习（Online-RL）评估过程中，这些记忆在不同任务间的演变情况。如图5所示，聚类模式证实了Gym、Darkroom和AnyMDP任务分布的差异性。值得注意的是，Darkroom和Gym的聚类主要位于左上区域，而AnyMDP占据更广阔的空间区域，反映出其更高的多样性。这种空间上的区分强调了AnyMDP的独特特性，并凸显了OmniRL在多样化任务中强大的泛化能力。

4.2 OmniRL在离线和在线学习中表现更好

为了进行消融研究并与包括AD、ADε和DPT在内的其他方法进行比较，我们收集了一个较小的数据集，其中包含128K条序列用于训练，其中ns = 16，na = 5，T = 8K，总共有10亿个时间步。图6总结了在Dsmall上训练的不同方法的性能，训练迭代次数相同。比较包括AD、ADε、DPT、OmniRL和OmniRL（无先验信息），其中从序列中移除了先验信息gt。

我们检查了不同初始上下文下不同方法的性能：(1) 在线RL：代理从空轨迹开始（h0 = ∅）。(2) 离线RL：代理从不完美的演示（例如，受干扰的专家策略）中派生的现有上下文开始（h0 = hπ）。(3) 模仿学习：代理从专家派生的现有上下文开始（h0 = h(exp)）。对于所有三类，后续交互会持续添加到上下文中。因此，这些模型仅在初始记忆或缓存方面有所不同。评估考察了代理在两个关键领域的能力：它们利用现有信息的能力以及持续探索和利用的能力。在图6的结果中，OmniRL和OmniRL（无先验信息）以较大差距超越了AD、ADε和DPT，验证了逐步监督（SS）的有效性。OmniRL（无先验信息）在所有三组中均落后于OmniRL，差距明显，这表明整合先验信息是有效的。表2和图17进一步证明了OmniRL的离线学习能力可以推广到Gymnasium环境。

4.3 随任务数量增加涌现出通用型ICRL

我们通过在四个独立数据集上进行元训练，验证了任务多样性在ICRL中的关键作用。每个数据集 D(Ttra(16, 5)) 均包含 128K 个序列，但在任务数量上有所不同（|Ttra| ∈ {100, 1K, 10K, 128K}）。需要注意的是，即使来自同一任务，不同的轨迹也可能因行为策略的差异以及决策和状态转移过程中的随机采样而产生。我们考察了验证损失 Lt 如何随着元训练迭代次数（外循环步数）和上下文内步骤数 t（内循环步数）的同时变化，在已见任务和未见任务上的表现，结果如图7所示。我们有如下几点观察：

任务数量对ICRL的泛化能力至关重要。以往关于上下文内学习（ICL）的研究 [26, 29] 强调了“突发性”（burstiness）的重要性。我们的结果首次表明，即使仅使用“突发性”序列，任务的数量及其整体多样性仍然极为关键。具体而言，在 |Ttrain| ≤ 10K 的组别中，过拟合导致在未见任务上出现ICL的“暂时性”现象 [19]，而在已见任务上的性能则持续提升。这些发现确认并拓展了Kirsch等人 [32] 以及Pan等人 [75] 所提到的“任务识别”阶段的结论。基于Chan等人 [31] 关于ICL和上下文内学习（IWL）的理论，一种可能的解释是：IWL主导了模型性能，即模型记忆了各个任务的特征，而ICL的作用是选择正确的任务，从而实现对已见任务的快速适应，但对未见任务的泛化能力较差。随着任务数量持续增加，由于记忆每个任务的特定信息变得不可行，模型对ICL的依赖程度逐渐增强。这表现为对未见任务的泛化能力提升，同时在已见和未见任务上的适应周期也变得更长，如图7所示。

ICRL的泛化能力可能与其零样本或少样本性能相矛盾。我们的结果揭示了关于ICRL评估的一个关键见解。以往大多数ICRL研究基于固定的、较短的上下文范围内的平均结果来评估性能。然而，我们的发现表明，更具泛化能力的上下文内学习器在零样本甚至少样本评估中可能表现更差，特别是当训练集与评估集存在显著重叠时——即评估集更接近于已见任务时。因此，我们认为，更关键的是关注学习器的渐近性能（asymptotic performance）。这可以通过考察在足够长的上下文的最后几步或最后若干轮次中的表现来有效评估，而非依赖短期指标。

5 结论与讨论

我们提出了一种可扩展、结构偏差较小的任务集合，用于基准测试和研究上下文逆向强化学习（Inverse Contextual Reinforcement Learning, ICRL）。我们的框架引入了两项关键创新：逐步监督（stepwise supervision）和先验信息引入（prior information induction），从而构建了一个更高效的ICRL系统。实验结果表明，与现有方法相比，我们的模型在强化学习任务中展现出更优越的泛化能力。所提出的模型比以往任何时候都覆盖更广泛的强化学习任务。本工作为ICRL研究提供了一个更具可扩展性的实验平台。

更广泛的影响：与以往研究相辅相成，我们的研究发现强调，任务多样性、序列长度——以及序列的“突发性”（burstiness）——是实现通用ICRL的关键决定因素。我们的结果还主张将评估指标转向渐近性能的衡量方式。本研究进一步推动了为大规模元训练量身定制、精心构建的合成数据集的发展。

局限性与未来工作：我们承认，离散的状态和动作空间是本工作最主要的局限性，这限制了其在连续状态和动作空间环境中的应用。将本工作扩展到连续空间和部分可观测环境，并引入更多模态信息，将显著拓宽其应用范围。

原文链接：https://arxiv.org/pdf/2502.02869