This paper introduces JAX-LOB, the first GPU-accelerated limit order book simulator built using JAX, designed to enable large-scale reinforcement learning for trading. The simulator addresses the low signal-to-noise ratio and overfitting issues in high-frequency trading by massively parallelizing order book processing across thousands of books on a single GPU. Key design decisions include using fixed-size arrays for order storage, avoiding continuous sorting, and leveraging JAX's vmap operator to parallelize message processing, which transforms conditional branches into select statements that execute all branches at runtime. The simulator supports limit orders, cancellations, and market orders with price-time priority matching, and can replay historical data from sources like LOBSTER.

The authors integrate JAX-LOB with the gymnax RL framework and PureJaxRL to create an end-to-end GPU-based training pipeline for optimal trade execution. The execution environment features a continuous action space with four price levels (far touch, mid-price, near touch, passive) and an observation space including intra-step L1 data, time, and execution progress. Using recurrent PPO, they demonstrate a 7x speedup in training (550 vs 74 steps per second) compared to CPU-based simulators on the same hardware. The work also provides detailed benchmarks of basic operations (add, cancel, match) and message processing times, showing that vmap eliminates runtime variance across message types. JAX-LOB is open-sourced and promises to unlock new research in agent-based models, generative models, and RL for trading by enabling massive parallel simulation and potentially differentiable LOB models for calibration.