# 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.
- : 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}
}
```