TargetGym

21 JAX environments for setpoint tracking,
with tuned PID and MPC baselines.

Reach the target. Then hold it.

_{One example task from each family, under PID control.}

TargetGym provides reinforcement learning environments for target MDPs: tasks where the objective is to reach a setpoint and hold it against disturbances, rather than to reach a goal state once. The environments model real plants, including an A320-like aircraft, a glass furnace, a nuclear reactor, a cement kiln, a grid battery and a wind turbine.

• Baselines included. Every environment ships a tuned PID; 20 of 21 also ship an MPC with full state access, which serves as an upper bound. Both are recorded over ten seeds in src/target_gym/data/baseline_returns.json.
• Validated physics. Each environment carries a PHYSICS.md with a sourced parameter table, published validation targets asserted by tests, and its documented approximations.
• JAX throughout. jit, vmap and scan compatible, end-to-end GPU, 0.6 M to 700 M steps/s on CPU depending on the plant.
• gymnax API, with a Gymnasium wrapper for non-JAX libraries and a JaxMARL interface for the multi-agent patrol task.

Learned-policy results are not published yet. The measurement protocol is defined in docs/rl-protocol.md.

Installation

pip install target-gym

Quickstart

Also available as a Colab notebook.

import jax
import numpy as np
from target_gym import Plane

env = Plane()
pid = env.make_pid()            # the shipped baseline, tuned

key = jax.random.PRNGKey(0)
obs, state = env.reset(key)

total = 0.0
for _ in range(env.default_params.max_steps_in_episode):
    action = pid(np.asarray(obs))
    obs, state, reward, terminated, truncated, _ = env.step(key, state, action)
    total += float(reward)
    if terminated or truncated:
        break

print("PID return:", total)

reset and step follow the gymnax API and take an optional params; each environment exports its parameter class (PlaneParams, and so on) for custom configurations. Every environment also exposes make_pid(), make_mpc() and save_video().

# doc: skip (trains for 10 000 steps; tests/plane/test_agent.py covers this path)
from target_gym import GymnasiumPlane
from stable_baselines3 import SAC

env = GymnasiumPlane()
model = SAC("MlpPolicy", env, verbose=1)
model.learn(total_timesteps=10_000, log_interval=4)

obs, info = env.reset()
while True:
    action, _states = model.predict(obs, deterministic=True)
    obs, reward, terminated, truncated, info = env.step(action)
    if terminated or truncated:
        break

Comparing against the baselines

For most environments the defaults are the scored configuration, so a plain rollout is comparable to the published numbers. The exceptions are the plants that host several task variants: PlaneParams is shared by plane, plane_sine and plane_energy, which are scored over 280, 480 and 1200 steps, so one class cannot default to all three. EnvSpec.test_params carries the per-variant settings, and spec.make_test_params() resolves them:

import numpy as np

from target_gym.provenance import load_recorded_baselines
from target_gym.registry import REGISTRY
from target_gym.runners.runners import baseline_policy, rollout

spec = REGISTRY["plane"]
params = spec.make_test_params()                     # the scored configuration
pid = baseline_policy(spec, "pid", params)

print("PID:", float(np.sum(rollout(spec, params, pid, seed=0)[2])))
print("MPC:", load_recorded_baselines()["plane"]["mpc_returns"][0])

The PID reads obs only, as a plant controller does. The MPC reads the full state, including quantities the observation withholds, which is what makes it an upper bound rather than a peer; docs/baselines.md quantifies the resulting advantage.

Vectorised rollouts, the registry API, the patrol interface and the wind model are covered in docs/getting-started.md.

Environments

Full environment reference → with observation and action shapes, tracked variables, baseline returns and physics contracts.

The CSTR, first-order and four-tank models are adapted from PC-gym and checked against its source term by term, as each PHYSICS.md provenance line records.

The two patrol environments are multi-agent: wingmen hold a slot on a lead flying its own route, so the reference is another aircraft and collision ends the episode.

Environments span six difficulty tiers, graded on dynamics (linearity, coupling, stiffness) and on the RL side (dimensionality, horizon, partial observability), from a first-order lag at tier 1 to the cement kiln and the multi-agent patrol at tier 6. See the complexity ladder.

