JaxWildfire: A GPU-Accelerated Wildfire Simulator for Reinforcement Learning

Fuente: arXiv
Guardado en:
Detalles Bibliográficos
Autores principales: Çakır, Ufuk, Darvariu, Victor-Alexandru, Lacerda, Bruno, Hawes, Nick
Formato: Preprint
Publicado: 2025
Materias:
Acceso en línea:
Etiquetas: Agregar Etiqueta
Sin Etiquetas, Sea el primero en etiquetar este registro!
_version_ 1866917129087877120
author Çakır, Ufuk
Darvariu, Victor-Alexandru
Lacerda, Bruno
Hawes, Nick
author_facet Çakır, Ufuk
Darvariu, Victor-Alexandru
Lacerda, Bruno
Hawes, Nick
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2512_06102
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle JaxWildfire: A GPU-Accelerated Wildfire Simulator for Reinforcement Learning
Çakır, Ufuk
Darvariu, Victor-Alexandru
Lacerda, Bruno
Hawes, Nick
Machine Learning
Artificial Intelligence
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.
title JaxWildfire: A GPU-Accelerated Wildfire Simulator for Reinforcement Learning
topic Machine Learning
Artificial Intelligence
url https://arxiv.org/abs/2512.06102