Skip to content

RL — wrappers and single-agent trainers

Single-agent training adapters. For multi-agent wrappers and trainers, see API → RL MARL. For the conceptual overview, see Training → Wrappers and Training → Trainers.

Wrappers

from powerzoojax.rl import (
    LogWrapper, LogEnvState,
    SafeRLWrapper, SafeRLState,
    RewardWrapper, RewardEnvState,
    bind,
)

flowchart LR
    ENV["pure env"] --> LOG["LogWrapper(env, params)"]
    ENV --> SAFE["SafeRLWrapper(env, params, cost_thresholds)"]
    LOG --> REW["RewardWrapper(wrapped, reward_fn)"]
    SAFE --> REW

LogWrapper

Binds params; tracks per-episode return and length; injects returned_episode_returns, returned_episode_lengths, returned_episode into info.

powerzoojax.rl.wrappers.LogWrapper(env, params=None)

PureJaxRL / Rejax adapter for single-agent PowerZooJax envs.

powerzoojax.rl.wrappers.LogEnvState

Wraps inner env state with episode tracking fields.

SafeRLWrapper

Returns a 6-tuple from step: (obs, state, reward, costs, done, info). Used by CMDP trainers.

powerzoojax.rl.wrappers.SafeRLWrapper(env, params=None, cost_threshold=25.0, *, selected_names=None, cost_thresholds=None)

CMDP wrapper that exposes selected constraint costs as a vector output.

powerzoojax.rl.wrappers.SafeRLState

Single-agent training state with vector CMDP cost tracking.

bind

Convenience constructor: safe=False returns LogWrapper, safe=True returns SafeRLWrapper.

powerzoojax.rl.wrappers.bind(env, params=None, safe=False, cost_threshold=25.0, *, selected_names=None, cost_thresholds=None)

Create a training-compatible wrapper around a core env.

RewardWrapper

Overlay a custom reward on top of an existing wrapper. The unmodified env reward is preserved in info["env_reward"].

powerzoojax.rl.reward.RewardWrapper(env, reward_fn)

Inject a custom reward function into a LogWrapper or SafeRLWrapper.

The custom reward_fn replaces (not adds to) the original env reward. Episode return statistics in info are computed from the custom reward.

Parameters:

Name Type Description Default
env

A LogWrapper or SafeRLWrapper instance.

 required
reward_fn Callable

Callable with signature::

reward_fn(obs, action, next_obs, reward, info) -> scalar reward_fn(obs, action, next_obs, reward, costs, info) -> scalar

obs: observation before the step (from state.prev_obs) action: action taken this step next_obs: observation after the step reward: original environment reward (may be used or ignored) costs: selected constraint vector when available info: environment info dict returns: scalar float32 — the new reward to use

 required
Note

Whether the inner env is a SafeRLWrapper (6-tuple step) is checked once at construction time (isinstance on env), so the static Python branch is JIT-safe under static_argnums=(0,).

Usage::

# Custom SOC-tracking reward
env = LogWrapper(BatteryEnv(), params)
env = RewardWrapper(env, lambda o, a, no, r, i: -jnp.abs(no[0] - 0.5))

# For CMDP — inner env is SafeRLWrapper, selected costs flow through unchanged
safe_env = SafeRLWrapper(
    TransGridEnv(),
    params,
    selected_names=("thermal_overload",),
    cost_thresholds=(0.0,),
)
env = RewardWrapper(safe_env, my_reward_fn)
obs, state, reward, costs, done, info = env.step(key, state, action)

step(key, state, action)

Returns 5-tuple or 6-tuple depending on the wrapped env.

powerzoojax.rl.reward.RewardEnvState

State for RewardWrapper — wraps inner state with custom-reward episode tracking.

Training configuration

from powerzoojax.rl import TrainConfig, load_config, save_config

powerzoojax.rl.config.TrainConfig(algo='ppo', total_timesteps=100000, num_envs=64, seed=42, learning_rate=0.0003, gamma=0.99, max_grad_norm=0.5, eval_freq=10000, eval_episodes=128, n_steps=128, n_epochs=4, n_minibatches=4, clip_eps=0.2, ent_coef=0.01, vf_coef=0.5, gae_lambda=0.95, cost_threshold=0.0, cost_thresholds=(), lambda_lr=0.005, cost_gamma=1.0, cost_scale=1.0, log_lambda_max=5.0, n_checkpoints=1, n_eval_envs=0, saute_horizon=None, saute_unsafe_reward=0.0, saute_use_reward_shaping=True, hidden_dims=(64, 64), continuous_action_dist='gaussian', normalize_observations=False, normalize_rewards=False, sac_num_epochs=None, sac_target_entropy_ratio=None, eval_num_episodes=128, record_eval_wall_time=False, wall_time_warmup=True, wall_time_warmup_timesteps=None) dataclass

