Skip to content

JaxWildfire: A GPU-Accelerated Wildfire Simulator for Reinforcement Learning

🕒 Published (v1): 2025-12-05 19:06 UTC · Source: Arxiv · Venue: NEURIPS 2025 · link

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

JaxWildfire is a JAX-native wildfire simulator built on a probabilistic cellular automata model that achieves 6–35× speedup over PyTorchFire by running fully on GPU via vmap. It supports differentiable parameter calibration from real wildfire data and serves as an RL training environment, demonstrated with a PPO air-tanker fire suppression agent.

Problem

Existing wildfire simulators are too slow for RL training: high-fidelity physics-based models (FARSITE, ELMFIRE) are computationally prohibitive, and even CA-based simulators incur GPU↔CPU data transfer overhead during RL agent interaction. No JAX-native wildfire simulator existed, preventing the efficient in-accelerator training loops that modern RL research increasingly requires.

Method

JaxWildfire implements a cellular automata fire spread model in JAX. Each cell carries fire state (unburned/burning/burned), wind (speed + direction), vegetation/fuel density modifiers, and slope. Fire propagation from a source cell to neighbors is scored via directional factors for wind (κ_wind) and slope (κ_slope) multiplied by vegetation/fuel terms. Ignition probability is derived by interpreting cumulative propagation potential as a Poisson arrival rate and applying a Poisson-CDF formula: p_ignite(v) = 1 − exp(−γ(v)λ), which is differentiable and bounded in [0,1]. The simulator has two modes: stochastic (Bernoulli sampling) and deterministic (raw probability propagation for gradient-based calibration via optax/Adam). JAX's vmap enables batched parallel simulations. Real-world environments are constructed automatically from ESA WorldCover vegetation maps and Copernicus DEM elevation data given a place name or bounding box. RL integration is via Gymnax compatibility.

Key Contributions

  • First JAX-native wildfire simulator designed for RL research, eliminating GPU↔CPU transfer bottleneck.
  • Principled Poisson-process ignition model replacing ad-hoc tanh squashing, enabling differentiable simulator calibration.
  • 6–35× wall-time speedup over PyTorchFire for single simulations; 18–51× using for_i (history-free) mode; peak throughput of 2,516 steps/s with 64 vectorized environments on a single A100.
  • Integration with ESA WorldCover and Copernicus DEM for geographically realistic environments via place-name lookup.
  • Demonstrated successful RL policy learning: PPO-trained air-tanker agent achieves fire suppression in a 20×20 synthetic environment.
  • Calibration on Bear 2020 real wildfire achieves IoU=0.424 vs. PyTorchFire's IoU=0.348.

Results

  • Speed: 6–35× over PyTorchFire (single simulation, GPU, grid sizes 100–900, NVIDIA A100); 18–51× with for_i mode.
  • Throughput: 2,516 steps/s with 64 vectorized environments on A100 at 64×64 grid.
  • Calibration accuracy: JaxWildfire IoU=0.424 vs. PyTorchFire IoU=0.348 on Bear 2020 wildfire perimeter prediction.
  • RL training: PPO agent successfully learns fire suppression (episode return converges toward +10 terminal reward) in synthetic 20×20 environment within ~10,000 training steps.

Limitations

  • RL validation is only on a small synthetic 20×20 environment; no real-world or large-scale RL benchmark.
  • No support for multi-agent scenarios (ground vehicles, drones) yet.
  • Fire spread model omits humidity and temperature effects.
  • Weather is provided as static or pre-specified time series; no dynamic weather coupling.
  • Calibration comparison is limited to one historical wildfire (Bear 2020).

Relevance to Agentic AI / LLM Agents

JaxWildfire is not directly an LLM agent paper, but it is highly relevant as infrastructure for sequential decision-making agents in safety-critical environments. High-throughput simulators are the rate-limiting factor for applying RL (and by extension, LLM-guided or hybrid agents) to real-world physical domains; a 6–51× simulation speedup directly expands the tractable complexity of wildfire management tasks that agents can be trained on. The differentiable simulator also opens the door to model-based RL and gradient-informed planning, which are increasingly explored in agentic frameworks operating in physical or simulation environments.