# JaxWildfire: A GPU-Accelerated Wildfire Simulator for Reinforcement Learning

arXiv: https://arxiv.org/abs/2512.06102 · HTML: https://mlsanity.com/paper/2512.06102/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Placed in RL for energy, networks and operations, one of 15 directions in Deep reinforcement learning methods.
- Builds on 6 of its 21 references that are in the map.
- First posted on arXiv 5 Dec 2025.

## Placed in

- Deep reinforcement learning methods / [RL for energy, networks and operations](https://mlsanity.com/deeprl/rl-for-energy-networks-and-operations/)

## Builds on (most cited first, 6 of 6)

- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Discovered Policy Optimisation](https://mlsanity.com/paper/2210.05639/) (2022)
- [JaxMARL: Multi-Agent RL Environments and Algorithms in JAX](https://mlsanity.com/paper/2311.10090/) (2023)
- [PyTorchFire: A GPU-Accelerated Wildfire Simulator with Differentiable Cellular Automata](https://mlsanity.com/paper/2502.18738/) (2025)
- [Reinforcement Learning for Wildfire Mitigation in Simulated Disaster Environments](https://mlsanity.com/paper/2311.15925/) (2023)

## Built on by (most cited first, 2 of 2)

- [Neural-Parameterized Cellular Automata for Wildfire Spread](https://mlsanity.com/paper/2606.11676/) (2026)
- [Aerial Wildfire Suppression Planning with a Hybrid CNN-Cellular Automata Fire Model](https://mlsanity.com/paper/2606.13633/) (2026)

## Abstract (arXiv)

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 textttJaxWildfire, a simulator underpinned by a principled probabilistic fire spread model based on cellular automata. It is implemented in JAX and enables vectorized simulations using textttvmap, allowing high throughput of simulations on GPUs. We demonstrate that textttJaxWildfire achieves 6-35x speedup over existing software and enables gradient-based optimization of simulator parameters. Furthermore, we show that textttJaxWildfire 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.