Plasmax: differentiable & parallelizable environments for kinetic control in Tokamaks

plasmax provides JAX-native fusion-control environments built on TORAX. The environments are inspired by real tokamak devices; ITER, SPARC, STEP, and KSTAR. Each tokamak contains different tasks for different scenarios and phases.

Environment design is inspired by recommendations from Challenges of Real World Reinforcement Learning, where missing, noisy observations are explicitly implemented.

Install

pip install plasmax

We did not package the training and agent stack to make the dependencies lighter. If you want to try training agents using our stack:

git clone https://github.com/TheodoreWolf/plasmax
cd plasmax
pip install -e . --group research

Quick start

import jax
import jax.numpy as jnp

import plasmax
from plasmax.wrappers import OracleWrappers, RealisticWrappers

env = plasmax.make(
    "iter/hybrid/flattop",
    backend="bohm_gyrobohm",
)
env = RealisticWrappers(env, max_steps=100)

state, info = env.init(jax.random.key(0))
action = jnp.zeros(env.action_space.shape, dtype=env.action_space.dtype)
state, info = env.step(state, action)

print(info.obs, info.reward)
print(info.terminated, info.truncated)

Environments are loaded through the make function. TORAX environment names are structured as {tokamak}/{scenario}/{phase} and require an explicit compatible backend. KSTAR is a standalone environment and takes no backend.

make returns the bare environment with physical actuator units and full observations. Apply RealisticWrappers(env) to add configured physics randomization, sensor effects, action scaling, history, and truncation. OracleWrappers(env) applies action scaling, history, and truncation while keeping nominal physics and full observations. Both accept max_steps and time_aware; RealisticWrappers also accepts quantize_bins. KSTAR supports RealisticWrappers, preserving its native actions and observations.

The environments use the Envelope API and contracts, this includes e.g. explicit truncation versus termination.

Tasks and backends

Task aliases	Compatible backend aliases
iter/{baseline,hybrid,advanced}/{rampup,flattop,rampdown}	cgm, qlknn, bohm_gyrobohm, tglfnn, tglfnn_nr
sparc/{prd,reduced_field}/{rampup,flattop,rampdown}	cgm, qlknn, bohm_gyrobohm, tglfnn, tglfnn_nr
step/spp_001_ec_hd/flattop	bohm_gyrobohm_step, tglfnn_spherical
mock/circular/smoke	mock
kstar_worldmodel	None; standalone environment

Unsupported environment/backend pairs are rejected before construction.

smoke = plasmax.make("mock/circular/smoke", backend="mock")
kstar = plasmax.make("kstar_worldmodel")

NB: the tglfnn_spherical backend requires a repository clone for now. TGLFNN-UKAEA is still an eager transitive TORAX dependency. TORAX 1.4.3 pins fusion-surrogates 0.4.6, whose TGLFNN extra pins the older 0.1.0 weights. Repository clones use a uv-only override to the final PyPI 0.2.0 weights until TORAX adopts fusion-surrogates 0.4.7. Published installs still follow TORAX's dependency metadata.

Equilibria generated with FreeGSNKE are committed artifacts, so FreeGSNKE is not a runtime dependency.

Every leaf task YAML owns its default reward:

task:
  reward: lh_transition

Omitting reward uses the task default; built-in TORAX rewards apply squareplus normally, its logarithm on disruption or solver failure, and zero on invalid state.

env = RealisticWrappers(plasmax.make("iter/advanced/rampup", backend="qlknn"))

oracle_ablation = OracleWrappers(
    plasmax.make(
        "iter/advanced/rampup",
        backend="qlknn",
        reward="Q_fusion",
    )
)

Individual wrappers also resolve their defaults from env.plasmax_config and accept their existing explicit arguments for custom compositions:

from plasmax.wrappers import NoiseWrapper, PhysicsRandomizationWrapper

env = plasmax.make("iter/baseline/flattop", "bohm_gyrobohm")
env = NoiseWrapper(PhysicsRandomizationWrapper(env))

Physics parameters live in EnvState.phys_params. Use env.with_physics(state, parameters) to replace selected entries, including through nested wrapper states. Values persist until overwritten or reset to nominal values. PhysicsRandomizationWrapper samples before every step, always using the configured nominals for relative ranges. Each stochastic wrapper owns its RNG; init(key) and reset(state, key) seed these streams.

Environment boundary

init, step, and reset return (state, info). A transition exposes:

• info.obs: the post-transition flat observation;
• info.reward: a scalar float32 RL-boundary reward;
• info.terminated: a physical, solver, or invalid-state termination;
• info.truncated: the configured time-limit cutoff;
• info.termination_code: the environment's termination reason.

If termination and the time limit coincide, termination wins.

Fixed-shape rollout collection is part of the installed library:

from plasmax import collect_episode

env = RealisticWrappers(plasmax.make("iter/hybrid/flattop", "bohm_gyrobohm"))


def act(obs, key):
    del obs, key
    return jnp.zeros(env.action_space.shape, dtype=env.action_space.dtype)


trajectory = collect_episode(
    act,
    env,
    jax.random.key(1),
    num_steps=env.max_steps,
)

The collector retains the first terminal transition, stops stepping the environment, and pads the remaining fixed-size output with valid=False.

Contribution and Development

We welcome contributions! To contribute, first fork the repository, then:

git clone {your_gh_username}/plasmax
cd plasmax
# We highly encourage uv for developement
uv sync --group dev
git checkout {name}/{what_you_are_changing}

Then you can open a PR in this repository. Make sure to run tests, CI will do this for you as well. Please have respect for the developer's time and do not submit PRs that can not be reasonably reviewed (even with the help of agents).

See Testing plasmax for the test commands and the local environment trajectory benchmark used by the pull-request report.

Repository layout

src/plasmax/        installed environments, tooling, configs, and data
agents/             clone-only baseline agents
training/           clone-only training launchers, adapters, and shared utilities
scripts/            generic evaluation and rollout launchers
tools/              artifact and equilibrium generation
benchmarks/         backend agreement and throughput benchmarks
experiments/        research studies and plotting
tests/              library and release tests

See training commands and evaluation and rollout commands for the clone-only entrypoints.

License and attribution

The library is licensed under the Apache License 2.0. TORAX and packaged third-party data/model assets retain their own attribution and license terms; the relevant notices are shipped adjacent to those assets.

Agents

An AGENTS.md file is included, which has my own personal code preferences. We recommend users who want to use agents to obtain an explicit JAX skill (I've written my own, that I will open-source, when I'm happy with it), as current agents are still not great at this. Agents were utilized throughout this work, while I did my best to check the code, mistakes remain. In my experience, the most dangerous are comments that state mistakes or bad assumption as facts, this then further reinforce the agents in their bad ideas.

Citation

Coming soon, once I get the paper out...