Learning Traffic Flow Dynamics with Stochastic Physics-Informed Neural Cellular Automata

Federica Bragone; Matthieu Barreau

cs.LGOct 7, 2026

Learning Traffic Flow Dynamics with Stochastic Physics-Informed Neural Cellular Automata

Authors: Federica Bragone, Matthieu Barreau

Organizations: Department of Decision and Control Systems, KTH Royal Institute of Technology, Stockholm, Sweden

Abstract

Traffic flow modeling is essential for understanding and predicting the collective dynamics of vehicles on road networks. Cellular automata provide a simple, interpretable yet powerful framework for representing these dynamics via local interaction rules, while retaining the ability to reproduce complex macroscopic traffic phenomena. However, learning local transition rules from data while preserving physically meaningful constraints remains challenging, particularly for stochastic models. In this work, we propose a physics-informed neural cellular automaton (PI-NCA) for data-driven traffic flow modeling. Building on the standard neural cellular automaton (NCA), we design a neural architecture that is physically consistent with the road topology and guarantees conservation of the total number of vehicles, thereby constraining the learned transition rules to physically admissible dynamics. We further extend this framework to stochastic dynamics by parameterizing probabilistic transition rules while preserving the same physics-informed constraints. We evaluate the proposed models on multiple traffic scenarios generated by the well-established Nagel-Schreckenberg and Kerner-Klenov-Wolf cellular automata. The results demonstrate that the PI-NCA successfully learns the dynamics of both traffic models and consistently outperforms a standard NCA, while the stochastic extension captures probabilistic transition rules without compromising the imposed physical constraints.

Figures & tables

Figure 1 : Traffic Cellular Automaton. Each cell contains the velocity of the car (

0

, …,

v_{\max}

) if there is a vehicle or

-1

if it is empty.

Figure 2 : An example of a one-step transition for the Nagel–Schreckenberg TCA. Colors indicate the vehicle’s speed in the corresponding cell. Each line corresponds to an operation: acceleration, slowing down, random deceleration, and movement to get from

\mathbf{s}_{t}

\mathbf{s}_{t+1}

Panel of Figure 3 — Figure 3 : Fundamental diagrams for the NaSch model. In (a), a point represents a local pair density–flow while the color refers to the vehicle velocity. In (b), two initial conditions have been tested, producing free-flow and congested regimes.

Figure 4 : Convolutional Neural Network structure for NCA, PI-NCA, and Stochastic PI-NCA with 2 hidden layers,

16

channels, and

v_{max}=5

. The operator

\mathcal{M}

represents the movement operation as described in Equation 12 .

Property	NCA	PI-NCA
Input states	$\mathbf{s}_{t-1}$	$\mathbf{s}_{t-1}$
Input channels $C$	$1$	$1$
Prediction target	State at $t$	Velocity at $t$
Loss	MSE	Masked MSE
Loss evaluated on	All cells	Occupied cells
Explicit update	No	Yes

Table 1 : Overview of the two deterministic variants: NCA and PI-NCA.

Figure 5 : Conceptual progression from deterministic to stochastic PI-NCA. The stochastic formulation extends the deterministic velocity prediction to a categorical velocity distribution while retaining the explicit movement update.

Figure 6 : Six initial states used for training and testing.

Parameter	Training	Testing
Road length $L$	100	50, 100, 150
Number of vehicles $N$	10, 15, 20	8, 10, 12, 15, 17, 20
Maximm velocity $v_{\text{max}}$	5	5
Trajectory length $T$	40	40
Initialization types	6	6

Table 2 : Data-generation settings using the NaSch model for training and testing.

Hyperparameter	Value
Optimizer	Adam
Learning rate	$10^{-3}$
Weight decay	$10^{-5}$
Batch size	64
Hidden channels	16
Kernel size	11

Table 3 : NCA and PI-NCA training configuration for deterministic and stochastic NaSch experiments.

Panel of Figure 7 — Figure 7 : Fundamental diagrams for KKW-1 model.

Traffic CA	Model	MAE	Car MAE	CE
NaSch	NCA	$0.0022\pm 0.0015$	$0.0070\pm 0.0062$	$0.0071\pm 0.0037$
NaSch	PI-NCA	$0.00006\pm 0.00011$	$0.00009\pm 0.00015$	$0.00018\pm 0.00040$
KKW-1	NCA	$0.1949\pm 0.0289$	$0.4667\pm 0.0768$	$0.3085\pm 0.0839$
KKW-1	PI-NCA	$0.0179\pm 0.0090$	$0.0391\pm 0.0187$	$0.00017\pm 0.00038$

