Model-aware reinforcement learning for high-performance Bayesian experimental design in quantum metrology

Abstract: Quantum sensors offer control flexibility during estimation by allowing manipulation by the experimenter across various parameters. For each sensing platform, pinpointing the optimal controls to enhance the sensor's precision remains a challenging task. While an analytical solution might be out of reach, machine learning offers a promising avenue for many systems of interest, especially given the capabilities of contemporary hardware. We have introduced a versatile procedure capable of optimizing a wide range of problems in quantum metrology, estimation, and hypothesis testing by combining model-aware reinforcement learning (RL) with Bayesian estimation based on particle filtering. To achieve this, we had to address the challenge of incorporating the many non-differentiable steps of the estimation in the training process, such as measurements and the resampling of the particle filter. Model-aware RL is a gradient-based method, where the derivatives of the sensor's precision are obtained through automatic differentiation (AD) in the simulation of the experiment. Our approach is suitable for optimizing both non-adaptive and adaptive strategies, using neural networks or other agents. We provide an implementation of this technique in the form of a Python library called qsensoropt, alongside several pre-made applications for relevant physical platforms, namely NV centers, photonic circuits, and optical cavities. This library will be released soon on PyPI. Leveraging our method, we've achieved results for many examples that surpass the current state-of-the-art in experimental design. In addition to Bayesian estimation, leveraging model-aware RL, it is also possible to find optimal controls for the minimization of the Cram\'er-Rao bound, based on Fisher information.