DES-GYMNAX: Fast Discrete-Event Simulator in JAX

In this work, we introduce DES-Gymnax, a novel high-performance discrete-event simulator implemented in JAX. By leveraging the just-in-time compilation, automatic vectorization, and GPU acceleration capabilities of JAX, DES-Gymnax can achieve 10x to 100x times performance improvement over traditional Python-based discrete-event simulators like Salabim. The proposed DES-Gymnax can feature a Gym-like API that facilitates seamless integration with reinforcement learning algorithms, addressing a critical gap between simulation engines and AI techniques. DES-Gymnax is validated on three benchmark models, i.e., an M/M/1 queue, a multi-server model, and a tandem queue model. Experimental results demonstrate that DES-Gymnax maintains simulation accuracy while significantly reducing execution time, enabling efficient large-scale sampling crucial for reinforcement learning applications in operations research areas. The open-source code is available in the DES-Gymnax repository (Yun, Jun, and Xiangfeng 2025).