XMob: Accelerated Differentiable Mesoscopic Traffic Simulation in JAX

This paper introduces XMob, a novel differentiable traffic simulation framework built in JAX to advance traditional models like SUMO’s mesoscopic simulator. By leveraging JAX’s capabilities for vectorized, hardware-accelerated computation (GPU/TPU), XMob achieves orders-of-magnitude speedups, enabling large-scale urban network simulations and extensive counterfactual analyses. A key innovation is XMob’s inherent differentiability, facilitating direct integration with gradient-based optimization for tasks such as demand calibration and network parameter estimation, significantly outperforming black-box approaches. Furthermore, XMob can be used in Physics-Informed Machine Learning (PIML) pipelines to enhance data-driven augmentation, embedding domain principles like flow conservation and shockwave theory. This ensures physically plausible and robust predictions, even for unobserved scenarios such as lane modifications. The hybrid architecture, combining a deterministic JAX core with incremental machine learning, offers a scalable and efficient solution for modern traffic simulation and optimization challenges.

• Algorithms and Theory