UNsim: Differentiable Traffic Simulator

• UNsim is an open-source Python/JAX traffic simulator based on the Link Transmission Model and DUO route-choice, enabling gradient-based optimization on realistic urban networks.
• It implements a fully differentiable simulation using arithmetic and piecewise-linear operations, which facilitates efficient calibration, sensitivity analysis, and optimal control without surrogate smoothing.
• The system’s end-to-end JAX-traced pipeline supports large-scale traffic modeling, as demonstrated by its successful Chicago-Sketch congestion pricing experiment that reduced travel time significantly.

UNsim is an open-source Python/JAX traffic simulator implementing an end-to-end differentiable network loading model based on the Link Transmission Model (LTM), a suite of node models, and a Dynamic User Optimum (DUO) route-choice model. Its stated objective is to expose a rich, physically-grounded network loading model—and its coupling with dynamic route choice—to automatic differentiation (AD), so that exact gradients of trajectory-level or network-level objectives with respect to model inputs can be computed in a single backward pass regardless of the number of decision variables. In the formulation reported for the Chicago-Sketch network, this design enables large-scale gradient-based optimization on a realistic urban network with around 2500 links and roughly 1 million vehicles (Seo, 13 Apr 2026).

1. Concept and system architecture

At its core, UNsim combines three components inside a single JAX-traced time evolution. The first is an LTM for first-order macroscopic traffic dynamics under a triangular fundamental diagram. The second is a suite of node models—origin, destination, dummy, diverge, merge, and the incremental node model for general junctions—written entirely through arithmetic and piecewise-linear $\min/\max$ operations. The third is a DUO route-choice model, optionally smoothed by a logit-based turning probability, that tracks per-destination cumulative counts. All components are written in pure Python, vectorized over links and nodes, and composed inside a single JAX scan over time (Seo, 13 Apr 2026).

This architecture is designed around differentiability rather than around post hoc sensitivity analysis. The intended consequence is that objectives such as total travel time, link delays, probe-vehicle travel times, and toll sensitivities can be differentiated directly with respect to quantities such as fundamental diagram parameters, origin–destination demands, merge priorities, time-dependent tolls, and turning fractions. The paper presents this as a way to make large-scale gradient-based calibration, sensitivity analysis, and optimal control practical on realistic urban networks.

2. Link Transmission Model formulation

UNsim’s LTM is Newell’s simplified kinematic wave formulation under a triangular fundamental diagram. On link $l$, the fundamental diagram is

$$Q_l(k)= \begin{cases} u_l\,k & \text{if } k \le k_l^* \ w_l\,(\kappa_l-k) & \text{if } k>k_l^* \end{cases}$$

with critical density $k_l^*=q_l^*/u_l$ and backward wave speed $w_l=q_l^*/(\kappa_l-k_l^*)$. The primary state variables are cumulative vehicle counts at the upstream and downstream boundaries, $N_{l,U}(t)$ and $N_{l,D}(t)$. By Newell’s formula, the cumulative count at any interior point $(t,x)$ is

$$N_l(t,x)=\min \left\{ N_{l,U}(t-x/u_l),\; N_{l,D}(t-(d_l-x)/w_l)+\kappa_l(d_l-x) \right\}.$$

Demand and supply are likewise written without logical branching:

$$D_l(t)=\min\left\{ \frac{N_{l,U}(t+\Delta t-d_l/u_l)-N_{l,D}(t)}{\Delta t},\; q_l^*,\; q_l^{out} \right\},$$

$l$0

After node inflows and outflows are determined, the cumulative counts update as

$l$1

A central claim of UNsim is that these piecewise-linear $l$2 and $l$3 operations admit subgradients almost everywhere, so no smooth relaxation is required for AD. The paper explicitly states that the kink points form a measure-zero set and do not impede AD in practice (Seo, 13 Apr 2026).

3. Dynamic route choice and per-destination flow tracking

UNsim couples the LTM with an endogenous route-choice model. At regular intervals it computes instantaneous link travel times $l$4 either from average density,

$l$5

