Skip to content

RL — multi-agent

Public API for multi-agent wrappers and trainers. For single-agent, see API → RL. For the conceptual overview, see Training → Trainers.

Multi-agent base contract

from powerzoojax.rl import MultiAgentEnvironment, MARLState

MultiAgentEnvironment declares the dict-API contract used by IPPO trainers:

  • reset(key, params) returns (obs_dict, state).
  • step(key, state, action_dict, params) returns (obs_dict, state, reward, done, info).

powerzoojax.rl.multi_agent.MultiAgentEnvironment

Bases: ABC

JaxMARL-compatible multi-agent environment base class.

Key differences from single-agent Environment: - obs, rewards, dones are Dict[str, Array] - actions are Dict[str, Array] - dones must include "all" key - params are bound at construction time

step(key, state, actions) abstractmethod

Returns (obs_dict, state, rewards_dict, dones_dict, info).

powerzoojax.rl.multi_agent.MARLState

Multi-agent environment state — thin wrapper around TransGridState.

All physical quantities (unit_power_mw, resource SOC, time_step, done, ...) are accessible via state.grid_state.

Grid MARL adapters

from powerzoojax.rl import GridMARLEnv, DistGridMARLEnv

GridMARLEnv wraps TransGridEnv. Generator agents are named unit_i; bundle devices are named by resource type (battery_0, pv_0, flexload_0, ...). Reward is shared across agents.

DistGridMARLEnv wraps DistGridEnv. Supports observation_mode="global" (default) and observation_mode="local" (Dec-POMDP K-hop neighborhood, used by the DERs benchmark).

powerzoojax.rl.multi_agent.GridMARLEnv(grid_env, grid_params, reward_mode='shared')

Bases: MultiAgentEnvironment

Multi-agent power grid environment.

Each generator unit and each resource device is an independent agent. Resources are read from grid_params.resources (ResourceBundle API).

Parameters:

Name Type Description Default
grid_env TransGridEnv

TransGridEnv instance.

 required
grid_params TransGridParams

TransGridParams with case data, load profiles, and optional resource bundles attached via resources=(bundle,).

 required
reward_mode str

'shared' (all agents get same reward).

 'shared'

Agent layout (action/obs): - Unit agents: action shape (1,) in [-1,1] normalised dispatch. - Resource agents: action shape (1,) per device in [-1,1].

Obs layout per agent
  • Shared obs: grid core obs (line flows, loads, unit dispatch, time)
  • Unit agent i: shared + unit_norm[i] (own dispatch fraction)
  • Resource device j: shared + device_obs_slice (e.g. [soc, p_norm])

Usage::

bundle = make_battery_bundle(case, bus_ids=[1, 3], power_mw=20)
params = make_trans_params(case, resources=(bundle,), max_steps=48)
env = GridMARLEnv(TransGridEnv(), params)
obs_dict, state = env.reset(key)
obs_dict, state, rewards, dones, info = env.step(key, state, actions_dict)

powerzoojax.rl.multi_agent.DistGridMARLEnv(dist_env, dist_params, voltage_penalty=0.0, soc_penalty=0.0, reward_mode='shared', observation_mode='global')

Bases: MultiAgentEnvironment

Multi-agent distribution grid environment — one agent per resource device.

Wraps DistGridEnv so each attached battery (or other resource) becomes an independent agent. Requires include_der=False in dist_params.

All physics are delegated to DistGridEnv.step(). Reward is shared: all agents receive the same scalar reward. An optional voltage_penalty adds a continuous voltage-deviation term on top of the base loss reward::

effective_reward = base_reward - voltage_penalty * cost_continuous

Parameters:

Name Type Description Default
dist_env DistGridEnv

DistGridEnv instance.

 required
dist_params DistGridParams

DistGridParams with resources attached and include_der=False.

 required
voltage_penalty float

Weight on the continuous voltage-deviation cost (v_under + v_over summed across buses). Set > 0 to shape rewards toward voltage regulation.

 0.0
soc_penalty float

Weight on the terminal SOC deviation penalty (info["soc_terminal_sq"], only non-zero at episode end). Set > 0 to discourage extreme terminal SOC.

 0.0
reward_mode str

"shared" — all agents get the same reward.

 'shared'
Agent naming

One agent per device, globally indexed per resource type: "battery_0", "battery_1", …

Obs layout per agent

[grid_core | own_device_obs] padded to uniform obs_dim. Grid core = v_norm(N) + p_flow_norm(L) + q_flow_norm(L) + p_load_norm(N) + q_load_norm(N) + sin + cos.

Action per agent

Shape (per_device_action_dim,) in [-1, 1]. P-only (default): (1,); P+Q (enable_q_control=True): (2,).

Usage::

bundle = make_battery_bundle(case, bus_ids=[18, 25, 33], power_mw=0.75)
params = make_dist_params(case, resources=[bundle], include_der=False)
env = DistGridMARLEnv(DistGridEnv(), params, voltage_penalty=8.0)
obs_dict, state = env.reset(key)
obs_dict, state, rewards, dones, info = env.step(key, state, actions)

Market MARL adapter

from powerzoojax.rl import MarketMARLEnv

See API → Market MARL for the GenCos task wrapper.

IPPO trainers

from powerzoojax.rl import make_ippo_train, make_ippo_act
  • algo="ippo" — independent PPO with full parameter sharing.
  • algo="ippo_typed" — typed parameter sharing: agents partitioned by name prefix (battery_*, renewable_*, flexload_*, ...). Each type has an independent SharedActorCritic.

For the DERs benchmark, algo="ippo_typed" is required (see Benchmarks → DERs).

make_ippo_act(env, params) returns a deterministic per-agent action function suitable for evaluation rollouts.

powerzoojax.rl.ippo.make_ippo_train(env, config)

Create a JIT-compiled IPPO training function for GridMARLEnv.

Parameters:

Name Type Description Default
env

GridMARLEnv instance (params bound at construction).

 required
config TrainConfig

TrainConfig with algo="ippo" and common PPO hyperparameters.

 required

Returns:

Type Description
Callable

train_fn(key: PRNGKey) -> TrainResult

The returned function is JIT-compiled. Call it once to compile + train: result = train_fn(jax.random.PRNGKey(42))

powerzoojax.rl.ippo.make_ippo_act(result)

Create a deterministic per-agent act function from an IPPO TrainResult.

Parameters:

Name Type Description Default
result

TrainResult returned by make_ippo_train or train(). result.params must be a SharedActorCritic pytree. result.config.hidden_dims sets the network architecture.

 required

Returns:

Type Description
Callable

act_fn(obs: Array) -> Array — shape (action_dim,) clipped to
Callable

[-1, 1]. Suitable for calling per-agent inside an evaluation loop.

Example::

act = make_ippo_act(result)
for name in env.agent_names:
    actions[name] = act(obs_dict[name])