Table 4 : Overall prediction performance over the complete NaSch and KKW-1 test sets.

Figure 8 : Comparison of NaSch (left) with the proposed NCA (middle) and their absolute error (right), using $L=100$ , 20 vehicles, $T=40$ , and random initialization.

Panel of Figure 9 — Figure 9 : Generalization across traffic initializations for the deterministic NCA (blue) and PI-NCA (orange) with NaSch and KKW-1 dynamics. (a) Mean absolute error and (b) conservation error. Each point corresponds to an independent training run.

Panel of Figure 10 — Figure 10 : Generalization across road lengths for the deterministic NCA (blue) and PI-NCA (orange) with NaSch and KKW-1 dynamics. (a) Mean absolute error and (b) conservation error for $L\in\{50,100,150\}$ . Shaded regions indicate road lengths not used during training. Each point corresponds to an independent training run.

Panel of Figure 11 — Figure 11 : Generalization across vehicle numbers for the deterministic NCA (blue) and PI-NCA (orange) with NaSch and KKW-1 dynamics. (a) Mean absolute error and (b) conservation error. Shaded regions indicate vehicle numbers not used during training. Each point corresponds to an independent training run.

Panel of Figure 12 — Figure 12 : Representative stochastic space-time trajectories generated by the reference traffic CA and the stochastic PI-NCA. (a) NaSch with $p=0.20$ , $L=100$ , $\rho=0.30$ , $T=100$ , and random initialization; (b) KKW-1 using the stochastic parameters in Table 5 with $L=100$ , $\rho=0.30$ , $T=100$ , and random initialization.

Panel of Figure 13 — Figure 13 : Local traffic behavior generated by the stochastic PI-NCA for NaSch and KKW-1 dynamics. (a)-(c) Local fundamental diagrams for NaSch and KKW-1, respectively. (b)-(d) Global fundamental diagrams for NaSch and KKW-1, respectively.

Appendix figures & tables8 assets

Supplementary material from the paper’s appendix.

Appendix

Parameter	Description	Value
$d$	Vehicle length	1
$v_{\mathrm{max}}$	Maximum vehicle velocity	5
$a$	Vehicle acceleration	1
$\tau$	Time-step duration	1
$k$	Synchronization parameter	2.04
$v_{p}$	Velocity threshold for random acceleration	2.33

Appendix

Table 5: Parameters used for the adapted KKW-1 model.

Traffic CA	Model	MSE	OA	RCE
NaSch	NCA	$0.0105\pm 0.0089$	$0.9995\pm 0.0003$	$0.0004\pm 0.0002$
NaSch	PI-NCA	$0.0002\pm 0.0003$	$0.9999\pm 0.0000$	$0.0000\pm 0.0000$
KKW-1	NCA	$0.7810\pm 0.1274$	$0.9472\pm 0.0071$	$0.0194\pm 0.0050$
KKW-1	PI-NCA	$0.0655\pm 0.0320$	$0.9946\pm 0.0028$	$0.00002\pm 0.00005$

Appendix

Table 6 : Overall prediction performance over the complete NaSch and KKW-1 test sets.

Init.	Traffic CA	Model	MAE	CE
$I_{1}$	NaSch	NCA	$0.0000\pm 0.0000$	$0.0000\pm 0.0000$
	NaSch	PI-NCA	$0.0000\pm 0.0000$	$0.0000\pm 0.0000$
	KKW-1	NCA	$0.0386\pm 0.0509$	$0.1350\pm 0.1154$
	KKW-1	PI-NCA	$0.0000\pm 0.0000$	$0.0000\pm 0.0000$
$I_{2}$	NaSch	NCA	$0.0000\pm 0.0000$	$0.0000\pm 0.0000$
	NaSch	PI-NCA	$0.0000\pm 0.0000$	$0.0000\pm 0.0000$

Appendix

Table 7 : Prediction performance across the different initialization strategies.

$L$	Traffic CA	Model	MAE	CE
50	NaSch	NCA	$0.0025\pm 0.0008$	$0.0160\pm 0.0076$
	NaSch	PI-NCA	$0.0002\pm 0.0003$	$0.0005\pm 0.0012$
	KKW-1	NCA	$0.4072\pm 0.0675$	$0.7136\pm 0.2020$
	KKW-1	PI-NCA	$0.0421\pm 0.0235$	$0.0005\pm 0.0012$
100	NaSch	NCA	$0.0031\pm 0.0032$	$0.0032\pm 0.0047$
	NaSch	PI-NCA	$0.0000\pm 0.0000$	$0.0000\pm 0.0000$

Appendix

