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

In one sentence

Summary

The paper addresses the computational bottleneck in training reinforcement learning agents for high-frequency trading tasks, specifically optimal trade execution, by introducing JAX-LOB, a GPU-enabled limit order book simulator. Traditional CPU-based simulators struggle with the low signal-to-noise ratio and the need for massive parallelization required to overcome overfitting and capture realistic market dynamics. The authors position this tool as a critical infrastructure component for agent-based models and RL environments, aiming to unlock scalable research into market microstructure and execution strategies without compromising the realism of LOB mechanisms. The core method involves implementing the LOB mechanics using JAX, leveraging its just-in-time compilation and automatic vectorization to process thousands of order books in parallel via the vmap operator. The architecture uses fixed-size arrays to represent bid and ask sides, avoiding the overhead of dynamic sorting during message processing. The simulator is integrated with the Gymnax RL framework and PureJaxRL, wrapping the environment to allow both experience rollout and network updates to occur entirely on the GPU. Experiments compare JAX-LOB against CPU-based implementations (including RL4MM and a custom linked-list simulator) using LOBSTER data for stocks like TSLA and AMZN, measuring processing times per message and training steps per second. Results demonstrate a significant performance advantage, with JAX-LOB achieving up to 75x faster per-message processing and a 7x speedup in RL training steps compared to CPU baselines. The paper provides a preliminary example of training a Recurrent PPO agent for optimal execution, showing it can outperform a TWAP benchmark in-sample, though out-of-sample validation is noted as future work. Limitations include the current lack of strategic agent interaction in the replay data (only direct market impact is captured) and the need for rigorous hyperparameter tuning and broader testing to establish robustness. The deliverable is an open-source framework designed to facilitate large-scale RL research in trading.

Research problem

The paper tackles the inefficiency of current limit order book (LOB) simulators in supporting large-scale reinforcement learning for trading. Specifically, it addresses the need for high-performance, parallelizable simulation environments that can handle the massive data throughput required for training RL agents in high-frequency trading contexts, such as optimal trade execution, without the communication bottlenecks inherent in CPU-GPU data transfer.

Core contributions

• Introduction of JAX-LOB, the first GPU-accelerated LOB simulator designed for parallel processing of thousands of books.
• Development of a Gymnax environment wrapper that enables end-to-end reinforcement learning training entirely on GPU, eliminating CPU-GPU transfer overhead.
• Demonstration of at least 7x speedup in RL training steps compared to equivalent CPU-based LOB simulators.
• Provision of a concrete implementation for optimal trade execution using Recurrent PPO, serving as a benchmark for future RL trading research.
• Open-sourcing of the JAX-LOB codebase and integration examples to facilitate community adoption.

Method

• Implements LOB state using fixed-size JAX arrays for bid and ask sides to enable JIT compilation.
• Uses the vmap operator to vectorize message processing across thousands of parallel order books.
• Designs a matching logic that avoids continuous sorting to reduce computational overhead during parallel execution.
• Integrates with Gymnax to create a differentiable RL environment where state transitions are GPU-native.
• Wraps the environment with PureJaxRL to facilitate end-to-end training of Recurrent PPO agents.
• Defines a reward function combining VWAP advantage and price drift to guide the execution agent.
• Pre-loads LOBSTER data into fixed-size time windows to optimize memory access patterns.

Results & findings

• JAX-LOB processes messages up to 75x faster per message than CPU baselines when parallelizing thousands of books.
• RL training speed is improved by at least 7x (550 vs 74 steps per second) compared to CPU implementations.
• The GPU-only pipeline avoids the latency of data transfer between CPU and GPU during network updates.
• The Recurrent PPO agent trained on JAX-LOB shows an episodic return that outperforms the TWAP benchmark in the limited in-sample test.
• Processing time for matching operations scales with order book capacity and incoming order size, but remains efficient under parallelization.
• The vmap operator forces execution of all conditional branches, making runtime dependent on the slowest branch (matching) rather than message type.

Data & experiments

• Utilizes LOBSTER high-frequency message data for stocks including TSLA and AMZN.
• Compares processing times of JAX-LOB against CPU-based implementations (RL4MM and a custom linked-list simulator) on Apple M1 and AMD EPYC hardware.
• Evaluates training speed (steps per second) of a Recurrent PPO agent using JAX-LOB on an Nvidia GPU versus a CPU-based Gym environment.
• Tests parallelization efficiency by processing 10 to 10,000 environments simultaneously on the GPU.
• Conducts preliminary in-sample training on a single day of data to validate the RL pipeline.

Limitations & caveats

• The RL example is trained on a single day of data, limiting the generalizability of the performance claims.
• The simulation uses historical replay data, which lacks strategic agent reactions to the RL agent's actions (no indirect market impact).
• The fixed-size array implementation requires pre-allocating memory for the maximum book capacity, which may be inefficient for sparse books.
• The vmap operator executes all conditional branches, potentially wasting compute on paths not taken in specific messages.
• Out-of-sample testing and hyperparameter tuning are identified as necessary future work for robust validation.

Curator notes

Discovery status: institutional_trading_round2_promoted. Source links: StockSim: A Dual-Mode Order-Level Simulator for Evaluating Multi-Agent LLMs in Financial Markets