Unified training configuration.

Fields map to Rejax or self-implemented CMDP backends:

  • Rejax PPO: num_steps ← n_steps, num_epochs ← n_epochs
  • Rejax SAC/TD3/DQN: only common fields are used
  • CMDP: all fields used, cost_thresholds defines the vector budget

algo = 'ppo' class-attribute instance-attribute

Algorithm: "ppo" | "sac" | "td3" | "dqn" | "ppo_lagrangian" | "saute_ppo" | "ippo" | "ippo_typed" | "ippo_typed_lagrangian" | "mappo".

"ippo_typed" — IPPO with type-specific parameter sharing: agents are partitioned by type prefix (e.g. "battery_*", "renewable_*", "flexload_*"); each type gets an independent SharedActorCritic. Used by the ders-medium and ders-medium-safe presets for heterogeneous DERs.

"ippo_typed_lagrangian" — typed IPPO-Lagrangian for heterogeneous cooperative MARL: the per-type actors remain decentralized, while explicit CMDP constraint costs are enforced through local cost critics and a shared team-level dual variable.

eval_episodes = 128 class-attribute instance-attribute

Episodes per in-training evaluation callback.

This controls lightweight monitoring during training only. It does not override the benchmark's post-training evaluation episode budget.

cost_threshold = 0.0 class-attribute instance-attribute

Scalar fallback used when cost_thresholds is empty and n_constraints == 1.

cost_thresholds = () class-attribute instance-attribute

Vector CMDP budget d_i in the same physical units as the raw env cost. Empty tuple falls back to cost_threshold broadcasting. The trainer divides both rollout costs and these thresholds by cost_scale before any dual update, so values must be given in raw (unscaled) units.

cost_scale = 1.0 class-attribute instance-attribute

Per-constraint cost scaling applied inside the trainer only. Divides raw env costs (and thresholds) before GAE and the dual update to keep per-step values numerically O(1). Does not change the CMDP problem definition.

n_eval_envs = 0 class-attribute instance-attribute

Number of environments to use for greedy eval inside each CMDP update step.

0 = skip eval (default, backward-compatible).

0 = run a greedy rollout of n_steps steps at every update, logs eval_returns (mean episodic return) and eval_cost_<name> (mean per-step cost) — comparable to rejax PPO's eval_returns. Adds ~(n_eval_envs / num_envs) overhead per update.

saute_horizon = None class-attribute instance-attribute

Episode horizon for budget normalisation. None → reads params.max_steps.

saute_unsafe_reward = 0.0 class-attribute instance-attribute

Reward substituted when any constraint budget is exhausted (z_i ≤ 0).

saute_use_reward_shaping = True class-attribute instance-attribute

If False, only augment obs without shaping reward (ablation).

continuous_action_dist = 'gaussian' class-attribute instance-attribute

Continuous-action actor family for Rejax PPO-style trainers.

"gaussian" keeps Rejax's default unbounded Gaussian actor. "beta" uses a bounded Beta policy whose support matches the Box action space directly.

sac_num_epochs = None class-attribute instance-attribute

If set and algo == "sac", passed to Rejax SAC as num_epochs (gradient steps per environment transition batch). If None, Rejax's default applies (commonly 1) — we do not mirror PPO's n_epochs for SAC to avoid changing legacy behaviour.

sac_target_entropy_ratio = None class-attribute instance-attribute

Optional override for Rejax SAC's target_entropy_ratio.

eval_num_episodes = 128 class-attribute instance-attribute

Number of parallel eval rollouts passed to rejax.evaluate (num_seeds).

Lower values speed up each eval checkpoint at the cost of noisier mean return.

record_eval_wall_time = False class-attribute instance-attribute

If True, Rejax backend records host wall time (s) at each eval via io_callback.

Single fused algo.train() is unchanged; overhead is one small host callback per eval (typically tens per run). Use for plots vs wall time instead of splitting train() into many short runs (which re-triggers compile / breaks scan fusion).

Note: Rejax runs one initial eval before any training, then evals every eval_freq; the first interval is therefore not comparable to later ones.

wall_time_warmup = True class-attribute instance-attribute

When record_eval_wall_time is True: run a short training pass first so XLA compiles, then start the wall clock at t0 for the main run (excludes most JIT cost from eval_wall_time_s). Set False for tests or if you want absolute time from process start.

wall_time_warmup_timesteps = None class-attribute instance-attribute

Warm-up horizon (env steps). None → eval_freq (one eval period, minimal scan length). Clamped to [1, total_timesteps). Ignored if wall_time_warmup is False or warm-up would equal the full run.

replace(**kwargs)

Return a new TrainConfig with fields replaced by kwargs.

resolved_cost_thresholds(n_constraints)