Table 8 : Prediction performance for different road lengths

L

. The values

L\in\{50,150\}

are out-of-distribution.

$N$	Traffic CA	Model	MAE	CE
8	NaSch	NCA	$0.0002\pm 0.0005$	$0.0000\pm 0.0000$
	NaSch	PI-NCA	$0.0000\pm 0.0000$	$0.0000\pm 0.0000$
	KKW-1	NCA	$0.0522\pm 0.0140$	$0.0479\pm 0.0131$
	KKW-1	PI-NCA	$0.0051\pm 0.0050$	$0.0010\pm 0.0023$
10	NaSch	NCA	$0.0008\pm 0.0015$	$0.0011\pm 0.0024$
	NaSch	PI-NCA	$0.0000\pm 0.0000$	$0.0000\pm 0.0000$

Appendix

Table 9 : Prediction performance for seen and unseen numbers of vehicles. The values

\{8,12,17\}

are out-of-distribution.

Figure 14 : Average probability assigned to each next velocity class for

p\in\{0.1,0.2,\dots\}

, comparing the reference NaSch distribution with the learned PI-NCA distribution.

Figure 15 : Total variation distance between the reference and predicted distributions. Points denote independent training runs.

Panel of Figure 16 — Figure 16 : Mean categorical next velocity probabilities predicted by PI-NCA as a function of local density for (a) NaSch and (b) KKW-1, respectively.

Explore similar work

An Empirical Markov Chain Car-Following (MC-CF) Model
Date pendingSungyong Chung, Yanlin Zhang, Nachuan Li +2Microscopic Traffic Simulations
Department of Civil and Environmental Engineering, University of Illinois Urbana-Champaign · Northwestern University Transportation Center, Northwestern University
CIWI-CKT: Chaos-Informed Wave Interference Feature Fusion and Cross-City Knowledge Transfer for Traffic Flow Forecasting
Jun 14, 2026Abdul Joseph Fofanah, Lian Wen, David Chen +1Accurate Traffic Forecasting Traffic
School of Information and Communication Technology, Griffith University, Brisbane, 4111, Australia · School of Information Engineering, Chang’an University, Xi’an, China
SpinFlow: A Physics-Informed Spin Field Framework for Traffic Phase Inference and Transition Detection
May 22, 2026Haopeng Deng, Fucheng Zheng, Xinhai XiaTraffic Congestion
School of Future Transportation, Guangzhou Maritime University, Guangzhou, China

Date pending, eess.SYJ/K move · Enter open · S save

An Empirical Markov Chain Car-Following (MC-CF) Model

Sungyong Chung, Yanlin Zhang, Nachuan Li, Dana Monzer+1

Department of Civil and Environmental Engineering, University of Illinois Urbana-Champaign · Northwestern University Transportation Center, Northwestern University

Microscopic Traffic Simulations

Open paper AskPDF arXiv

Car-following behavior is fundamental to traffic flow theory, yet traditional models often fail to capture the stochasticity of naturalistic driving. This paper proposes an empirical probabilistic sampling approach to car-following modeling that bypasses conventional parametric assumptions. Under this approach, we introduce the Markov Chain Car-Following (MC-CF) model, which represents state transitions as a Markov process and predicts behavior by randomly sampling accelerations from empirical distributions within discretized state bins. Evaluation on the Waymo Open Motion Dataset (WOMD) demonstrates that MC-CF variants significantly outperform all physics-based baselines (IDM, Gipps, FVDM, and SIDM) across both one-step and open-loop trajectory prediction metrics, and remain competitive with modern data-driven baselines including neural network and Gaussian mixture model approaches. Zero-shot generalization on the Naturalistic Phoenix (PHX) dataset further confirms cross-domain transferability. Finally, microscopic ring road simulations validate the framework's scalability: by incrementally integrating unconstrained free-flow trajectories and high-speed freeway data (TGSIM) alongside a conservative inference strategy, the model substantially reduces collisions across most tested scenarios and successfully reproduces naturalistic and stochastic shockwave propagation, though crashes persist under severe shockwave conditions. Overall, the proposed MC-CF model provides a robust and scalable foundation for simulating population-level stochastic traffic behavior that requires no behavioral parameter calibration, making it well-suited for the data-rich future of intelligent transportation.

Abstract

Explore similar work

An Empirical Markov Chain Car-Following (MC-CF) Model

CIWI-CKT: Chaos-Informed Wave Interference Feature Fusion and Cross-City Knowledge Transfer for Traffic Flow Forecasting

SpinFlow: A Physics-Informed Spin Field Framework for Traffic Phase Inference and Transition Detection