Constrained collaborative optimization of charged particle tracking with multi-agent reinforcement learning

Reinforcement learning (RL) demonstrated immense success in modeling complex physics-driven
systems, providing end-to-end trainable solutions by interacting with a simulated or real environ-
ment, maximizing a scalar reward signal. In this work, we propose, building upon previous work,
an end-to-end multi-agent RL approach with assignment constraints for reconstructing particle
tracks in pixelated particle detectors. Our approach optimizes collaboratively a parameterized
policy, functioning as a heuristic to a multidimensional assignment problem, by jointly minimiz-
ing the total amount of particle scattering over the reconstructed tracks in a readout frame. To sat-
isfy constraints, guaranteeing a unique assignment of particle hits, we propose a safety layer solv-
ing a linear assignment problem for every joint action. Further, to enforce cost margins, increas-
ing the distance of the local policies predictions to the decision boundaries of the optimizer map-
pings, 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 simu-
... mehr
lated 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.