Constrained Optimization of Charged Particle Tracking with Multi-Agent Reinforcement Learning
作者: Tobias Kortus, Ralf Keidel, Nicolas R. Gauger, Jan Kieseler
分类: physics.comp-ph, cs.AI, cs.LG
发布日期: 2025-01-09
💡 一句话要点
提出基于多智能体强化学习与约束优化的带电粒子径迹重建方法
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)
关键词: 多智能体强化学习 粒子径迹重建 分配约束 约束优化 质子成像
📋 核心要点
- 现有粒子径迹重建方法难以在复杂物理系统中实现端到端训练和优化。
- 提出一种多智能体强化学习方法，通过协同优化策略最小化粒子散射，并使用安全层保证分配约束。
- 实验表明，该方法在模拟数据上优于多种基线方法，且约束优化提高了泛化能力和预测稳定性。
📝 摘要（中文）
本文提出了一种基于多智能体强化学习（MARL）的带电粒子径迹重建方法，该方法利用分配约束来优化像素化粒子探测器中的粒子径迹。通过协同优化参数化策略，该方法旨在最小化重建径迹上的总粒子散射量。为了满足约束条件，保证粒子命中的唯一分配，该方法引入了一个安全层，为每个联合动作解决一个线性分配问题。此外，为了强制成本裕度，增加局部策略预测与优化器映射决策边界的距离，建议在黑盒梯度估计中添加一个组件，迫使策略找到总分配成本更低的解决方案。在为质子成像开发的粒子探测器模拟数据上，实验结果表明，与多个单智能体和多智能体基线相比，该方法的有效性。此外，还证明了具有成本裕度的约束在优化和泛化方面的有效性，这体现在具有高重建性能的更广泛区域以及降低的预测不稳定性。该研究为基于强化学习的径迹重建的进一步发展奠定了基础，通过约束策略增强了性能，并通过个体和团队奖励的选项提供了优化跟踪算法的更大灵活性。
🔬 方法详解
问题定义：论文旨在解决像素化粒子探测器中带电粒子径迹重建问题。现有方法在处理复杂物理系统时，难以进行端到端的训练和优化，并且难以保证粒子命中的唯一分配，从而影响重建精度。
核心思路：论文的核心思路是将粒子径迹重建问题建模为一个多智能体强化学习问题，每个智能体负责一部分粒子的径迹重建。通过协同优化所有智能体的策略，可以最小化重建径迹上的总粒子散射量。同时，引入分配约束来保证粒子命中的唯一分配，从而提高重建精度。
技术框架：该方法的技术框架主要包括以下几个模块：1) 多智能体强化学习模块：使用多个智能体协同优化参数化策略，每个智能体负责一部分粒子的径迹重建。2) 安全层：为每个联合动作解决一个线性分配问题，以满足分配约束，保证粒子命中的唯一分配。3) 黑盒梯度估计模块：用于估计策略梯度，并引入成本裕度，增加局部策略预测与优化器映射决策边界的距离，从而提高重建性能。
关键创新：该方法最重要的技术创新点在于将多智能体强化学习与分配约束相结合，从而实现了对粒子径迹重建问题的有效求解。与现有方法相比，该方法可以进行端到端的训练和优化，并且可以保证粒子命中的唯一分配，从而提高了重建精度。
关键设计：在多智能体强化学习模块中，使用了参数化的策略函数来表示每个智能体的策略。在安全层中，使用了线性分配问题来求解粒子命中的唯一分配。在黑盒梯度估计模块中，引入了成本裕度，以增加局部策略预测与优化器映射决策边界的距离。具体的参数设置、损失函数和网络结构等技术细节在论文中有详细描述。
🖼️ 关键图片
📊 实验亮点
实验结果表明，该方法在模拟的质子成像粒子探测器数据上，与多个单智能体和多智能体基线相比，具有更好的重建性能。特别地，引入成本裕度的约束优化方法，不仅提高了重建性能，还增强了泛化能力和预测稳定性，表现为具有高重建性能的区域更广，预测不稳定性降低。
🎯 应用场景
该研究成果可应用于高能物理实验、医学成像（如质子治疗）、以及其他需要精确粒子径迹重建的领域。通过提高粒子径迹重建的精度和效率，可以帮助科学家更好地理解物质的微观结构，并为疾病诊断和治疗提供更准确的信息。此外，该方法还可以推广到其他需要解决分配约束的优化问题。
📄 摘要（原文）
Reinforcement learning demonstrated immense success in modelling complex physics-driven systems, providing end-to-end trainable solutions by interacting with a simulated or real environment, maximizing a scalar reward signal. In this work, we propose, building upon previous work, a multi-agent reinforcement learning approach with assignment constraints for reconstructing particle tracks in pixelated particle detectors. Our approach optimizes collaboratively a parametrized policy, functioning as a heuristic to a multidimensional assignment problem, by jointly minimizing the total amount of particle scattering over the reconstructed tracks in a readout frame. To satisfy constraints, guaranteeing a unique assignment of particle hits, we propose a safety layer solving a linear assignment problem for every joint action. Further, to enforce cost margins, increasing the distance of the local policies predictions to the decision boundaries of the optimizer mappings, we recommend the use of an additional component in the blackbox gradient estimation, forcing the policy to solutions with lower total assignment costs. We empirically show on simulated data, generated for a particle detector developed for proton imaging, the effectiveness of our approach, compared to multiple single- and multi-agent baselines. We further demonstrate the effectiveness of constraints with cost margins for both optimization and generalization, introduced by wider regions with high reconstruction performance as well as reduced predictive instabilities. Our results form the basis for further developments in RL-based tracking, offering both enhanced performance with constrained policies and greater flexibility in optimizing tracking algorithms through the option for individual and team rewards.