JaxWildfire: A GPU-Accelerated Wildfire Simulator for Reinforcement Learning

AI summary

Overview

• Research area: Reinforcement learning for natural hazard management, specifically wildfire suppression, plus GPU-accelerated environmental simulation.
• Technical level: Advanced (assumes familiarity with reinforcement learning, cellular automata, JAX, and gradient-based optimization).
• Scope: The paper introduces JaxWildfire, a JAX-based cellular automata wildfire simulator built for high-throughput, differentiable simulation and RL training, benchmarked against PyTorchFire and validated on a real wildfire and a synthetic suppression task.

What This Paper Is About

Training reinforcement learning agents for wildfire management requires millions of environment interactions, but existing wildfire simulators are too slow, and even faster cellular automata simulators must shuttle data back and forth between the simulator and the RL algorithm in main memory. The authors build JaxWildfire, a wildfire simulator written in JAX that runs entire vectorized simulations on the GPU, so that simulation and RL training can happen in the same place. The goal is to make wildfire environments fast enough, and differentiable enough, to be practical for RL research and for tuning simulator parameters against real fire data.

Key Contributions

1. A JAX-native cellular automata wildfire simulator whose spread model is built on Alexandridis et al. (2008) and Xia and Cheng (2025), but replaces the earlier ad-hoc tanh squashing function with a Poisson-process formulation that converts fire propagation potential into a bounded, differentiable ignition probability.
2. Large speed improvements over the closest comparable framework, PyTorchFire, with a reported 6-35x speedup for single simulations and 18-51x when using the for_i routine, plus vectorized vmap simulations that reach a peak throughput reported as 2.5^16 steps/s with 64 environments on a single NVIDIA A100 GPU at grid size 64x64.
3. Support for real-world and synthetic environments, including the European Space Agency's WorldCover dataset and Digital Elevation Model (DEM), where a place-name string (for example "Cape Town, South Africa") is geocoded into a bounding box and used to download the relevant tiles. Data can also be exported to PyTorch format for benchmarking machine learning models of fire spread.
4. Demonstrations of both calibration and RL training, using gradient-based optimization to calibrate the simulator against the real Bear 2020 wildfire, and training a PPO air tanker agent in a synthetic 20x20 environment to extinguish fires.

Main Findings

• Simulation speed: Benchmarked on an NVIDIA A100 GPU across different grid sizes and 5 random seeds, JaxWildfire achieves a 6-35x speedup over PyTorchFire for single simulations. This rises to 18-51x when using for_i, which does not track history and is therefore more efficient.
• Vectorized throughput: Using JAX's vmap, JaxWildfire runs multiple vectorized simulations on the GPU, reaching a peak throughput reported as 2.5^16 steps/s with 64 environments on a single A100 GPU at grid size 64x64. The paper states the framework benefits from increasing degrees of parallelism.
• Calibration accuracy on real data: Following the setup of Xia and Cheng (2025) with a combined binary cross-entropy (BCE) loss for fire spread shape and a mean squared error (MSE) loss after 2D average pooling for fire scale, JaxWildfire was calibrated on the Bear 2020 wildfire and compared to PyTorchFire. The simulators achieve comparable results measured by Intersection over Union, with a slight advantage for JaxWildfire. Results were averaged over 5 seeds, with variance reported as negligible. JaxWildfire used the optax library with Adam and the raw probabilities routine.
• RL suitability: A PPO air tanker agent trained in a synthetic 20x20 environment learned a fire suppression policy, successfully moving toward the fire and controlling the valve to extinguish it. The agent receives partial egocentric observations, controls the valve (open/closed) to drop limited water, and water transforms a cell from burning to burned.
• Reward structure in the RL scenario: A penalty is applied at each timestep proportional to the number of burning cells, with a terminal +10 reward if the fire is completely extinguished. Initial fire and agent positions are randomized during training, with fixed wind conditions.

Methodology in Plain English

The simulator represents the landscape as a grid of cells. Each cell carries four pieces of information: whether it is unburned, burning, or burned; wind speed and direction; vegetation-related modifiers (a canopy ignition modifier and a ground-fuel density modifier, both ranging from -1 to 1); and the slope of the terrain in the direction of each of its neighbors. Slope is static because elevation does not change, while the other fields vary over time.

Fire spreads from a burning cell to its neighbors according to a "propagation potential" score that combines a tunable base propagation factor, the vegetation and fuel density of the source cell, and separate exponential factors for wind and slope (each governed by tunable parameters). Wind is modeled as directional, so spreading is stronger when it aligns with the wind direction.

