JaxMARL-HFT: GPU-Accelerated Large-Scale Multi-Agent Reinforcement Learning for High-Frequency Trading

TL;DR

This work introduces JaxMARL-HFT, the first GPU-accelerated open-source MARL environment for high-frequency trading on market-by-order data, built atop JAX-LOB and JaxMARL to support heterogeneous agent types. It achieves up to $240\times$ faster end-to-end training, enabling large-scale hyperparameter sweeps and processing of granular LOBSTER data, which makes MARL research in HFT feasible at scale. Using Independent PPO in a two-agent setup (market making and order execution) with one year of AMZN LOBSTER data, the framework yields learned policies that often surpass baseline strategies, illustrating the practical potential of MARL in realistic, data-rich trading environments. The codebase is open-source and designed to be extended to more agents and tasks, paving the way for robust, scalable MARL research in financial markets.

Abstract

Agent-based modelling (ABM) approaches for high-frequency financial markets are difficult to calibrate and validate, partly due to the large parameter space created by defining fixed agent policies. Multi-agent reinforcement learning (MARL) enables more realistic agent behaviour and reduces the number of free parameters, but the heavy computational cost has so far limited research efforts. To address this, we introduce JaxMARL-HFT (JAX-based Multi-Agent Reinforcement Learning for High-Frequency Trading), the first GPU-accelerated open-source multi-agent reinforcement learning environment for high-frequency trading (HFT) on market-by-order (MBO) data. Extending the JaxMARL framework and building on the JAX-LOB implementation, JaxMARL-HFT is designed to handle a heterogeneous set of agents, enabling diverse observation/action spaces and reward functions. It is designed flexibly, so it can also be used for single-agent RL, or extended to act as an ABM with fixed-policy agents. Leveraging JAX enables up to a 240x reduction in end-to-end training time, compared with state-of-the-art reference implementations on the same hardware. This significant speed-up makes it feasible to exploit the large, granular datasets available in high-frequency trading, and to perform the extensive hyperparameter sweeps required for robust and efficient MARL research in trading. We demonstrate the use of JaxMARL-HFT with independent Proximal Policy Optimization (IPPO) for a two-player environment, with an order execution and a market making agent, using one year of LOB data (400 million orders), and show that these agents learn to outperform standard benchmarks. The code for the JaxMARL-HFT framework is available on GitHub.

Figures (5)

• Figure 1: The two levels of parallelisation on the GPU with vectorised mapping (vmap) in JAX.
• Figure 2: Speed Comparison of the Reinforcement Learning Training Pipeline
• Figure 3: Training curves using a proxy reward to aid comparison across reward functions. MARL experiment with a market making and an execution agent.
• Figure 4: Detailed plots considering the behaviour of the learned policies for each agent. Results are given for policies learned using the FixedQuant action space with the Spooner reward function and a Quadratic inventory penalty for the market maker.
• Figure 5: Evaluation of the learned agents when facing baseline implementations (TWAP and AvSt avellaneda2008high). The bottom left has both agents play a baseline policy, whilst the top right has both play a learnt policy. The bottom triangles represent the market maker portfolio value, whilst the top represent the slippage experienced by the execution agent. The learnt policies improve over the baseline, and the execution agent performs worse when facing a learnt market making policy than when facing the baseline.