JaxWildfire: A GPU-Accelerated Wildfire Simulator for Reinforcement Learning

TL;DR

JaxWildfire is a GPU-accelerated wildfire simulator built with JAX, enabling fast, vectorized simulations and gradient-based optimization, which facilitates reinforcement learning for wildfire management.

Contribution

It introduces a novel, high-speed wildfire simulator with GPU acceleration and differentiable capabilities, advancing RL applications in natural hazard management.

Findings

Achieves 6-35x speedup over existing wildfire simulators.

Enables gradient-based optimization of wildfire models.

Facilitates training RL agents for wildfire suppression policies.

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.