Return vector cost thresholds, broadcasting legacy scalar config.

to_rejax_kwargs()

Map to Rejax create() kwargs (algo-dependent subset).

powerzoojax.rl.config.load_config(source, config_cls=TrainConfig)

Load a TrainConfig from a YAML file path or a plain dict.

Parameters:

Name Type Description Default
source

str path to a YAML file, or a dict of field values.

 required
config_cls

Config class to instantiate (default: TrainConfig).

 TrainConfig

Returns:

Type Description
TrainConfig

A TrainConfig instance.

powerzoojax.rl.config.save_config(config, path)

Save a TrainConfig to a YAML file.

Parameters:

Name Type Description Default
config TrainConfig

TrainConfig instance to save.

 required
path str

Destination file path (.yaml).

 required

Trainers

from powerzoojax.rl import TrainResult, make_train, make_cmdp_train, train

make_train(env, config) is the dispatcher. It picks Rejax for algo in {ppo, sac, td3, dqn}, make_cmdp_train for algo == "ppo_lagrangian", or the IPPO backends for MARL envs.

powerzoojax.rl.trainer.make_train(env, config=TrainConfig())

Unified training factory — dispatch to the appropriate backend.

Parameters:

Name Type Description Default
env

Wrapped environment (LogWrapper, SafeRLWrapper, or RewardWrapper).

 required
config TrainConfig

TrainConfig specifying the algorithm and hyperparameters.

 TrainConfig()

Returns:

Type Description
Callable

train_fn(key) -> TrainResult
Backend routing
  • "ppo" | "sac" | "td3" | "dqn" → Rejax (via _RejaxAdapter)
  • "ppo_lagrangian" → self-implemented CMDP
  • "ippo" → shared-param IPPO (all agents one net)
  • "ippo_typed" → type-specific IPPO (DERs default)
  • "ippo_typed_lagrangian" → type-specific IPPO-Lagrangian
  • "saute_ppo" → Rejax PPO on a SauteWrapper env
  • "mappo" → NotImplementedError (pending)

powerzoojax.rl.trainer.TrainResult(params, metrics, config, env_name='', checkpoints=None)

Structured container for training results.

Attributes:

Name Type Description
params

Trained network parameters (pytree).
metrics

Dict of metric arrays with shape (n_updates, ...).
config

The TrainConfig used for this run.
env_name

Optional name of the environment.

summary property

Final training metrics as a plain dict (JSON-serializable).

save(path)

Save summary + config to a JSON file.

powerzoojax.rl.cmdp.make_cmdp_train(env, config)

Create a JIT-compiled multi-constraint PPO-Lagrangian trainer.

One-line entry point

from powerzoojax.rl import train
result = train("dso-nflex", seed=0)

powerzoojax.rl.train.train(preset_or_env, config=None, reward_fn=None, seed=42, **config_overrides)

Train an RL agent on a PowerZooJax environment.

Parameters:

Name Type Description Default
preset_or_env

Either a preset name string (e.g. "battery-soc-tracking") or a wrapped env (LogWrapper / SafeRLWrapper).

 required
config TrainConfig

TrainConfig instance. When using a preset this defaults to the preset's built-in config. When passing an env directly it defaults to TrainConfig() (PPO).

 None
reward_fn

Optional reward function with signature (obs, action, next_obs, reward, info) -> scalar or (obs, action, next_obs, reward, costs, info) -> scalar. Wraps the env in a RewardWrapper when provided.

 None
seed int

Random seed (overrides config.seed).

 42
**config_overrides

Additional TrainConfig field overrides (e.g. total_timesteps=500_000, learning_rate=1e-3).

 {}

Returns:

Type Description
TrainResult

TrainResult with params, metrics, config, env_name.

Examples::

# Preset + seed override
result = train("battery-soc-tracking", seed=0)

# Preset + config override
result = train("case5-economic-dispatch",
               config=TrainConfig(algo="sac", total_timesteps=500_000))

# Direct env
result = train(env, config=TrainConfig(algo="ppo_lagrangian",
                                       cost_thresholds=(0.0,)))

# Custom reward
result = train("battery-soc-tracking",
               reward_fn=lambda o, a, no, r, i: -jnp.abs(no[0] - 0.5))

Presets

from powerzoojax.rl import list_presets, get_preset

powerzoojax.rl.presets.list_presets()

Return all presets as a list of dicts (AI-friendly JSON output).

Each dict has keys: name, description, algo, total_timesteps.

powerzoojax.rl.presets.get_preset(name)

Get a preset by name.

Parameters:

Name Type Description Default
name str

Preset name (e.g. "battery-soc-tracking").

 required

Returns:

Type Description
PresetDef

PresetDef namedtuple.

Raises:

Type Description
KeyError

If the name is not found.