or from a segment-level approximation in which the link is divided into $l$6 cells, $l$7 is evaluated at the segment points, densities $l$8 are inferred, and travel time is accumulated as $l$9. A reverse Bellman–Ford on the resulting $$Q_l(k)= \begin{cases} u_l\,k & \text{if } k \le k_l^* \ w_l\,(\kappa_l-k) & \text{if } k>k_l^* \end{cases}$$0 yields, for each node $$Q_l(k)= \begin{cases} u_l\,k & \text{if } k \le k_l^* \ w_l\,(\kappa_l-k) & \text{if } k>k_l^* \end{cases}$$1 and destination $$Q_l(k)= \begin{cases} u_l\,k & \text{if } k \le k_l^* \ w_l\,(\kappa_l-k) & \text{if } k>k_l^* \end{cases}$$2, the next-link indicator $$Q_l(k)= \begin{cases} u_l\,k & \text{if } k \le k_l^* \ w_l\,(\kappa_l-k) & \text{if } k>k_l^* \end{cases}$$3 (Seo, 13 Apr 2026).

Let $$Q_l(k)= \begin{cases} u_l\,k & \text{if } k \le k_l^* \ w_l\,(\kappa_l-k) & \text{if } k>k_l^* \end{cases}$$4 denote the number of vehicles bound for destination $$Q_l(k)= \begin{cases} u_l\,k & \text{if } k \le k_l^* \ w_l\,(\kappa_l-k) & \text{if } k>k_l^* \end{cases}$$5 on link $$Q_l(k)= \begin{cases} u_l\,k & \text{if } k \le k_l^* \ w_l\,(\kappa_l-k) & \text{if } k>k_l^* \end{cases}$$6. The DUO diverge ratio at node $$Q_l(k)= \begin{cases} u_l\,k & \text{if } k \le k_l^* \ w_l\,(\kappa_l-k) & \text{if } k>k_l^* \end{cases}$$7 toward outlink $$Q_l(k)= \begin{cases} u_l\,k & \text{if } k \le k_l^* \ w_l\,(\kappa_l-k) & \text{if } k>k_l^* \end{cases}$$8 is

$$Q_l(k)= \begin{cases} u_l\,k & \text{if } k \le k_l^* \ w_l\,(\kappa_l-k) & \text{if } k>k_l^* \end{cases}$$9

where $k_l^*=q_l^*/u_l$0 for intermediate nodes, and $k_l^*=q_l^*/u_l$1 at origins. Destination-specific outflows then follow FIFO:

$k_l^*=q_l^*/u_l$2

The route-choice mapping remains differentiable almost everywhere because the diverge ratios are continuous functions of per-destination vehicle counts and all filters remain arithmetic or $k_l^*=q_l^*/u_l$3 operations. To provide nonzero derivatives even when the shortest-path indicator $k_l^*=q_l^*/u_l$4 does not change under small cost perturbations, UNsim also offers a logit smoothing:

$k_l^*=q_l^*/u_l$5

The paper describes this as preserving the network-consistent DUO baseline while ensuring smooth gradients with respect to link costs (Seo, 13 Apr 2026).

4. Reverse-mode automatic differentiation

UNsim implements the full simulator as a single JAX scan over $k_l^*=q_l^*/u_l$6 timesteps. The one-step transition is written as $k_l^*=q_l^*/u_l$7, where the step computes demands and supplies via the LTM equations, node flows through the node models and fixed-$k_l^*=q_l^*/u_l$8 incremental node model, per-destination splits, and cumulative-count updates. The full forward simulation is then $k_l^*=q_l^*/u_l$9, a scalar objective is defined as $w_l=q_l^*/(\kappa_l-k_l^*)$0, and grad(J) returns $w_l=q_l^*/(\kappa_l-k_l^*)$1 in one backward pass (Seo, 13 Apr 2026).

