Wrappers¶

All wrappers are importable from powerzoo.wrappers.

Task-aware PettingZoo wrapping for benchmark tasks now lives under powerzoo.tasks.interfaces.TaskPettingZooWrapper. It remains re-exported from powerzoo.wrappers for backward compatibility.

from powerzoo.wrappers import (
    GymnasiumWrapper,
    NormalizationWrapper,
    SafeRLWrapper,
    GymnasiumSafeWrapper,
    ForecastWrapper,
    MARLWrapper,
    FlattenWrapper,
)

`powerzoo.wrappers.gym_wrappers.GymnasiumWrapper(env)` ¶

Bases: Wrapper

Wrap a PowerZoo GridEnv to produce a standard Gymnasium interface.

The inner env's step() returns a state dict as observation; this wrapper calls env.obs(state) to obtain a flat numpy array and passes it through as the Gymnasium observation.

It also exposes: - env.observation_space / env.action_space (from inner env) - env.obs_names / env.action_names (human-readable labels) - Correct (obs, info) return from reset()

Parameters:

Name	Type	Description	Default
`env`		Any `GridEnv` subclass (`TransGridEnv`, `DistGridEnv`, …).	required

Example::

from powerzoo.envs.grid.trans import TransGridEnv
from powerzoo.wrappers import GymnasiumWrapper

raw = TransGridEnv()
env = GymnasiumWrapper(raw)

obs, info = env.reset(seed=0)
obs, reward, terminated, truncated, info = env.step(env.action_space.sample())

`reset(*, seed=None, options=None, **kwargs)` ¶

Reset inner env and return (obs_array, info).

`step(action)` ¶

Step inner env; convert state dict → flat obs array.

Numpy array actions are auto-converted to the dict format expected by GridEnv (e.g. a unit-dispatch vector → {'unit_power_mw': action}).

Adapts any PowerZoo GridEnv to the standard Gymnasium 5-tuple API:

from powerzoo.envs.grid.trans import TransGridEnv
from powerzoo.wrappers import GymnasiumWrapper

env = GymnasiumWrapper(TransGridEnv())
obs, info = env.reset(seed=42)
obs, reward, terminated, truncated, info = env.step(env.action_space.sample())

`powerzoo.wrappers.gym_wrappers.NormalizationWrapper(env, clip=True)` ¶

Bases: ObservationWrapper

Normalise observations to [-1, 1] using fixed bounds.

Bounds are derived automatically from the inner environment's case data (physical limits), and are stable across episodes — safe to use during training.

Parameters:

Name	Type	Description	Default
`env`	`Env`	A `GymnasiumWrapper` or any `gym.Env` with a `Box` obs space.	required
`clip`	`bool`	Whether to clip the normalised observation to `[-1, 1]`. Default `True`.	`True`

The raw bounds can be inspected via env.obs_low and env.obs_high.

Example::

from powerzoo.wrappers import GymnasiumWrapper, NormalizationWrapper

env = NormalizationWrapper(GymnasiumWrapper(TransGridEnv()))
obs, info = env.reset(seed=0)   # obs ∈ [-1, 1]

Normalises observations (and optionally actions) to [−1, 1] using running statistics. Stacks on top of GymnasiumWrapper:

from powerzoo.wrappers import GymnasiumWrapper, NormalizationWrapper

env = NormalizationWrapper(GymnasiumWrapper(TransGridEnv()))

`powerzoo.wrappers.safe_rl_wrapper.SafeRLWrapper(env, cost_threshold=None)` ¶

Bases: Wrapper

Wrap a Gymnasium env to emit (obs, reward, cost, terminated, truncated, info).

`step(action)` ¶

Returns a 6-tuple (obs, reward, cost, terminated, truncated, info) compatible with OmniSafe and Safety-Gymnasium.

Cost extraction priority:

info['selected_constraint_costs'] — task-selected CMDP vector
info['constraint_costs'] — core env full vector
info['cost_sum'] or compatibility info['cost']
0.0 — safe fallback

