面向核聚变等离子体控制的离线强化学习：代码库与基准测试

# 面向核聚变等离子体控制的离线强化学习：代码库与基准测试
来源：https://arxiv.org/html/2606.07550
杨福¹，†† 包浩民²，†† 罗希特·松克³ 胡晓燕³ 阿拉文德·韦努戈帕尔³ 杰夫·施耐德³ 陈嘉宇⁴，‡‡
¹中南大学
²重庆大学
³卡内基梅隆大学
⁴香港大学
†† 同等贡献。‡‡ 通讯作者。

###### 摘要

https://github.com/LucasCJYSDL/Offline-RL-Kit-for-Nuclear-Fusion
陈嘉宇 ([email protected])
离线强化学习（RL）为利用历史托卡马克数据开发等离子体控制器提供了一条有前景的途径，因为在真实设备上进行在线试错代价高昂且风险巨大。然而，由于缺乏针对核聚变中现实的多执行器、长时域等离子体控制问题的标准化离线强化学习基准，这一方向的进展仍难以衡量。我们推出了 RL4F，一个用于核聚变等离子体控制的离线强化学习基准，它提供了闭环评估环境和跨四个全剖面跟踪任务（旋转、密度、温度和压力）的基线比较。该评估环境所基于的动态函数来自真实托卡马克 DIII-D 的历史放电数据。我们在统一协议下评估了多种模仿学习和离线强化学习基线方法。我们发现，离线基于模型的强化学习方法在大多数目标上获得了最佳平均性能，但没有任何单一方法能在所有任务中占优，这突显了动力学建模在复杂、长时域等离子体控制任务中的重要性。为了促进进一步研究，我们开源了代码库、数据集和评估框架，该基准不仅面向聚变社区，也适用于离线强化学习算法的发展。

## 1 引言

核聚变通过利用为恒星提供能量的反应，为获得丰富、低碳的能源提供了一条潜在途径 (Giet al., 2020 (https://arxiv.org/html/2606.07550#bib.bib34))。托卡马克是实现可控核聚变最有前途的约束装置之一，但其运行需要对高温、不稳定且强耦合的等离子体进行实时控制。最近的研究表明，强化学习可用于训练此类控制器。特别是，Degraveet al. (2022 (https://arxiv.org/html/2606.07550#bib.bib5)) 在 TCV 托卡马克上演示了基于深度强化学习的磁控制，后续研究已将基于强化学习的等离子体控制扩展到撕裂模避免、剖面跟踪和下降控制 (Traceyet al., 2024 (https://arxiv.org/html/2606.07550#bib.bib6); Seoet al., 2024 (https://arxiv.org/html/2606.07550#bib.bib7); Charet al., 2023 (https://arxiv.org/html/2606.07550#bib.bib9); Wanget al., 2025 (https://arxiv.org/html/2606.07550#bib.bib8))。这些进展表明，强化学习可以通过在高维托卡马克控制问题中直接优化反馈策略来补充传统的等离子体控制设计。

直接在真实托卡马克上开发强化学习控制器难以规模化：托卡马克运行昂贵、时间有限且对安全性要求极高。一种自然的方式是在部署前先训练候选策略并在仿真中评估其闭环行为。基于物理的仿真器，例如 RAPTOR (Feliciet al., 2011 (https://arxiv.org/html/2606.07550#bib.bib32))、前向 Grad-Shafranov 静态 (FGE) 仿真器 (Carpanese, 2021 (https://arxiv.org/html/2606.07550#bib.bib14))，以及更近期的 TORAX (Citrinet al., 2024 (https://arxiv.org/html/2606.07550#bib.bib33))，为前向建模、轨迹优化和控制器开发提供了可靠工具。先前的研究已经展示了在这些仿真器上训练的强化学习策略具有良好的控制性能 (Degraveet al., 2022 (https://arxiv.org/html/2606.07550#bib.bib5); Traceyet al., 2024 (https://arxiv.org/html/2606.07550#bib.bib6))。然而，基于物理的仿真器在计算上比传统的强化学习仿真器昂贵得多，尤其是在强化学习智能体在训练过程中随机探索执行器空间时，这常常导致迭代求解器收敛较慢。此外，将这些仿真器适配到特定的托卡马克可能需要大量的建模选择、参数辨识和校准。另一种方法是从历史实验数据中学习面向控制的动力学模型，从而产生与该特定设备更直接相关的仿真环境 (Charet al., 2023 (https://arxiv.org/html/2606.07550#bib.bib9); Sonkeret al., 2026 (https://arxiv.org/html/2606.07550#bib.bib41))。尽管最近取得了进展，基于强化学习的剖面控制仍然具有挑战性：等离子体剖面是高维空间量，其动力学是非线性的，并且跨多个执行器耦合，实际的强化学习控制器在奖励规范、稳态跟踪偏差和样本效率方面仍然面临困难 (Traceyet al., 2024 (https://arxiv.org/html/2606.07550#bib.bib6))。这些挑战激发了对托卡马克等离子体控制进行数据驱动训练和离线评估的标准化基准的需求。

我们的主要贡献是 RL4F，一个用于托卡马克剖面控制的离线强化学习统一基准。剖面指的是等离子体量从芯部（即等离子体最内区域）到边缘（即最外区域）的径向空间变化。它具有四个关键特征。1) **真实的预部署工作流程。** 我们首先从历史 DIII-D 实验放电中训练一个参考动力学模型，然后使用该模型生成用于离线策略学习的轨迹。候选算法仅从模型生成的数据集中学习，并在参考动力学模型上进行闭环评估，这反映了在实际机器测试之前必须从过去实验数据开发控制器的实际设置。2) **多任务剖面跟踪。** 该基准涵盖四个全剖面跟踪任务，包括旋转、密度、温度和压力，这些任务展现了不同的控制难度和等离子体响应通道。3) **场景相关的执行器空间。** 策略通过一个共享的动作空间执行，该空间由中性束功率、中性束扭矩、气体注入和电子回旋加热组成，涵盖了与 DIII-D（位于圣地亚哥的托卡马克装置）剖面控制相关的主要加热、动量输入和加料通道。4) **大规模数据。** 该基准包含 5,882 次放电和 945,828 个转换（经过过滤后），并具有固定的训练、验证和测试集划分。据我们所知，这是首个专门为聚变等离子体控制中的离线强化学习设计的基准。

