GymTorax Documentation

Python Version PyPI Version License

A Gymnasium environment for reinforcement learning in tokamak plasma control

Gym-TORAX is a Python package that provides reinforcement learning (RL) environments for plasma control, built on top of the TORAX plasma simulator.

Its purpose is to bridge the gap between plasma physics simulation and RL research:

  • For RL users, it exposes ready-to-use environments following the Gymnasium API, abstracting away the plasma physics.

  • For plasma physicists, it provides a framework to design and customize new control tasks, making it easy to test RL algorithms in different operational scenarios.

Key Features

  • Gymnasium Integration: Standard RL environment interface compatible with the Gymnasium ecosystem

  • TORAX Physics: 1D transport equations solved with TORAX plasma physics models

  • Configurable Environment: Flexible action spaces, observation spaces, and reward functions

Getting Started

Installation

pip install gymtorax

Basic Usage

import gymnasium as gym
import gymtorax

# Create environment
env = gym.make('gymtorax/IterHybrid-v0')

# Reset and run with a random agent
observation, info = env.reset()
action = env.action_space.sample()
observation, reward, terminated, truncated, info = env.step(action)