Plasmax: Differentiable & Parallelizable Environments for Tokamak Kinetic Control

Controlling plasma to obtain the best nuclear fusion performance is an open problem, where progress would have a significant impact on the nuclear fusion research field. While Reinforcement Learning (RL) has great potential to tackle such complex control problems, research on its application to fusion remains sparse. One significant barrier is the lack of fast, configurable RL environments tailored to fusion applications. In this work, we introduce \textit{Plasmax}: 17 differentiable, parallelizable environments for plasma control with 5 different backends for dynamics simulation. Plasmax allows benchmarking reinforcement learning versus traditional control methods, evaluating robustness to different modes of real-world inspired observation corruption, and studying transfer between simulators with different fidelities. We hope Plasmax encourages bridging the machine learning and nuclear fusion research fields, in order to advance nuclear fusion research.