我们在相同的闭环协议和剖面级跟踪指标下，评估了涵盖模仿学习、无模型离线强化学习以及基于模型离线强化学习的代表性基线方法。我们的评估表明，基于模型的方法通常优于无模型基线方法，而没有任何单一算法在所有剖面目标上占优，这突显了托卡马克剖面控制作为一个具有挑战性的离线强化学习基准。

## 2 相关工作

**聚变等离子体控制。** 核聚变是可再生能源发电的主要候选方案之一。一个核心挑战是控制等离子体剖面以实现稳定、高性能的运行。传统的等离子体控制通常依赖于预先计算的馈送线圈电流轨迹 (Walker and Humphreys, 2006 (https://arxiv.org/html/2606.07550#bib.bib1))，并结合针对单个目标量的反馈回路。这种剖面控制系统已在多个托卡马克上实现，包括 JET 的 safety q-剖面控制 (Moreauet al., 2003 (https://arxiv.org/html/2606.07550#bib.bib2))、TCV 的 q-剖面和电子温度控制 (Bartonet al., 2015 (https://arxiv.org/html/2606.07550#bib.bib3))，以及 EAST 的 q-剖面控制 (Wanget al., 2021 (https://arxiv.org/html/2606.07550#bib.bib4))。这些方法在广泛的放电场景中展现了良好的性能，但设计可能具有挑战性且耗时，尤其是在控制量为高维或强耦合的等离子体场景中。最近，强化学习作为设计聚变系统反馈控制器的新框架出现。Degraveet al. (2022 (https://arxiv.org/html/2606.07550#bib.bib5)) 在 Maximum-a-Posteriori (MPO) (Abdolmalekiet al., 2018 (https://arxiv.org/html/2606.07550#bib.bib15)) 框架下训练了一个基于强化学习的控制器，用于跟踪各种等离子体配置下的位置、电流和形状。这一研究方向由 Traceyet al. (2024 (https://arxiv.org/html/2606.07550#bib.bib6)) 进一步扩展。Charet al. (2023 (https://arxiv.org/html/2606.07550#bib.bib9)) 开发了一个用于跟踪 βN 和等离子体旋转量的离线强化学习框架，并使用近端策略优化 (PPO) 训练控制器，该控制器后来被用于设计贝叶斯优化 (BO) 风格的控制器 (Sonkeret al., 2025 (https://arxiv.org/html/2606.07550#bib.bib12)) 以减轻撕裂不稳定性。Seoet al. (2024 (https://arxiv.org/html/2606.07550#bib.bib7)) 应用深度确定性策略梯度 (DDPG) (Lillicrapet al., 2016 (https://arxiv.org/html/2606.07550#bib.bib16)) 方法，在保持 H 模性能的高压等离子体的同时避免撕裂不稳定性。Wanget al. (2025 (https://arxiv.org/html/2606.07550#bib.bib8)) 训练强化学习策略以避免在下降阶段发生破裂。

**离线强化学习基准。** 著名的 D4RL 数据集 (Fuet al., 2021a (https://arxiv.org/html/2606.07550#bib.bib10)) 在具有偏差数据分布的挑战性机器人控制场景中基准测试离线强化学习。Dope 基准 (Fuet al., 2021b (https://arxiv.org/html/2606.07550#bib.bib17)) 建立在 D4RL 和 RL Unplugged (Gulcehreet al., 2020 (https://arxiv.org/html/2606.07550#bib.bib11)) 之上，侧重于离策略评估。Qinet al. (2022 (https://arxiv.org/html/2606.07550#bib.bib13)) 提出 NeoRL 以缩小早期离线强化学习基准与实际场景之间的差距。Liuet al. (2023 (https://arxiv.org/html/2606.07550#bib.bib19)) 提出了一个基准测试套件，以促进在训练和部署阶段离线安全强化学习算法的发展与评估。Parket al. (2025 (https://arxiv.org/html/2606.07550#bib.bib18)) 提出了用于离线目标条件强化学习的 OGBench。与这些现有基准相比，聚变等离子体控制任务呈现出独特的挑战，包括高度非线性和随机的动力学、部分可观测性以及安全关键的操作约束。据我们所知，RL4F 是首个专门针对聚变等离子体控制这一关键真实世界操作场景的离线强化学习基准。

## 3 仿真器

**问题设置。** 我们将托卡马克剖面控制建模为一个有限时域马尔可夫决策过程 (MDP) M = (S, A, P, r, ρ0, γ, H)，其中 S 是状态空间，A 是动作空间，P(s'|s,a) 是转移分布，r(s,a) 是奖励函数，ρ0 是初始状态分布，γ 是折扣因子，H 是情节时域。策略 π(a|s) 旨在最大化期望折扣回报

J(π) = E_{π,P,ρ₀} [∑_{t=0}^{H-1} γ^t r(s_t, a_t)]。 (1)
参考图标题
图 1: 参考 RPNN 动力学集成从历史 DIII-D 运行数据中训练，并用作评估训练策略的闭环环境。
在离线强化学习中，智能体在训练期间无法与环境交互 (Levineet al., 2020 (https://arxiv.org/html/2606.07550#bib.bib42))，而是被提供一个固定的轨迹数据集 D = {τ_i}_{i=1}^N，其中每条轨迹由未知行为策略收集的转换 (s_t, a_t, r_t, s_{t+1}) 组成。这种设置与托卡马克控制非常匹配，因为在线试错代价高昂且风险巨大。

我们使用循环概率神经网络 (RPNN) 从轨迹数据中学习等离子体动力学，遵循先前数据驱动的托卡马克动力学建模和离线基于模型的控制工作 (Charet al., 2023 (https://arxiv.org/html/2606.07550#bib.bib9); Sonkeret al., 2026 (https://arxiv.org/html/2606.07550#bib.bib41))。给定当前等离子体状态 s_t、先前的执行器设置 a_{t-1} 以及执行器增量 Δa_t = a_t - a_{t-1}，模型预测下一个状态变化 Δs_t = s_{t+1} - s_t 的分布。具体来说，RPNN 输出高斯分布的参数 μ_t, log σ_t² = f_θ(s_t, a_{t-1}, Δa_t)，并且下一个状态通过将预测的状态变化添加到当前状态来自回归地推进。这种循环概率公式允许模型捕捉依赖于历史的等离子体演化，同时为长时域 rollout 提供不确定性估计。动力学模型训练工作流程的示意图如图 1 (https://arxiv.org/html/2606.07550#S3.F1) 所示。

**动力学建模。** 鉴于真实托卡马克设备的可及性有限，我们训练一个参考动力学模型作为数字孪生。该参考动力学模型从历史 DIII-D 实验放电中训练，并用于生成离线训练数据以及提供闭环评估环境。如图 2 (https://arxiv.org/html/2606.07550#S3.F2) 所示，我们采用两阶段 RPNN 训练程序 (Sonkeret al., 2026 (https://arxiv.org/html/2606.07550#bib.bib41))。

第一阶段使用均方误差目标优化预测均值。第二阶段从第一阶段检查点初始化，冻结预测主干，并使用负对数似然目标训练对数方差头，从而在不改变已学习的均值动力学的情况下校准不确定性。遵循先前基于模型的控制实践 (Chenet al., 2025 (https://arxiv.org/html/2606.07550#bib.bib44))，我们训练了一个自助 (bootstrapped) RPNN 集成；预测的对数方差捕获了偶发不确定性，而集成成员之间的不一致性提供了认知不确定性的估计。更多训练细节见附录 B (https://arxiv.org/html/2606.07550#A2)。

参考图标题
图 2: 动力学模型的两阶段训练程序。
我们在大约 18,000 次历史 DIII-D 实验放电上训练了一个 25 成员的 RPNN 集成，这些数据跨越了近十年的数据收集。每次放电包含大约四秒的数据，采样间隔为 20 毫秒。对于剖面量，包括电子温度、离子温度、密度、压力、旋转和安全因子 q 剖面，我们使用 ZipFIT 重构 (Loganet al., 2018 (https://arxiv.org/html/2606.07550#bib.bib43))，它提供平滑、受物理约束的剖面估计。遵循先前的托卡马克动力学建模工作 (Charet al., 2023 (https://arxiv.org/html/2606.07550#bib.bib9))，我们在动力学模型训练之前使用 PCA 降低剖面量的维度。训练集成的保留预测保真度在附录 B.1 (https://arxiv.org/html/2606.07550#A2.SS1) 中报告。为了模拟托卡马克实验，我