Getting started¶

Install¶

pip install target-gym
# or
uv add target-gym

Python 3.11 through 3.14. For development, see CONTRIBUTING.

An episode, directly¶

Environments follow the gymnax interface: pure functions taking a PRNG key, a state and parameters, and returning the next state. Nothing is hidden in instance attributes, so the whole thing composes with jit, vmap and scan.

import jax
import jax.numpy as jnp
from target_gym import CSTR, CSTRParams

env = CSTR()
params = CSTRParams(max_steps_in_episode=100)

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

step = jax.jit(env.step_env)
total = 0.0
for _ in range(params.max_steps_in_episode):
    action = jnp.array([0.0])
    obs, state, reward, terminated, info = step(key, state, action, params)
    total += float(reward)
    if bool(terminated):
        break
print(total)

step_env reports natural termination only -- the plant reaching a state it cannot come back from. The time limit is separate, and gymnax's six-value env.step applies it, returning terminated and truncated separately. Use step_env when you want to reason about the physics, and step when you want the standard RL episode boundary.

Many episodes at once¶

The reason for the functional interface: rollouts vectorise.

import jax
import jax.numpy as jnp
from target_gym import CSTR, CSTRParams

env, params = CSTR(), CSTRParams(max_steps_in_episode=100)

def episode_return(key):
    _, state = env.reset_env(key, params)

    def body(carry, _):
        state, total = carry
        _, state, reward, _, _ = env.step_env(key, state, jnp.zeros(1), params)
        return (state, total + reward), None

    (_, total), _ = jax.lax.scan(body, (state, 0.0), None, params.max_steps_in_episode)
    return total

keys = jax.random.split(jax.random.PRNGKey(0), 256)
returns = jax.jit(jax.vmap(episode_return))(keys)
print(returns.mean())

The Gymnasium interface¶

For libraries that expect the classic API, every environment has a wrapper:

from target_gym import GymnasiumPlane

env = GymnasiumPlane()
obs, info = env.reset(seed=0)
obs, reward, terminated, truncated, info = env.step((0.8, 0.0))

gym_wrapper_factory builds the same wrapper for any of the JAX environments. Note that the wrapper is stateful and returns NumPy -- unlike step_env, it must not be wrapped in jax.jit.

Working from the registry¶

The registry is how the library refers to environments generically -- it is what the test suite, the baselines and the figure generation are all driven by, and the most convenient way to write code that works across environments.

from target_gym.registry import REGISTRY

spec = REGISTRY["boiler_drum"]
env, params = spec.make_env(), spec.params_cls()
pid = spec.make_pid()          # the shipped baseline, or None
pid.reset()

See the environment reference for every registry name.

Rendering¶

Each environment renders a control-room dashboard: a live plant schematic, an instrument stack with limit and setpoint markers, and strip charts. Quantities the controller cannot measure are marked, so a frame shows both what the agent knows and what it is up against.

Full guide, including the two rendering toolkits, headless setup and how to regenerate the shipped media: Rendering.

from target_gym import Plane, PlaneParams

env = Plane()
params = PlaneParams(max_steps_in_episode=1_000)
env.save_video(
    lambda obs: (0.8, 0.0), seed=42,
    folder="videos", episode_index=0, params=params, format="gif",
)

The multi-agent patrol¶

(both aircraft learn a cooperative formation):

import jax
from target_gym import PlanePatrolMARL

env = PlanePatrolMARL(num_wingmen=4)         # 5 planes: 1 lead + 4 wingmen
params = env.default_params
key = jax.random.PRNGKey(0)

obs, state = env.reset(key, params)          # obs = {"lead": ..., "wingman_0": ..., ...}
actions = {agent: env.action_space(agent).sample(key) for agent in env.agents}
obs, state, rewards, dones, info = env.step(key, state, actions, params)
# rewards/dones are dicts keyed by agent (+ dones["__all__"]); reward is shared.

Wind and turbulence¶

Every plane-based env (Plane, Plane3D, PlanePatrol, PlanePatrolMARL) inherits a wind model applied to the air-relative aerodynamics. It is an unobservable disturbance by default — pass observe_wind=True for a fully-observable baseline. Formations feel a single shared gust field.

from target_gym import Plane, PlaneParams

params = PlaneParams(
    wind_x=-15.0,          # steady mean wind (m/s), world frame
    wind_shear_x=0.02,     # + linear altitude shear: +0.02 m/s per metre above...
    shear_ref_alt=5000.0,  # ...this reference altitude
    turbulence_sigma=3.0,  # + Ornstein-Uhlenbeck gusts (0 = off); theta = turbulence_theta
)

hidden = Plane()                     # wind is a hidden disturbance (POMDP)
baseline = Plane(observe_wind=True)  # appends the realized wind to the observation