Gym-TORAX: Open-source software for integrating RL with plasma control simulators

2025-10-13

AI Summary

This work addresses the challenge of jointly optimizing performance and stability in tokamak plasma control. Methodologically, it introduces the first open-source Python framework that deeply integrates the standard reinforcement learning (RL) interface (Gymnasium) with the plasma dynamics simulator TORAX. Through modular design, it enables flexible specification of control actions, state observations, and reward functions, automatically generating reusable RL environments; it further establishes a standardized benchmark environment tailored to ITER discharge initiation scenarios. The key contributions are: (i) the first formalized interface bridging plasma control simulation and the RL ecosystem, significantly enhancing algorithm portability and experimental reproducibility; and (ii) seamless support for rapid integration and validation of mainstream RL algorithms, providing a scalable, reproducible, and open platform for fusion control research.

Technology Category

• Machine Learning: Reinforcement Learning
• Intelligent Robots: Learning & Optimization for ROB
• Search and Optimization: Sampling/Simulation-based Search

Application Category

• User Modeling, Personalization and Recommendation: User modeling and simulation for interactive and conversational systems
• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications
• Responsible Web: Machine-in-the-loop, human agency and autonomy

Abstract

This paper presents Gym-TORAX, a Python package enabling the implementation of Reinforcement Learning (RL) environments for simulating plasma dynamics and control in tokamaks. Users define succinctly a set of control actions and observations, and a control objective from which Gym-TORAX creates a Gymnasium environment that wraps TORAX for simulating the plasma dynamics. The objective is formulated through rewards depending on the simulated state of the plasma and control action to optimize specific characteristics of the plasma, such as performance and stability. The resulting environment instance is then compatible with a wide range of RL algorithms and libraries and will facilitate RL research in plasma control. In its current version, one environment is readily available, based on a ramp-up scenario of the International Thermonuclear Experimental Reactor (ITER).

Problem

• Integrating reinforcement learning with plasma control simulators
• Optimizing plasma performance and stability in tokamaks
• Creating compatible RL environments for plasma dynamics simulation

Innovation

• Open-source Python package for RL plasma control
• Creates Gymnasium environments wrapping TORAX simulator
• Compatible with various RL algorithms and libraries