from powerzoo.wrappers import GymnasiumWrapper, SafeRLWrapper

env = SafeRLWrapper(GymnasiumWrapper(TransGridEnv()), cost_threshold=25.0)
obs, info = env.reset(seed=0)
obs, reward, cost, terminated, truncated, info = env.step(env.action_space.sample())

Parameter	Default	Description
`cost_threshold`	`25.0`	Scalar threshold exposed as `env.cost_threshold` for OmniSafe

`powerzoo.wrappers.safe_rl_wrapper.GymnasiumSafeWrapper(env, cost_threshold=None)` ¶

Bases: Wrapper

Keep Gymnasium's 5-tuple API while projecting vector costs to info['cost'].

`step(action)` ¶

Returns the standard Gymnasium 5-tuple but injects cost into info['cost']. Use this when your algorithm reads cost from info rather than as a separate return value.

from powerzoo.wrappers import GymnasiumWrapper, GymnasiumSafeWrapper

env = GymnasiumSafeWrapper(GymnasiumWrapper(TransGridEnv()))
obs, reward, terminated, truncated, info = env.step(env.action_space.sample())
print(info['cost'])  # non-negative scalar

`powerzoo.wrappers.forecast_wrapper.ForecastWrapper(env, horizon=6, mode='perfect', noise_std=0.02, normalize=True)` ¶

Bases: ObservationWrapper

Augment observations with a look-ahead demand forecast.

Parameters¶

env : gym.Env A Gymnasium-wrapped PowerZoo GridEnv. horizon : int Number of future time steps to append to the observation. Default 6. mode : str 'perfect', 'noisy', or 'none'. Default 'perfect'. noise_std : float Fractional Gaussian noise std for mode='noisy' (e.g. 0.02 = 2 %). Ignored for other modes. Default 0.02. normalize : bool Whether to normalise forecast values by the maximum demand in the dataset (so each value is in [0, ~1]). Default True.

`reset(*, seed=None, options=None, **kwargs)` ¶

`observation(obs)` ¶

Append forecast to base observation.

Appends a horizon-length demand forecast to every observation. Extends observation_space automatically (base dim + horizon).

Parameter	Default	Description
`horizon`	`6`	Number of future steps to append
`mode`	`'perfect'`	`'perfect'` (ground truth), `'noisy'` (Gaussian noise), or `'none'` (zeros)
`noise_std`	`0.02`	Fractional noise std for `mode='noisy'` (e.g. 0.02 = 2 %)
`normalize`	`True`	Divide forecast values by dataset maximum

from powerzoo.wrappers import GymnasiumWrapper, ForecastWrapper

env = ForecastWrapper(GymnasiumWrapper(TransGridEnv()), horizon=6, mode='noisy')
obs, info = env.reset(seed=0)
# obs[-6:] contains the next 6 half-hour demand values

`powerzoo.wrappers.marl_wrapper.MARLWrapper(env, agent_type='generators', render_mode=None)` ¶

Bases: ParallelEnv

PettingZoo Parallel wrapper for PowerZoo GridEnv.

Parameters¶

env : GridEnv An initialised (and optionally resource-populated) PowerZoo grid env. agent_type : str 'generators' or 'resources'. render_mode : str or None Passed through to env.render_mode if set.

`reset(seed=None, options=None)` ¶

`step(actions)` ¶

Converts a single-agent PowerZoo env to the PettingZoo Parallel API. Each resource registered in the underlying GridEnv becomes an independent agent.

from powerzoo.wrappers import MARLWrapper

env = MARLWrapper(TransGridEnv(), agent_type='generators')
obs, info = env.reset(seed=0)
actions = {a: env.action_space(a).sample() for a in env.agents}
obs, rewards, terminations, truncations, info = env.step(actions)

`powerzoo.wrappers.flatten.FlattenWrapper(env, resource_names=None, obs_keys=None, custom_obs_fn=None, custom_action_fn=None)` ¶

Bases: Wrapper

