Quantum control of solid state qubits via Reinforcement Learning

The gist The use of reinforcement learning (RL) in the design of quantum control strategies for pulse level control of NISQ era quantum devices is attracting growing interest.

NV Quantum Registers and Indirect Control

The work investigates the use of model-free RL to design pulse sequences for controlling nuclear spins in a nitrogen-vacancy (NV) center while accounting for its microscopic spin-bath. NV centers in diamond are promising systems for solid state quantum memories and networks, and they can be indirectly controlled via the electronic spin or directly via radio frequency (rf) fields. The system involves a multiqubit register consisting of the NV center – electron spin and the nitrogen 14N nuclear spin – together with a set of 13C nuclear spins coupled to the electronic one. The full Hamiltonian describing the system is given by H = HR + HB + HI, where HR is the register’s Hamiltonian, HB is the bath Hamiltonian, and HI is the register-bath interaction.

Spin-Bath Modeling and Simulation

To model the NV center registers and their surrounding spin-bath microscopically, the generalized cluster correlation expansion (gCCE) technique is exploited. The gCCE allows going beyond a simple mean field treatment of the spin-bath and retaining its microscopic structure and memory, which are lost in a Lindblad master equation approach. The simulation involves calculating the n-th order gCCE contribution to the density matrix by tracing out bath spins, where the Hamiltonian entering the free evolution propagator contains H(E0) = HR + H(E0)B,m + H(E0)I,m. This process is used to simulate open quantum dynamics and reward RL agents.

Reinforcement Learning for Pulse Sequence Design

The RL agent acts on the environment via the application of mw pulses to the NV electronic spin to accomplish state preparation or gate design tasks. The action space is continuous, parametrized by idle time, pulse duration, and pulse phase. Given these specifications, the Proximal Policy Optimization (PPO) algorithm is a natural choice for the problem at hand. The agent learns optimal strategies to maximize the cumulative expected reward thus accomplishing its task(s).

Sparse Reward Strategies and Fidelity Estimation

The agents are trained using sparse reward strategies, where rewards are assigned only at the end of an episode. For state engineering, the state fidelity F = Tr(ρρ˜) can be estimated via a single-shot reward given by 1/dTr[ρPin]e(in)m /Pr(in). For gate engineering, the process fidelity F = Tr(U†V)2/d2 can be estimated using Full direct fidelity estimation (DFE) or Hofmann classical fidelity bounds. The 0-fidelity estimator is also formulated as a single-shot reward strategy.

Domain Randomization and Sim-to-Real Transfer

Domain randomization (DR) is used at two levels: randomizing the state configuration of the spin-bath within the gCCE, and randomizing relevant physical parameters like the Rabi frequency of mw pulses or coupling constants. DR helps produce robust policies that remain effective when transferred to the real world. Furthermore, pre-training an agent on a simulated environment and then fine-tuning it with single-shot fidelity rewards in different environments serves as a blueprint for experimental use.

Benchmarking and Comparison

The performance of RL agents is benchmarked against black-box optimization methods such as Bayesian Optimization (BO) and Evolution strategies like CMA-ES. The comparison often involves using a 500-shot fidelity estimate to emulate the completely mixed state of the spin-bath for agents A(ss) and B(ss). Results show that RL agents can discover qualitatively different high-fidelity protocols and reach parameter space regions not explored by other quantum optimization strategies. The work concludes that RL agents can learn to cope with the presence of a non-Markovian spin-bath and find solutions robust to uncertainty in the theoretical modeling of the environment.

Key Findings on Performance

Agents trained with the gCCE1 approximation (agent B) perform better than those without a spin-bath (agent A) when evaluated in environments accounting for exact open system dynamics. Agents B learn to use shorter times and larger magnetic fields, which is attributed to the stiffening of the quantization axis by the larger magnetic field. The use of domain randomization combined with fine-tuning allows agents to adapt to new physical realizations, restoring high fidelities for most bath realizations. The final protocol lays out a clear pipeline for the ready use of our agents in real-world experimental situations.

Hyperparameters and Robustness

The PPO algorithm used for training employed specific hyperparameters, including 2000 steps, batch size 200, and a learning rate of 5 × 10−4. The analysis of Bloch sphere trajectories shows how often the trajectories from different optimization methods remain near the poles or approach the equator. The results demonstrate that agents B(gCCE1) are more robust than agents A(ss) when tested under domain randomization, showing a clear increase in performance when evaluated in environments sampled outside the training distribution.

State Preparation Success

The RL agents successfully found high fidelity pulse sequences for state preparation tasks, such as preparing a Bell state between two 13C nuclear spins using the NV center as a bus. The work demonstrates that RL agents can find high fidelity pulse sequences for this task, which is not true for other black-box optimization methods even in the absence of a spin-bath. The ability of RL to learn robust protocols suggests potential applications in real solid-state qubit devices.

Benchmarking Summary

