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.


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.


RL agent

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

RL environment

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


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.


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

      title={Self-Correcting Quantum Many-Body Control using Reinforcement Learning with Tensor Networks}, 
      author={Friederike Metz and Marin Bukov},


View Github