轨道校正强化学习项目 (v1.2)

本项目使用TD3（Twin Delayed Deep Deterministic Policy Gradient）强化学习算法来校正粒子加速器中的电子束轨道。通过强化学习智能体控制水平和垂直校正器的角度，使电子束轨道尽可能接近理想轨道。

新增功能 (v1.2)

• 奖励函数重新设计：改进奖励函数，增强对小误差的敏感性，提供更精细的完成奖励，并增加稳定性奖励
• 增强型奖励机制：鼓励BPM数值与目标轨道越接近越好，根据BPM与目标轨道的差距动态调整动作奖励策略
• 超参数优化更新：贝叶斯优化现在支持优化稳定性奖励权重参数，扩展了优化参数空间

新增功能 (v1.1)

• 贝叶斯超参数优化：集成Optuna库实现强化学习超参数的自动优化，包括奖励权重、学习率、噪声参数等
• 并行优化支持：支持多线程并行执行超参数优化，提高优化效率
• 动态评估策略：根据优化进度自动调整训练回合数，前期快速筛选，后期精细评估
• 早期停止机制：在训练过程中实现性能监控，避免无效训练

项目结构

• fodo_lattice.py: 定义FOFO晶格结构，包括磁铁、漂移段、监测器和校正器等组件
• td3_agent.py: 实现TD3强化学习算法，包括环境、智能体和训练过程
• BayesianOptimize.py: 实现基于Optuna的贝叶斯超参数优化功能
• Bayesian_method.ipynb: Jupyter笔记本文件，用于执行贝叶斯超参数优化和结果分析
• lattice_simulator.py: 实现晶格模拟器，用于高效计算轨道响应矩阵
• main.ipynb: 原始的Jupyter笔记本文件，用于基础训练和测试过程

工作原理

该项目通过强化学习算法控制粒子加速器中的束流轨道。环境中，智能体观察当前的轨道偏差状态，并决定如何调整水平和垂直校正器的角度来减小偏差。奖励函数设计鼓励减小轨道偏差，同时惩罚过大的校正动作。

贝叶斯优化通过构建概率模型来智能地搜索最优超参数组合，相比于网格搜索和随机搜索，能够更高效地找到接近最优的参数配置。

依赖库及版本

• Python 3.10
• OCELOT 24.03.0
• PyTorch 1.12.1
• NumPy 1.21.5
• Gym 0.21.0
• Optuna 3.1.0 或更高版本

安装说明

1. 确保已安装Python 3.10
2. 安装OCELOT:

   conda install -c ocelot-collab ocelot
   

   git clone https://github.com/ocelot-collab/ocelot.git
   

   python setup.py install

3. 安装其他依赖:

   pip install torch==1.12.1 numpy==1.21.5 gym==0.21.0 optuna>=3.1.0
   

使用方法

基础训练

1. 运行训练: 打开Bayesian_method.ipynb并在Jupyter Notebook中执行所有单元格来训练智能体。

2. 查看结果: 训练过程中的奖励、轨道误差等指标将在Notebook中显示。

超参数优化

1. 运行贝叶斯优化: 在BayesianOptimize.py中配置优化参数，然后运行该脚本开始超参数优化。

2. 配置优化选项:

   • n_calls: 优化迭代次数
   • random_state: 随机种子，确保结果可重现
   • fast_eval: 是否使用快速评估模式
   • n_jobs: 并行任务数，设置为大于1的值可启用并行优化

3. 查看优化结果: 优化完成后，将显示最佳超参数组合和对应的性能指标。

功能特性

动态奖励权重调整

在训练过程中，系统会根据训练进度和智能体性能动态调整奖励函数中各部分的权重，以平衡不同目标。

校正器动态权重

根据当前轨道误差分布动态调整不同校正器的控制权重，使误差较大的区域对应的校正器获得更高的权重。

数据归一化

对BPM读数进行动态归一化处理，基于实际观测数据的统计特性进行标准化，提高训练稳定性。

多维度奖励函数

奖励函数综合考虑轨道改善、动作幅度、稳定性和完成情况等多个维度，引导智能体学习高质量的控制策略。

当前问题和后续改进方向

根据最新的研究进展，我们已经完成了以下改进：

1. 晶格设计优化：

   • 已在晶格出口处增加BPM监测器，提供更完整的轨道信息
   • 优化了磁铁和校正器布局，提高了校正效率

2. 算法和模型改进：

   • 重新设计了奖励函数，提供更密集和有效的奖励信号
   • 调整了神经网络结构，支持更复杂的策略学习

3. 计算性能优化：

   • 重写了轨道生成模块，使用矩阵运算替代OCELOT调用
   • 实现了响应矩阵方法，显著提高了计算效率
   • 支持GPU加速计算，充分利用硬件资源

4. 并行化和分布式训练：

   • 解决了OCELOT框架的多进程限制问题
   • 实现了真正的并行训练以大幅提升训练速度

尽管取得了上述进展，目前仍存在以下关键问题需要解决：

1. 轨道发散问题：尽管进行了多次优化，算法在训练过程中仍然经常出现轨道发散现象，无法稳定收敛到目标轨道。这可能是由于奖励函数设计、动作空间定义或环境动态特性等因素导致。

2. 训练稳定性不足：算法在不同随机种子下的表现差异很大，有些情况下能取得较好效果，但在大多数情况下训练过程不稳定。

3. 泛化能力有限：训练得到的策略在特定初始条件下表现良好，但在面对不同的初始轨道偏差时，校正效果往往不佳。

后续计划从以下几个方向进行深入研究：

1. 深入分析发散原因：

   • 详细分析奖励函数的有效性，可能需要引入更多维度的奖励信号
   • 检查动作空间设计是否合理，是否存在动作幅值或频率问题
   • 研究环境状态表示是否足够完整，是否缺少关键信息

2. 改进训练策略：

   • 尝试使用不同的RL算法，如PPO、SAC等进行比较
   • 引入课程学习机制，从简单场景逐步过渡到复杂场景
   • 设计更好的探索策略，避免陷入局部最优

3. 增强算法鲁棒性：

   • 增加训练场景的多样性，提高策略泛化能力
   • 引入对抗训练机制，提高算法在各种扰动下的稳定性