The distinctive piece is how that potential becomes a probability. Instead of squashing the score with a fixed function, the authors treat the accumulated potential from all burning neighbors as the intensity rate of a Poisson process; multiplying by the target cell's susceptibility gives the effective rate of successful ignitions, and the ignition probability is one minus the exponential of the negative of that rate. This keeps the probability between 0 and 1 and lets gradients flow through the simulator.

The simulator has two modes. In the deterministic mode, it propagates raw continuous probabilities using an update that tracks the probability a cell is unburned, burning, or burned-out. In the stochastic mode, it draws independent Bernoulli samples for ignition and for whether a burning cell continues burning or becomes burned. Because everything is written in JAX and vectorized with vmap, many simulations can run in parallel on a GPU without moving data to main memory between the simulator and the RL algorithm.

For calibration, the authors optimized simulator parameters with gradient descent using the deterministic raw-probabilities routine, minimizing a combination of binary cross-entropy and pooled mean squared error against the real historical fire perimeter. For the RL demonstration, they used the JAX-based RL library Gymnax to connect JaxWildfire to existing algorithm implementations and trained an agent with PPO.

Why This Matters

The paper argues that AI for natural hazards needs to shift from reactive pattern recognition to anticipatory sequential decision-making under uncertainty, and that reinforcement learning is a promising framework for that shift. Its practical bottleneck is interaction speed. By making wildfire simulation roughly an order of magnitude faster and keeping it on the GPU alongside the learning algorithm, JaxWildfire lowers a barrier that had made high-fidelity simulators like FARSITE and ELMFIRE too expensive for RL, and even limited faster cellular automata simulators through memory transfer overhead.

Real-world applications implied by the work:

• Wildfire suppression planning: training policies that decide when and where to drop limited water or retardant from aircraft, as demonstrated with the synthetic air tanker scenario.
• Proactive decision support systems: the authors frame the work as a contribution to broader efforts to build proactive decision support for natural hazard management.
• Simulator parameter calibration: using gradients to fit spread-model parameters to historical fire perimeters, as done with the Bear 2020 wildfire, to keep simulators aligned with observed fire behavior.
• Benchmarking machine learning models of fire spread: the PyTorch export functionality lets researchers test models that learn fire spread dynamics on the same environments.

Industry relevance: the work targets the RL research community directly, and is compatible with Gymnax and with real geospatial data sources (WorldCover, DEM) plus weather inputs supplied as static values or time series. The authors state they aim to publicly release JaxWildfire under a permissive license with documentation and unit tests, at https://github.com/ori-goals/JaxWildfire. They describe JaxWildfire as the first wildfire simulator implemented in JAX specifically designed for RL research, to the best of their knowledge.

Future Directions

• Heterogeneous multi-agent teams: adding support for combinations such as ground vehicles and drones to simulate more complex and practical wildfire management situations.
• Richer fire physics: extending the underlying model to account for the effects of humidity and temperature on fire spread.
• Scaling task complexity: the paper notes that implementation speed is currently a limiting factor for the scale and complexity of RL training tasks, and positions JaxWildfire as a step toward removing that limit.
• Open question of calibration vs. simulation fidelity: the calibration comparison against PyTorchFire on Bear 2020 yielded comparable results with only a slight advantage for JaxWildfire, leaving room to examine how the new Poisson-based ignition model performs across a wider range of real fires. The paper does not report a broader multi-fire calibration study.

Target Audience

Reinforcement learning researchers looking for fast, GPU-resident, differentiable environment simulators; machine learning practitioners working on sequential decision-making under uncertainty for natural hazards; wildfire and emergency-management researchers interested in calibrating or benchmarking spread models against historical perimeters; and JAX ecosystem developers, since the simulator integrates with Gymnax and supports Exported PyTorch data. Readers need comfort with RL algorithms (PPO), cellular automata, and JAX primitives such as vmap.

Authors’ abstract

Artificial intelligence methods are increasingly being explored for managing wildfires and other natural hazards. In particular, reinforcement learning (RL) is a promising path towards improving outcomes in such uncertain decision-making scenarios and moving beyond reactive strategies. However, training RL agents requires many environment interactions, and the speed of existing wildfire simulators is a severely limiting factor. We introduce $\texttt{JaxWildfire}$, a simulator underpinned by a principled probabilistic fire spread model based on cellular automata. It is implemented in JAX and enables vectorized simulations using $\texttt{vmap}$, allowing high throughput of simulations on GPUs. We demonstrate that $\texttt{JaxWildfire}$ achieves 6-35x speedup over existing software and enables gradient-based optimization of simulator parameters. Furthermore, we show that $\texttt{JaxWildfire}$ can be used to train RL agents to learn wildfire suppression policies. Our work is an important step towards enabling the advancement of RL techniques for managing natural hazards.