Skip to content
/ qdx Public

Quantum error correction code AI-discovery with Jax

License

Notifications You must be signed in to change notification settings

jolle-ag/qdx

Repository files navigation

overview

QEC AI-discovery with JAX ⚛️🤖🚀

arXiv Open In Colab

Code repository for the paper "Simultaneous Discovery of Quantum Error Correction Codes and Encoders with a Noise-Aware Reinforcement Learning Agent" by Jan Olle, Remmy Zen, Matteo Puviani and Florian Marquardt.

Description

This library can be used to train Reinforcement Learning (RL) agents to codiscover quantum error correction (QEC) codes and their encoding circuits from scratch, without any additional domain knowledge except how many errors are not detected given the quantum circuit it has built.

The RL agent can be made noise-aware, meaning that it learns to produce encoding strategies simultaneously for a range of noise models, making it applicable in very broad situations.

overview

The whole RL algorithm, including the Clifford simulations of the quantum circuits, are implemented in Jax. This enables parallelized training of multiple agents on vectorized environments in a single GPU. Our implementation builds upon PureJaxRL.

Installation

QDX can be installed by:

  1. Cloning the repository
git clone https://github.com/jolle-ag/qdx.git
cd qdx
  1. Installing requirements
pip install -r requirements.txt

GPU compatibility:

To install JAX with NVIDIA GPU support, use:

# CUDA 12 installation
pip install -U "jax[cuda12]"

Usage Example

We include a demo jupyter notebook for two different situations: [[7,1,3]] code discovery in a fixed symmetric depolarizing noise channel and noise-aware [[6,1]] code discovery in a biased noise channel.

Open In Colab

License

The code in this repository is released under the MIT License.

Citation

@article{olle_simultaneous_2023,
  title={Simultaneous Discovery of Quantum Error Correction Codes and Encoders with a Noise-Aware Reinforcement Learning Agent},
  author={Olle, Jan and Zen, Remmy and Puviani, Matteo and Marquardt, Florian},
  url = {http://arxiv.org/abs/2311.04750},
  journal={arXiv preprint arXiv:2311.04750},
  urldate = {2024-04-01},
  publisher = {arXiv},
  month = nov,
  year = {2023},
  note = {arXiv:2311.04750 [quant-ph]},
}