JAX-LOB: A GPU-Accelerated limit order book simulator to unlock large
  scale reinforcement learning for trading
Sascha Frey Kang Li Peer Nagy Silvia Sapora Chris Lu Stefan Zohren Jakob Foerster Anisoara Calinescu
Financial exchanges across the world use limit order books (LOBs) to process
orders and match trades. For research purposes it is important to have large
scale efficient simulators of LOB dynamics. LOB simulators have previously been
implemented in the context of agent-based models (ABMs), reinforcement learning
(RL) environments, and generative models, processing order flows from
historical data sets and hand-crafted agents alike. For many applications,
there is a requirement for processing multiple books, either for the
calibration of ABMs or for the training of RL agents. We showcase the first
GPU-enabled LOB simulator designed to process thousands of books in parallel,
with a notably reduced per-message processing time. The implementation of our
simulator - JAX-LOB - is based on design choices that aim to best exploit the
powers of JAX without compromising on the realism of LOB-related mechanisms. We
integrate JAX-LOB with other JAX packages, to provide an example of how one may
address an optimal execution problem with reinforcement learning, and to share
some preliminary results from end-to-end RL training on GPUs.