AI在可控核聚变领域取得重大突破！DeepMind最新研究通过强化学习算法，将等离子体形状精度提升65%，为“人造太阳”实现精准放电指明了方向。
编者按：本文来自微信公众号新智元（ID：AI_era），微新创想经授权转载，头图来源摄图网。一旦人类掌握了可控核聚变能，将可拥有无穷无尽的清洁能源。可控核聚变对于“托卡马克装置”的运行至关重要，反馈控制在其中扮演着关键角色。控制系统会主动管理磁线圈，以控制拉长离子体的不稳定性，防止破坏性的垂直事件发生。此外，若能实现对等离子体电流、位置和形状的精确控制，还可实现热排放，甚至对其能量进行管理。科学家们一直致力于研究等离子体配置变化对这些相关量的影响，因此需要能够用于新配置，并围绕标称场景快速变化的系统。
传统上，等离子体的精确控制是通过等离子体电流、形状和位置的连续闭环来实现的。在这种模式下，控制设计者预先计算出一组前馈线圈电流，然后为每个受控量建立反馈回路。然而，等离子体形状和位置无法直接测量，必须通过磁场测量实时间接估算。尤其是等离子体的形状，必须使用平衡重构代码进行实时估算。这类系统已成功稳定了大范围的放电，但其设计不仅具有挑战性，还耗时，特别是针对新型等离子体情况。
值得强调的是，强化学习（RL）已成为构建实时控制系统的另一种全新范式。2022年，DeepMind团队在Nature上发表论文表明，RL设计的系统能够成功实现“托卡马克磁控制”的主要功能。这项工作提出了一个系统，RL智能体通过与FGE托卡马克模拟器交互，学习控制托卡马克配置变量（TCV）。智能体学习的控制策略随后被集成到TCV控制系统中，通过观察TCV的磁场测量，并为所有19个磁控线圈输出控制指令。研究人员展示了RL智能体控制各种情况的能力，包括高度拉长的等离子体、雪花，甚至还展示了同时在真空室中，使用两个独立等离子体稳定“液滴”配置的新方法。
尽管RL方法有许多优点，但其局限性限制了其作为控制托卡马克等离子体的实用解决方案的应用。最新研究中，DeepMind决定要解决三个挑战：指定一个既可学习又能激发精确控制器性能的标量奖励函数、追踪误差的稳态误差、较长的训练时间。团队提出了“奖励塑形”的方法，通过向智能体提供明确的错误信号，和集成错误信号来解决积分器反馈中的稳态误差问题。此外，在片段分块和迁移学习中，解决了生成控制策略所需的训练时间问题。研究人员针对复杂的放电情况采用了多重启动方法，使得训练时间大幅缩减。
强化学习控制等离子体的最新论文中，研究人员采用与Nature那篇论文相同的基本实验。RL通过与模拟环境的交互，学习特定实验的控制策略，然后TCV上部署由此产生的放电策略。具体来讲，使用自由边界模拟器FGE进行动态建模，并添加了额外随机性，以模拟传感器值和电源的噪声，并改变等离子体的参数。研究人员使用最大后验优化（MPO）算法来制定控制策略。智能体与1000份FGE环境进行交互，收集看到的观察结果、采取的行动，以及获得的奖励。每一步获得的奖励，都是根据等离子体状态与参考值中包含的目标值的接近程度来计算的，并辅以其他因素，如避免不良等离子体状态。
从最优控制范式到强化学习的直接转换是，为每个要最小化的误差项设置一个奖励分量，其中每个分量都被映射为一个标量值。然后将这些值合并为一个标量奖励值。根据观察、行动和奖励的记录序列，智能体使用正则化损失函数上的梯度下降交替更新策略和critic网络。对于等离子体放电，actor网络被限制在一个能以10kHz频率执行的小型架构中，但critic网络只在训练过程中使用，因此可以足够复杂地学习环境动态。
面向实用的强化学习控制器在具体任务实操中，研究人员演示了智能体具体训练过程。首先讨论了通过奖励塑形来提高控制精度。然后介绍了通过积分观测来减少稳态误差的工作，讨论了使用“episode chunking”来改善现实的训练时间。最后探讨了迁移学习作为提高训练效率的手段。
奖励塑形（reward shaping）传统控制算法用各种办法来最小化主动测量（或估计）的数量误差，而强化学习（RL）算法则旨在最大化一个通用定义的奖励信号。在训练过程中，这种奖励最大化目标能推动智能体行为的演化，但是在部署时不会计算奖励值。在经典控制算法中，控制器的性能可以通过显式调整控制增益（例如，修改响应性或干扰抑制）和调整多项输入多项输出（MIMO）系统的权衡权重来进行调整。相比之下，在强化学习中，奖励函数对于被学习的控制器行为至关重要。因此，需要仔细设计奖励函数来调整控制器行为。
研究人员发现，通过调整奖励函数的设计，他们可以快速适应智能体的行为，并权衡目标的不同方面。此外，研究人员证明了塑形奖励函数对于创建准确的强化学习控制策略是必不可少的。而且他们进一步展示了通过使用更新后的奖励函数继续训练，可以将智能体应用到新的目标上。
研究人员在先前研究的基础上修改了为磁控而设计的奖励函数。研究人员使用加权的SmoothMax函数来组合奖励组件的值。在某些情况下，一个单独的奖励组件由多个相关的误差量构成，比如在多个控制点处的形状误差。研究人员还利用SmoothMax函数将这些误差组合成一个单一的标量奖励组件。奖励组件并不受限于从传感器测量中获得，这在构建中就能提供了额外的灵活性。奖励组件还可以是多模态的，例如鼓励智能体远离状态空间中不理想或模拟器建模较差的区域。
研究人员使用用SoftPlus转换来获得标量奖励组件。理论上，许多参数的选择应该是近似等效的，因为它们是奖励的单调调整，不应该对最优策略产生很大影响。然而，在实践中，研究者依赖于梯度下降（gradient descent），并没有一个完美的全局优化器（global optimizer）。研究人员需要在面对随机回报的情况下探索全局空间。很好和很差的紧密值使得很难找到任何可观的奖励区域（或者在如何改进方面有明显的梯度）。另一方面，较宽松的很差值使得更容易找到奖励信号，但更难以发现精确的控制，因为改进时奖励变化较小。直观上，因此，“紧密”奖励参数可能更适用于初始条件接近目标状态的情况，因此奖励不需要塑造目标发现，而应更注重精确性。
在简单环境中的奖励塑形在研究人员的初始实验中，考虑了三种训练方法，重点是通过修改“shape_70166”任务中形状误差的奖励组件的超参数来最小化形状误差。1. 基准线：采用之前研究的默认奖励参数 – good = 0.005，bad = 0.05。2. 窄化奖励：将参数更新为good = 0和bad = 0.025。3. 奖励调度（reward schedule）：将good和bad的值在训练过程中逐渐调整为更加尖峰（more Peaked），good = 0，bad从0.1逐渐减少到0.025，共进行600万次策略更新步骤。该调度在训练开始时提供了一个较宽的奖励区域来帮助探索，随着训练的进行逐渐收紧奖励函数，以鼓励准确性。
积分器（integrator）反馈积分误差的近似可以通过递归神经网络来计算，然而，它们更容易过度拟合仿真动态。在这项工作中，研究人员采用了一种更简单的解决方案：没有让策略（policy）学习积分误差，而是手动计算它，并将其附加到前馈策略所观察到的观测集中。他们特别关注了减少等离子体电流（??）的稳态误差，之前研究的训练策略表现出明显的偏差，并且该误差可以很容易地计算。研究人员在“shape_70166”任务中评估了将平均误差信号纳入考虑的好处。在该任务中，等离子体电流和形状的参考值是恒定的，环境初始化后实际值接近参考值。因此，智能体的主要目标是控制稳态误差（steady-state）。研究人员发现，积分器反馈显著降低了等离子体电流偏差。
Episode Chunking在TCV上的实验持续1-2秒，相当于以10kHz的控制频率进行10,000 – 20,000个时间步。FGE模拟器（如上所述用于训练智能体）在训练过程中使用一颗AMD EPYC 7B12 CPU核心，每个典型的模拟步骤大约需要2秒钟，使用随机动作。因此，FGE生成包含10,000个步骤的一次完整episode大约需要5小时的时间。这意味着在最理想的情况下，即智能体在第一次尝试之前已经知道最佳策略，训练时间仍然会约为5小时（以观察高质量的结果）。实际上，强化学习智能体需要探索动作空间以找到最佳策略。因此，根据任务复杂性，训练时间可能从几天到几周不等。此外，研究人员的任务结构使得智能体需要按顺序学习相对独立的“技能”。例如，在“showcase_xpoint”任务中，智能体必须先使等离子体变形，然后移动其垂直位置，然后改变其流向，最后恢复原始形状。研究人员观察到该任务的学习过程发生在两个明显的阶段。将分块（chunking）技术应用于展示_x点（showcase_xpoint）任务，并分别使用两个/三个块，可以显著缩短训练时间。
转移学习在试图减少训练时间时，一个自然的问题是问是否可以重用之前放电时训练的模型，也就是说，智能体在解决一个初始任务时积累的知识在多大程度上可以转移到一个相关的目标任务上。研究人员以两种形式考察迁移学习的性能：1. 零样本（Zero-shot）：在目标任务上运行在初始任务上学习的策略，而无需进行任何额外的数据收集或策略参数更新。2. 微调（Fine tuning）：使用在初始任务上学习的模型的权重来初始化策略和值函数，然后使用这些权重在新的目标任务上通过与环境交互进行训练，其中目标任务作为奖励。需要注意的是，这要求在两个任务中使用相同的架构（actor和critic网络）。在两种情况下，研究人员使用在showcase_xpoint任务上训练的智能体参数作为迁移的初始参数。
研究人员在TCV托卡马克上对一些上述的仿真改进进行了测试。通过这种方式，研究人员可以评估当前工作的优势和局限性，并为下一步的改进提供方向。等离子体形状精度的奖励塑形研究人员检查了奖励塑形在两种不同配置和目标上所带来的精度改进：减少形状稳定任务中的LCFS误差和提高“snowflake_to_perfect”任务配置中的X点精度。研究人员将模拟结果与TCV上的实验结果以及来自Degrave等人（2022）的可比实验进行了比较。与先前的研究一样，研究人员通过将演员网络（由JAX图定义）创建为共享库对象来部署控制策略，其中命令的动作是输出高斯分布的均值。
研究人员首先测试了一个控制策略，该策略通过在奖励塑形部分中讨论的奖励塑形方法来减少shape_70166稳定任务中的LCFS误差。对于这个稳定任务，研究人员使用了TCV的标准击穿过程和初始等离子体控制器。在0.45秒时，控制权移交给学习的控制策略，然后它试图在1秒的持续时间内维持固定的等离子体电流和形状。放电后，研究人员使用LIUQE代码计算重构的平衡态。在1秒的放电过程中的每个0.1毫秒时间片内，研究人员计算等离子体形状的误差。研究人员比较了三个实验的精度，分别从模拟放电和TCV放电中测量形状误差。
X点位置精度的奖励塑形接下来，研究人员将比较奖励塑形对更复杂的“snowflake”配置的影响。该策略的训练奖励被塑形以增加X点控制的准确性。与稳定实验中一样，等离子体是通过标准的TCV程序创建和初始控制的，在0.45秒时将控制权移交给强化学习控制器。在这个实验中，RL训练的策略成功地建立了一个两个X点距离为34厘米的“snowflake”。然后，该策略成功将两个X点带到了目标距离6.7厘米的位置，接近建立一个所谓的“完美snowflake”。然而，在1.0278秒（即交接后的0.5778秒），等离子体因垂直不稳定性而发生破裂。
通过“Episode Chunking”来加速训练最后，研究人员验证了使用“Episode Chunking”来减少训练时间，特别是验证在TCV放电中是否出现可能的“不连续性”。研究人员进行了一个在showcase配置下使用3个块进行训练的实验。这个实验的重建平衡态的时间轨迹可以看到。研究人员发现实验按预期进行，没有因为“episode chunking”而产生明显的伪影。这证明了这种训练加速方法没有损失质量。
20世纪50年代起，众多科学家们致力于探索、攻克可控核聚变这一难题。DeepMind最新研究，用强化学习算法大幅提升了等离子体的精度，极大缩短了学习新任务的训练时间。这为可控核聚变在未来实现“精准放电”，能量管理铺平了道路。在为人类获取海量清洁能源，以改变未来的能源路线图上，DeepMind再次点亮了一盏明灯。
