Plasma Shape Control via Zero-shot Generative Reinforcement Learning

TL;DR

This work tackles the limited adaptability of PID control and generalization gaps in RL for tokamak plasma shape control by proposing a zero-shot framework that learns a foundation policy from a large offline PID dataset. The method combines Generative Adversarial Imitation Learning with a Hilbert space latent encoder to produce a goal-conditioned policy that can track diverse plasma configurations without task-specific retraining, trained in a data-driven HL-3 dynamics environment. Key contributions include the Hilbert encoder that encodes target states into a navigable latent space, an adversarial discriminator to enforce PID-like behavior, and a PPO-based actor-critic that optimizes a hybrid imitation/directed-reward objective, demonstrated by precise, stable zero-shot tracking across multiple current and shape scenarios in high-fidelity simulations. The results suggest a data-efficient path toward flexible intelligent plasma control suitable for future fusion devices, reducing the need for retraining when targets or operating conditions change.

Abstract

Traditional PID controllers have limited adaptability for plasma shape control, and task-specific reinforcement learning (RL) methods suffer from limited generalization and the need for repetitive retraining. To overcome these challenges, this paper proposes a novel framework for developing a versatile, zero-shot control policy from a large-scale offline dataset of historical PID-controlled discharges. Our approach synergistically combines Generative Adversarial Imitation Learning (GAIL) with Hilbert space representation learning to achieve dual objectives: mimicking the stable operational style of the PID data and constructing a geometrically structured latent space for efficient, goal-directed control. The resulting foundation policy can be deployed for diverse trajectory tracking tasks in a zero-shot manner without any task-specific fine-tuning. Evaluations on the HL-3 tokamak simulator demonstrate that the policy excels at precisely and stably tracking reference trajectories for key shape parameters across a range of plasma scenarios. This work presents a viable pathway toward developing highly flexible and data-efficient intelligent control systems for future fusion reactors.

Figures (9)

• Figure 1: Overview of the proposed generative reinforcement learning framework for plasma current and shape control on the HL-3 Tokamak System. The framework consists of a data collection module, a training stage, and an inference stage. In the training phase, a Hilbert Encoder maps goals to a structured latent space and generates intrinsic directional rewards. A discriminator incorporates PID data to provide imitation rewards, and an RL Policy is trained to generate control actions based on these intrinsic and imitation rewards. This training is conducted by having the agent interact with a pretrained data-driven dynamics model, which serves as a simulation environment. During the inference stage, the trained Hilbert encoder and RL agent are deployed to achieve zero-shot control over diverse reference plasma configurations.
• Figure 2: Correspondence between physical plasma states and the Hilbert latent space.
• Figure 3: Network architectures for the Encoder, Actor, Critic, and Discriminator modules.
• Figure 4: Overview of the model training process for the generative RL framework.
• Figure 5: Key performance metrics during the training process. This figure illustrates the evolution of critical metrics over 2000 training epochs. (a) Encoder loss, (b) Discriminator loss, and (c) Discriminator scores for both the agent policy and the PID data.