Each environment renders a control-room dashboard: plant schematic, gauges with limit and setpoint markers, strip charts, and explicit marking of quantities the controller cannot measure. See the rendering guide.

Why setpoint tracking

Holding a setpoint indefinitely exposes failure modes that episodic goal-reaching does not:

Also modelled: actuator lag, competing objectives, and scheduled setpoints that reward anticipation (the building's night setback, the furnace's crown schedule, the battery's dispatch blocks). Aircraft fly in steady wind, altitude shear and Ornstein-Uhlenbeck turbulence, unobservable by default.

Baselines

Every environment ships a tuned PID; 20 of 21 also ship an MPC. The exception is patrol_bearing_only, which withholds the slot error a planner would read.

Controller structure is chosen per plant:

• Boiler drum: three-element control, so feedwater tracks measured steam flow and shrink-and-swell cannot mislead the level loop.
• Cement kiln: cascade, because integral action on a 25-minute-old measurement oscillates at the delay period.
• Four-tank: crossed loops, since the negative RGA element makes the diagonal pairing unstable.

Three MPC implementations are used: CasADi/IPOPT where a symbolic model exists, gradient-based planning through the JAX dynamics elsewhere, and cross-entropy sampling for the cement kiln. Solver convergence is recorded alongside every result.

Each baseline must beat the best constant action on its environment, a deliberately low bar that a mis-wired controller fails. Weak baselines are documented as such on their environment page.

A PID losing on a task says something about that task, not about PID control. These environments are selected for problems where anticipation pays. Where reacting to the reference and the disturbances is sufficient, a PID is optimal or close enough that the difference does not appear in a return. Two environments were corrected after measurement showed exactly this: a battery whose dispatch signal was so noisy that no controller could exceed 0.43 of the reward ceiling, and a patrol follower tracking a lead at one fixed turn rate, which a single feedforward term cancels.

docs/baselines.md covers the MPC implementations, tuning and caching, solver reporting and per-environment coverage.

Physics validation

Each environment carries a PHYSICS.md stating what it models and where that holds, a parameter table citing or deriving every constant (flagged TUNED - not sourced when neither applies), the published numbers it must reproduce, and its known deviations.

Tests assert consequences rather than formulas: ISA table values, L/D ratios, thermal time constants, energy balances, equilibria. A test that recomputes the implementation's own expression would pass on a wrong one.

All 21 environments are covered by fifteen contracts, since aircraft variants share a plant. A shared conformance suite additionally checks determinism, PRNG handling, jit/vmap/scan compatibility and numerical health over full episodes.

docs/PHYSICS_METHODOLOGY.md documents the method and lists what each model is validated against.

Documentation

Rendered with search and navigation at yannberthelot.github.io/TargetGym. The links above point at the Markdown in this repository, which reads on GitHub and is rewritten for the site at build time.

Related projects

• gymnax: the JAX environment API TargetGym implements.
• Gymnasium: supported through a wrapper for non-JAX libraries.
• PC-gym: process-control environments; source of the CSTR, first-order and four-tank models.
• safe-control-gym: closest in intent, comparing classical control, MPC and RL on shared tasks. It covers more controllers, TargetGym covers more plants.
• JaxMARL: the multi-agent API the patrol formation task follows.

Roadmap

Rewards, physics, baselines and the measurement protocol are settled. Outstanding before 1.0: published learned-policy results and hosted documentation. Known gaps are recorded rather than omitted.

Full roadmap →

Contributing

Bug reports, new environments, improved baselines and physics corrections are welcome.

git clone https://github.com/YannBerthelot/TargetGym.git
cd TargetGym
uv sync --group dev
make ci                # ruff, black --check, docs, fast tests

Further targets: make test, make test-all, make figures, make videos, make tuning.

Adding an environment means registering an EnvSpec, which inherits every shared check, and writing a PHYSICS.md sourcing its numbers. CONTRIBUTING.md has the details.

Citation

@misc{targetgym2025,
  title        = {TargetGym: Reinforcement Learning Environments for Target MDPs},
  author       = {Yann Berthelot},
  year         = {2025},
  url          = {https://github.com/YannBerthelot/TargetGym},
  note         = {Lightweight physics-based RL environments for aircraft, process control, and industrial systems}
}

License

MIT.