Self-Correcting Quantum Many-Body Control using Reinforcement Learning with Tensor Networks

This repository contains the code and data for the corresponding preprint article arXiv:2201.11790.

Summary

We present a novel Q-learning framework (QMPS) specifically designed for controlling 1d spin chains in which the RL agent is represented by a combination of a matrix product state (MPS) and a neural network (NN). The algorithm can be used to find optimal control protocols that prepare a target (ground) state starting from a set of initial states, and as an example we implement the paradigmatic mixed-field Ising model. To reach system sizes which lie beyond exact simulation techniques, we employ matrix product states as a representation for the quantum state and as a trainable machine learning ansatz. The hybrid MPS+NN architecture is then optimized via backpropagation and conventional gradient descent.

Content

RL agent

• dqn/main.py: Script that performs one full instance of training for specified (hyper)parameters and plots/saves the results.
• dqn/dqn.py: QMPS agent (essentially a DQN agent where the ansatz is composed of a hybrid NN+MPS network).
• dqn/dqn_utils.py: Functions for training & evaluating a QMPS agent.
• dqn/models.py: QMPS ansatz with forward and backward passes.
• dqn/replay_buffers.py: DQN replay buffer.

RL environment

• environment/env.py: Gym-style RL environment for controlling a spin chain
• environment/tlfi_tn.py: Functions for simulating a spin chain. In this case the transverse longitudinal field Ising (TLFI) model.
• environment/tlfi_tn_model.py: MPO implementation of the TLFI Hamiltonian

trained_models/ contains the corresponding data of the control studies presented in the paper (including the optimized QMPS parameters).

Requirements

The code is written in Python and apart from the usual libraries (numpy, scipy, matplotlib) you need to have the following packages installed:

• JAX: For performance enhancenment via just-in-time compilation.
• TensorNetwork: For the spin chain simulations.

Citation

If you use our code/models for your research, consider citing our paper:

@misc{metz2022,
      title={Self-Correcting Quantum Many-Body Control using Reinforcement Learning with Tensor Networks}, 
      author={Friederike Metz and Marin Bukov},
      year={2022},
      eprint={2201.11790},
      archivePrefix={arXiv},
      primaryClass={quant-ph}
}