粒子加速器中基于硬件在线强化学习的微秒级延迟反馈技术
引言：科学实验控制面临的挑战
在前沿科学领域，未来大规模科学实验的调试与运行面临着巨大挑战。以粒子加速器为例，众多待调整的参数使得手动调优极为困难，不仅耗费大量人力，还会大幅增加项目预算，因为设施的调试和运行需要大量资源。尽管自动算法（优化器）已应用于粒子加速器，但随着任务中变量增多，“维度灾难”问题凸显，导致其性能下降。
与此同时，数据驱动的机器学习（ML）技术备受关注，强化学习（RL）作为一种极具潜力的方法，通过训练智能体学习策略以最大化奖励函数来实现对环境的控制。然而，传统的机器学习库主要针对吞吐量性能优化，不适用于微秒级延迟应用；多数可编程逻辑实现侧重计算加速，无法满足实时环境的要求。因此，将强化学习部署在边缘设备上成为解决这些问题的关键。
强化学习的理论基础
在强化学习中，环境被建模为马尔可夫决策过程（MDP），由状态空间（$S$）、动作空间（$A$）、转移动态（$P$）和奖励集合（$R$）构成的四元组。智能体的目标是最大化期望回报（$G$），其公式为$G_t = \sum_{k=0}^{\infty} \gamma^k R_{t+k+1}$，其中$\gamma \in [0,1]$是折扣因子，用于在无限 horizon 问题中约束累积奖励，进而将 MDP 扩展为$(S, A, P, R, \gamma)$。
由于环境的完整状态往往不能直接获取，部分可观测马尔可夫决策过程（POMDP）应运而生，它在 MDP 的基础上增加了可能观测的空间（$O$）和观测到$o_t$时转移到新状态$s_t$的概率（$\epsilon$）。
现代许多强化学习算法采用演员-评论家架构，演员函数负责选择动作，评论家函数用于估计给定状态的期望回报。评论家可实现值函数（$v_\pi: S \rightarrow \mathbb{R}$）或动作值函数（$Q_\pi: S \times A \rightarrow \mathbb{R}$）。在深度强化学习中，演员和评论家均通过神经网络（NN）进行近似。
经验积累器架构与 KINGFISHER 系统
为克服传统方法的局限，论文引入了经验积累器架构。该架构将训练和推理分离，在边缘设备上实现实时策略（$\pi^{(\text{edge})}$）的推理，利用低延迟计算平台执行前向传播，即根据观测信号预测下一个动作。同时，策略神经网络与环境的交互被记录下来，生成状态-动作-奖励元组，用于在传统计算平台（如 CPU）上异步训练智能体的仿真副本（$\pi^{(\text{CPU})} \gets \pi^{(\text{edge})}$）。这种方式在保证低推理延迟的同时，显著提高了系统的灵活性和抽象性。
基于经验积累器架构，构建了 KINGFISHER 系统。它是一个基于 AMD-Xilinx Versal VCK190 的异构计算平台，集成了 FPGA、ARM 处理器和 AI 引擎（AIE）阵列。该系统通过 40 Gbps 的 Aurora 64b/66b 链路接收数据，用于创建观测数据流，并通过控制数模转换器（DAC）生成模拟控制信号来控制粒子加速器。此外，系统具备数据存储和传输功能，通过直接内存访问（DMA）块将观测-动作数据流写入 DDR 内存，ARM 处理器再将数据复制到网络文件中。为实现对水平贝塔振荡（HBO）的精确控制，系统还结合了 KAPTURE 系统等硬件设施，形成了完整的控制链路。
PPO 算法在系统中的应用
论文选择近端策略优化（PPO）算法来实现强化学习控制。PPO 算法对超参数变化稳定性好，尽管其样本效率相对离策略算法（如软演员-评论家算法，SAC）较低，但在本应用中，由于经验收集速率高，这一劣势得到了弥补。
演员网络在 AIE 中实现，通过神经网络选择高斯分布的均值来确定动作，概率分布的标准差作为可训练参数与神经网络系数一同更新。在实验中，系统与加速器进行了 2048 次交互，外部冲击器激发振荡后，每次情节结束都进行一次训练步骤，更新神经网络的系数，并将新的权重和偏置上传到智能体。整个推理循环的延迟仅为 2.8 微秒。
实验结果分析
- 与传统 FIR 控制器对比：通过实验验证，强化学习智能体在控制水平贝塔振荡方面性能卓越。与传统的有限冲激响应（FIR）控制器相比，智能体能够自动适应束流电流的变化，而 FIR 控制器因其线性特性，性能会受到束流电流变化的显著影响。
- 非线性响应优势：从非线性响应的角度看，神经网络智能体能够展现出非线性输出。通过计算总谐波失真加噪声（THD+N）指标发现，当输入信号幅度较高时，智能体能够学习应用更复杂的动作，以更好地控制振荡。
- 训练方式比较：在训练方式的比较中，与模拟训练相比，在加速器上直接训练虽然存在数据访问开销，但对于计算密集型模拟的系统，能够大幅缩短训练时间。而且，在加速器上训练的智能体可以直接应用于实际运行，而模拟训练的智能体在迁移到现实环境时可能会受到未建模现象的影响，导致性能下降。
结论与展望
本文通过在粒子加速器中实现基于硬件在线强化学习的微秒级延迟反馈，成功突破了传统控制方法的瓶颈。经验积累器架构和训练时奖励定义的引入，为基于强化学习的控制器在现实边缘环境中的应用提供了重要的技术支持。
借助 KINGFISHER 系统，展示了智能控制超快速非线性动力学的巨大潜力，这种方法有望在粒子加速器、聚变实验等更多类似系统中得到广泛应用，推动科学实验和工业设施控制技术迈向新的高度。
未来，还可以进一步研究如何优化系统性能，降低数据访问开销和能耗，以及如何更好地处理复杂环境中的各种不确定性，以实现更高效、更智能的控制。
引用
Scomparin, L., Caselle, M., Santamaria Garcia, A., Xu, C., Blomley, E., Dritschler, T., … & Weber, M. (n.d.). Microsecond-Latency Feedback at a Particle Accelerator by Online Reinforcement Learning on Hardware.