# JAX-LOB: A GPU-Accelerated limit order book simulator to unlock large scale reinforcement learning for trading

arXiv: https://arxiv.org/abs/2308.13289 · HTML: https://mlsanity.com/paper/2308.13289/

## Key facts

- Cited by 14 papers in the map (data through 14 Oct 2026).
- Builds on 15 of its 45 references that are in the map.
- First posted on arXiv 25 Aug 2023.

## Builds on (most cited first, 15 of 15)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Deep Reinforcement Learning with Double Q-learning](https://mlsanity.com/paper/1509.06461/) (2015)
- [Machine learning accelerated computational fluid dynamics](https://mlsanity.com/paper/2102.01010/) (2021)
- [Discovered Policy Optimisation](https://mlsanity.com/paper/2210.05639/) (2022)
- [Market Making via Reinforcement Learning](https://mlsanity.com/paper/1804.04216/) (2018)
- [Double Deep Q-Learning for Optimal Execution](https://mlsanity.com/paper/1812.06600/) (2018)
- [Get Real: Realism Metrics for Robust Limit Order Book Market Simulations](https://mlsanity.com/paper/1912.04941/) (2019)
- [Reinforcement Learning for Market Making in a Multi-agent Dealer Market](https://mlsanity.com/paper/1911.05892/) (2019)
- [ABIDES-Gym: Gym Environments for Multi-Agent Discrete Event Simulation and Application to Financial Markets](https://mlsanity.com/paper/2110.14771/) (2021)
- [Fast Agent-Based Simulation Framework with Applications to Reinforcement Learning and the Study of Trading Latency Effects](https://mlsanity.com/paper/2008.07871/) (2020)
- [Market Making with Scaled Beta Policies](https://mlsanity.com/paper/2207.03352/) (2022)
- [Model-based gym environments for limit order book trading](https://mlsanity.com/paper/2209.07823/) (2022)
- [Some challenges of calibrating differentiable agent-based models](https://mlsanity.com/paper/2307.01085/) (2023)
- [Imitate then Transcend: Multi-Agent Optimal Execution with Dual-Window Denoise PPO](https://mlsanity.com/paper/2206.10736/) (2022)

## Built on by (most cited first, 14 of 14)

- [JaxMARL: Multi-Agent RL Environments and Algorithms in JAX](https://mlsanity.com/paper/2311.10090/) (2023)
- [XLand-MiniGrid: Scalable Meta-Reinforcement Learning Environments in JAX](https://mlsanity.com/paper/2312.12044/) (2023)
- [minimax: Efficient Baselines for Autocurricula in JAX](https://mlsanity.com/paper/2311.12716/) (2023)
- [StockSim: A Dual-Mode Order-Level Simulator for Evaluating Multi-Agent LLMs in Financial Markets](https://mlsanity.com/paper/2507.09255/) (2025)
- [Deep Limit Order Book Forecasting](https://mlsanity.com/paper/2403.09267/) (2024)
- [LOB-Bench: Benchmarking Generative AI for Finance -- an Application to Limit Order Book Data](https://mlsanity.com/paper/2502.09172/) (2025)
- [FinEvo: From Isolated Backtests to Ecological Market Games for Multi-Agent Financial Strategy Evolution](https://mlsanity.com/paper/2602.00948/) (2026)
- [Reinforcement Learning-Based Market Making as a Stochastic Control on Non-Stationary Limit Order Book Dynamics](https://mlsanity.com/paper/2509.12456/) (2025)
- [A Research Agenda for Usability and Generalisation in Reinforcement Learning](https://mlsanity.com/paper/2412.16970/) (2024)
- [Revisiting Ensemble Methods for Stock Trading and Crypto Trading Tasks at ACM ICAIF FinRL Contest 2023-2024](https://mlsanity.com/paper/2501.10709/) (2025)
- [PB2: Preference Space Exploration via Population-Based Methods in Preference-Based Reinforcement Learning](https://mlsanity.com/paper/2506.13741/) (2025)
- [Constrained Policy Optimization for Provably Fair Order Matching](https://mlsanity.com/paper/2604.06522/) (2026)
- [ReLOBGen: Replayable Limit Order Book Message Generation](https://mlsanity.com/paper/2609.35867/) (2026)
- [Mixture-of-Experts for Cryptocurrency Order Execution: Training Stability, Tail Risk, and Failure Modes](https://mlsanity.com/paper/2610.03369/) (2026)

## Abstract (arXiv)

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.