The paper reports several computational properties of this construction. Once the forward pass is JIT-compiled, reverse-mode AD computes exact subgradients almost everywhere at a cost only $w_l=q_l^*/(\kappa_l-k_l^*)$2–$w_l=q_l^*/(\kappa_l-k_l^*)$3 that of one forward run, independent of the total number of parameters. The time complexity is $w_l=q_l^*/(\kappa_l-k_l^*)$4, and the memory cost is reported as $w_l=q_l^*/(\kappa_l-k_l^*)$5 via checkpointing. Because the entire computation graph is static, with no Python if branches inside $w_l=q_l^*/(\kappa_l-k_l^*)$6, the implementation avoids the usual need to derive model-specific adjoint equations by hand.

A common misconception in differentiable simulation is that nonsmooth traffic primitives must be replaced by smooth surrogates before gradient-based optimization becomes possible. UNsim argues against that requirement for this model class: because the network loading is built from arithmetic and piecewise-linear $w_l=q_l^*/(\kappa_l-k_l^*)$7 operations on continuous aggregate state variables, AD can operate directly on the original formulation rather than on a smoothed approximation.

5. Chicago-Sketch congestion pricing experiment

The principal demonstration in the paper is a dynamic congestion toll optimization on the Chicago-Sketch dataset. The network contains 927 nodes, 2 557 links, approximately 1.17 million vehicles, and 17 963 OD pairs. The decision variables are tolls on 383 congested links over 40 time intervals, giving 15 320 variables. The optimization objective is to minimize total travel time plus $w_l=q_l^*/(\kappa_l-k_l^*)$8, subject to $w_l=q_l^*/(\kappa_l-k_l^*)$9 (Seo, 13 Apr 2026).

Optimization is performed with Adam using $N_{l,U}(t)$0, $N_{l,U}(t)$1, $N_{l,U}(t)$2, and gradient clipping at $N_{l,U}(t)$3, for 10 000 iterations. The reported total runtime is 8 527 s, corresponding to approximately $N_{l,U}(t)$4 s per forward-plus-backward pass; the abstract summarizes this as about 2 hours for 10 000 iterations, with one simulation run and gradient derivation taking about $N_{l,U}(t)$5 s. The simulated horizon is 3 hours with $N_{l,U}(t)$6 s, implying $N_{l,U}(t)$7 steps. The result reported is a drop in total travel time from $N_{l,U}(t)$8 veh·hr to $N_{l,U}(t)$9 veh·hr, a reduction of $N_{l,D}(t)$0. This performance was achieved on a single NVIDIA GH200 GPU with 96 GB memory, using Python/JAX parallelism over links and nodes (Seo, 13 Apr 2026).

In the context of the paper’s broader claims, this experiment is the main evidence that the differentiable formulation scales beyond toy networks. A plausible implication is that the method is particularly suited to high-dimensional control problems in which derivative-free methods would be too slow or too sample-inefficient.

6. Software organization, workflow, and research uses

UNsim is released as an open-source Python package at github.com/toruseo/UNsim. Its core modules are unsim.network, which builds a directed graph with FD parameters, OD demands, and merge priorities; unsim.ltm, which implements LTM link update functions; unsim.nodes, which contains vectorized node models; unsim.route, which provides DUO and logit-DUO route-choice routines; and unsim.engine, which executes the JAX scan and provides utilities for defining objectives (Seo, 13 Apr 2026).

The minimal workflow in the paper consists of loading a network, defining a scalar loss on the simulator state, obtaining a JIT-compiled gradient with grad(loss), and iteratively updating the decision variables with an optimizer. Because the simulator is fully differentiable, the same pipeline can be used not only for toll optimization but also for calibration, sensitivity analysis, and optimal control tasks such as congestion pricing, ramp metering, and signal timing. The paper further notes that neural-net modules, including learned OD patterns and learned turning-fraction surrogates, may be embedded inside the one-step transition and trained end-to-end against observed data via the same AD pipeline.

Within the nomenclature of recent simulation literature, UNsim refers specifically to the traffic simulator of Seo and collaborators (Seo, 13 Apr 2026). This is distinct from unrelated systems with similar names, such as the closed-loop neural sensor simulator “UniSim” for self-driving vehicles (Yang et al., 2023) and the multimodal perceptual-metric framework “UniSim” discussed in the vision-language literature (Ghazanfari et al., 2024).