The comparison with black-box optimization methods shows that while RL agents can find high fidelity solutions, they are comparable to, and often better than, the solutions found by other black-box optimization algorithms. The different methods seem to explore distinct parameter regimes reflecting different optimization strategies for reaching high-fidelity pulse sequences. The comparison with single-shot RL agents using the DFE estimator is significant because it shows that they still manage to learn high-fidelity solutions with the direct fidelity estimate strategy outperforming the other two during evaluation.

Conclusion

In this work, we have shown that RL agents can learn to cope with the presence of a non-Markovian spin-bath and find solutions robust to uncertainty in the theoretical modeling of the environment. The crucial element is using domain randomization during training, combined with fine-tuning at deployment and the use of single-shot fidelity rewards, which lays down a blueprint for experimental use. The potential for applications in real solid-state qubit devices is highlighted.

Acknowledgments

This project was made possible by the DLR Quantum Computing Initiative and the Federal Ministry for Research, Technology and Space; qci.dlr.de/projects/qlearning and the Science Sphere concept of Ulm University. The authors acknowledge funding through the IQST and BMFTR (QSENS-QMAT).

References

[1] V. V. Sivak, A. Eickbusch, H. Liu, B. Royer, I. Tsioutsios, and M. H. Devoret.

[2] Luigi Giannelli, Sofia Sgroi, Jonathon Brown, Gheorghe Sorin Paraoanu, Mauro Paternostro, Elisabetta Paladino, and Giuseppe Falci.

[3] Valentin Gebhart, Raffaele Santagati, Antonio Andrea Gentile, Erik M Gauger, David Craig, Natalia Ares, Leonardo Banchi, Florian Marquardt, Luca Pezze, and Cristian Bonato.

[4] Mario Krenn, Jonas Landgraf, Thomas Foesel, and Florian Marquardt.

[5] Shreyasi Mukherjee, Dario Penna, Fabio Cirinna, Mauro Paternostro, Elisabetta Paladino, Giuseppe Falci, and Luigi Giannelli.

[6] Yuri Alexeev, Marwa H Farag, Taylor L Patti, Mark E Wolf, Natalia Ares, Alan Aspuru-Guzik, Simon C Benjamin, Zhenyu Cai, Shuxiang Cao, Christopher Chamberland.

[7] Marin Bukov and Florian Marquardt.

[8] Tobias Kiermeyer, Thomas Heydenreich, Leo Van Damme, Sebastian Hohenemser, Florian Marquardt, and Steffen J Glaser.

[9] Callum W Duncan, Pablo M Poggi, Marin Bukov, Nikolaj Thomas Zinner, and Steve Campbell.

1. Bold header: Model-Free RL for Indirect Control of Nuclear Spins

This allows an AI system to learn tailored pulse sequences with higher fidelities compared to benchmark solutions we obtain with other state-of-the-art black-box optimization methods or the same agents trained without the spin-bath, by accounting for the microscopic structure of the bath via generalized cluster correlation expansion (gCCE).

2. Bold header: Sparse Reward Strategies for Hardware Feasibility

The AI can be trained using single(few)-shot(s) fidelity estimate rewards, such as 1/dTr[ρPin]e(in)m /Pr(in) for state preparation, which is explicitly noted as being amenable to be implemented directly on an experimental setup via local Pauli measurements.

3. Bold header: Domain Randomization for Robust Sim-to-Real Transfer

The AI can be trained with domain randomization, where a different state-configuration of the spin-bath is uniformly sampled and the agent learns a strategy dealing with this uncertainty, leading to robust policies against changes in relevant parameters like the Rabi frequency of the mw pulses or the coupling constants.

4. Bold header: Pre-training for Fine-Tuning on Real Devices

The system can utilize a two-stage learning pipeline where an agent is first pre-trained on a ‘realistic’ model of a nuclear register, accounting for variability in the physical parameters via DR, and then fine-tuned by continuing learning with either exact fidelity reward or single-shot fidelity reward on new environments.

5. Bold header: Benchmarking Against Black-Box Control

The AI can be compared against established methods like Bayesian Optimization (BO) and CMA-ES, allowing researchers to determine that RL agents can discover qualitatively different high-fidelity protocols and reach parameter space regions not explored by other quantum optimization strategies.

The use of reinforcement learning (RL) in the design of quantum control strategies for pulse level control of NISQ era quantum devices is attracting growing interest. Model-free RL, not requiring preexisting knowledge of the physical system, acts as a black-box optimization tool well suited for noisy devices. In this work, we provide a simulation-based analysis of the performance of model-free RL for state preparation and qubit gate design in a challenging setup: indirect control of NV nuclear spins in a spin-bath. In particular, the spin-bath is modeled via the generalized cluster correlation expansion accounting for magnetic impurities and intra-bath interactions. We analyze the performance of the RL agents by performing simulations with minimally observable environments and noisy, sparse rewards as introduced in (Phys. Rev. X 12, 011059), and compare them to other state-of-the-art black-box optimization tools.