Combined wrapper for flattening both observation and action spaces

Intelligently flattens dict spaces based on controlled resources.

Parameters:

Name	Type	Description	Default
`env`	`Env`	Environment to wrap (should be PowerEnv instance)	required
`resource_names`	`Optional[List[str]]`	List of resource names to control (e.g., ['bat0', 'wind_0']) If None, controls all resources	`None`
`obs_keys`	`Optional[List[str]]`	List of observation keys to include (e.g., ['grid', 'resources', 'time']) If None, includes all observations	`None`
`custom_obs_fn`	`Optional[Callable]`	Optional custom observation flattening function	`None`
`custom_action_fn`	`Optional[Callable]`	Optional custom action mapping function	`None`

Example

from powerzoo.envs.power_env import PowerEnv from powerzoo.wrappers.flatten import FlattenWrapper

env = PowerEnv(config) env = FlattenWrapper(env, resource_names=['bat0'])

from stable_baselines3 import PPO model = PPO('MlpPolicy', env, verbose=1) model.learn(total_timesteps=10000)

`reset(**kwargs)` ¶

Reset and flatten observation

`step(action)` ¶

Convert flat action to dict and step

Flattens dict/nested observation_space and action_space into 1-D Box spaces. Useful for algorithms that require flat vector inputs.

from powerzoo.wrappers import GymnasiumWrapper, FlattenWrapper

env = FlattenWrapper(GymnasiumWrapper(TransGridEnv()))

Wrappers¶

powerzoo.wrappers.gym_wrappers.GymnasiumWrapper(env) ¶

reset(*, seed=None, options=None, **kwargs) ¶

step(action) ¶

powerzoo.wrappers.gym_wrappers.NormalizationWrapper(env, clip=True) ¶

powerzoo.wrappers.safe_rl_wrapper.SafeRLWrapper(env, cost_threshold=None) ¶

step(action) ¶

powerzoo.wrappers.safe_rl_wrapper.GymnasiumSafeWrapper(env, cost_threshold=None) ¶

step(action) ¶

powerzoo.wrappers.forecast_wrapper.ForecastWrapper(env, horizon=6, mode='perfect', noise_std=0.02, normalize=True) ¶

Parameters¶

reset(*, seed=None, options=None, **kwargs) ¶

observation(obs) ¶

powerzoo.wrappers.marl_wrapper.MARLWrapper(env, agent_type='generators', render_mode=None) ¶

Parameters¶

reset(seed=None, options=None) ¶

step(actions) ¶

powerzoo.wrappers.flatten.FlattenWrapper(env, resource_names=None, obs_keys=None, custom_obs_fn=None, custom_action_fn=None) ¶

reset(**kwargs) ¶

step(action) ¶

`powerzoo.wrappers.gym_wrappers.GymnasiumWrapper(env)` ¶

`reset(*, seed=None, options=None, **kwargs)` ¶

`step(action)` ¶

`powerzoo.wrappers.gym_wrappers.NormalizationWrapper(env, clip=True)` ¶

`powerzoo.wrappers.safe_rl_wrapper.SafeRLWrapper(env, cost_threshold=None)` ¶

`step(action)` ¶

`powerzoo.wrappers.safe_rl_wrapper.GymnasiumSafeWrapper(env, cost_threshold=None)` ¶

`step(action)` ¶

`powerzoo.wrappers.forecast_wrapper.ForecastWrapper(env, horizon=6, mode='perfect', noise_std=0.02, normalize=True)` ¶

`reset(*, seed=None, options=None, **kwargs)` ¶

`observation(obs)` ¶

`powerzoo.wrappers.marl_wrapper.MARLWrapper(env, agent_type='generators', render_mode=None)` ¶

`reset(seed=None, options=None)` ¶

`step(actions)` ¶

`powerzoo.wrappers.flatten.FlattenWrapper(env, resource_names=None, obs_keys=None, custom_obs_fn=None, custom_action_fn=None)` ¶

`reset(**kwargs)` ¶

`step(action)` ¶