Optimizing Spin-Magnetometers with Reinforcement Learning

Quantum sensors have been widely successful in both research and industry, leveraging quantum resources to surpass their classical counterparts in measurement precision. Despite this, there is still room for improvement in achieving maximal sensing precision. Sensors operate in four stages: state preparation, controlled parametrization, measurement, and parameter estimation. In order to reach theoretical bounds in sensitivity, each stage in the sensing pipeline requires careful optimization. Here, we focus on the first two stages, using tools from reinforcement learning (RL). Using the soft actor-critic (SAC) algorithm, we train an RL agent on simulated data of a spin-based magnetometer in the presence of an ambient magnetic field, which we aim to measure. The output of the agent is a stochastic policy, from which we may sample control fields to apply during the parametrization stage, so as to maximize the quantum Fisher information (QFI) accumulated by the system. We evaluate the performance of the agent in a limited control context, as it varies with the spin quantum number, and in the presence of decoherence, finding that it converges to an optimal solution. Additionally, we apply the agent to situations not seen in training, and find it is able to generalize. This architecture is agnostic to the particular system, making it relatively easy to adapt to various different sensors and control contexts. We ultimately find that RL is a valuable resource for quantum optimal control, directly applicable to quantum sensing, and to the broader quantum information community.