JaxWildfire: A GPU-Accelerated Wildfire Simulator for Reinforcement Learning

2025-12-05

AI Summary

Low computational efficiency of existing wildfire simulators severely hinders reinforcement learning (RL) training for uncertain decision-making tasks such as fire management. To address this, we propose the first end-to-end differentiable wildfire simulator built on JAX: it models fire spread using a probabilistic cellular automaton and leverages GPU parallelism and `vmap`-based vectorization to enable efficient gradient computation and batched simulation. Compared to established tools (e.g., FARSITE, Prometheus), our simulator achieves 6–35× speedup while natively supporting RL training frameworks. Experiments demonstrate its effectiveness in training policy-gradient-based agents that generalize across diverse terrains and wind conditions, significantly improving fire suppression performance. This work pioneers the use of JAX in wildfire modeling, establishing a new paradigm for high-fidelity, differentiable, and scalable simulation in natural disaster response.

Technology Category

• Search and Optimization: Sampling/Simulation-based Search
• Machine Learning: Imitation Learning & Inverse Reinforcement Learning
• Cognitive Modeling & Cognitive Systems: Simulating Human Behavior

Application Category

• User Modeling, Personalization and Recommendation: User modeling and simulation for interactive and conversational systems
• Search and Retrieval-Augmented AI: Agentic search
• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications

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 $ exttt{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 $ exttt{vmap}$, allowing high throughput of simulations on GPUs. We demonstrate that $ exttt{JaxWildfire}$ achieves 6-35x speedup over existing software and enables gradient-based optimization of simulator parameters. Furthermore, we show that $ exttt{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.

Problem

• Speeds up wildfire simulation for RL training
• Enables gradient-based optimization of simulator parameters
• Trains RL agents for wildfire suppression policies

Innovation

• GPU-accelerated wildfire simulator using JAX
• Vectorized simulations via vmap for high throughput
• Gradient-based optimization of simulator parameters