用JAX加速野火模拟，让强化学习更快训练出灭火策略。
JaxWildfire: A GPU-Accelerated Wildfire Simulator for Reinforcement Learning
- 基于概率细胞自动机的野火模型，支持GPU并行仿真。
- 相比现有工具提速6至35倍，可对模拟参数进行梯度优化。
- 已用于训练强化学习灭火策略，适合做灾害管理算法研究者。
人工智能方法正被越来越多地应用于管理野火等自然灾害。特别是强化学习（RL）在不确定性决策场景中展现巨大潜力，有望超越被动应对策略。然而，训练RL代理需要大量环境交互，而现有野火模拟器速度过慢，严重制约了进展。我们提出$ exttt{JaxWildfire}$，一个基于细胞自动机的原理性概率野火传播模型，使用JAX实现，通过$ exttt{vmap}$支持向量化仿真，在GPU上实现高吞吐量。实验表明，$ exttt{JaxWildfire}$相较现有软件提升6至35倍性能，并支持模拟参数的梯度优化。此外，我们展示了该模拟器可用于训练强化学习代理，学习有效的野火抑制策略。本工作为推动强化学习在自然危害管理中的应用迈出关键一步。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。