Distribution environments

Public API for the radial distribution envs and their backward / forward sweep solvers. For physics, see Physics → Distribution.

DistGridEnv

from powerzoojax.envs import DistGridEnv, make_dist_params

Step contract

• reward = -loss_penalty_weight * p_loss_MW.
• costs = [cost_voltage_violation, cost_thermal_overload, cost_resource].
• info["cost_sum"] = sum(costs) is the aggregate diagnostic.
• The DSO benchmark selects only "voltage_violation" at the task / wrapper layer. cost_mode is deprecated and kept only for backward-compatible config loading.
• info["cost_continuous"] is a continuous voltage / apparent-power overload diagnostic.
• info["soc_terminal_sq"] is populated on the terminal transition when a battery bundle is the first attached bundle.

Observation layout

[v_mag_norm | p_branch_norm | q_branch_norm | p_load_norm | q_load_norm | sin(t) | cos(t) | <bundle_obs>]. This is the policy input vector seen by DistGridEnv. For field semantics and the DSO-specific FlexLoadBundle slice, see Physics → Distribution.

powerzoojax.envs.grid.dist.DistGridEnv

Bases: Environment

Distribution grid environment — BFS power flow.

Observation layout

node voltages (n_nodes,) — (v - 1) / 0.1 line active flows (n_lines,) — p_branch / p_load_ref line reactive flows (n_lines,) — q_branch / q_load_ref node active loads (n_nodes,) — p_load / p_load_ref node reactive loads (n_nodes,) — q_load / q_load_ref [time_sin, time_cos] (2,) — appended per bundle

Reward

-loss_penalty_weight * p_loss_MW (operational cost only)

Constraint cost channels

cost_voltage_violation: count of nodes violating voltage limits cost_thermal_overload: count of lines violating thermal limits (|S| > cap) cost_resource: bundle clipping costs cost_sum: total violation count

reset(key, params)

step(key, state, action, params)

Run one environment step.

Parameters:

Name	Type	Description	Default
key	PRNGKey	PRNG key.	required
state	DistGridState	Current grid state.	required
action	Array	When include_der=True: (n_nodes-1 + sum(bundle.action_dim),) in [-1, 1]. When include_der=False: (sum(bundle.action_dim),) — bundle actions only.	required
params	DistGridParams	Environment parameters.	required

observation_space(params)

action_space(params)

powerzoojax.envs.grid.dist.DistGridState

DistGridEnv state — pure grid quantities only.

Resource state (battery SOC, etc.) lives in resource_states. episode_batt_soc_init stores SOC at episode start for terminal penalties.

powerzoojax.envs.grid.dist.DistGridParams

DistGridEnv parameters.

Scalar configuration fields are marked pytree_node=False so they are treated as static compile-time constants by jax.jit.

Action scaling note: der_low and der_high are in per-unit (p.u.) relative to base_mva. With the default base_mva = 100, der_low = -0.005 corresponds to −500 kW and der_high = 0.005 to +500 kW. Do not set them to physical MW values directly.

powerzoojax.envs.grid.dist.make_dist_params(case, max_steps=48, steps_per_day=48, v_slack=1.0, v_min=0.9, v_max=1.1, loss_penalty_weight=0.1, resources=(), include_der=True, cost_mode='aggregate')

Create DistGridParams from a CaseData.

Parameters:

Name	Type	Description	Default
case	CaseData	CaseData instance.	required
max_steps	int	Maximum steps per episode.	48
steps_per_day	int	Steps per day for time encoding.	48
v_slack	float	Slack bus voltage magnitude [p.u.].	1.0
v_min	float	Minimum allowed voltage [p.u.].	0.9
v_max	float	Maximum allowed voltage [p.u.].	1.1
loss_penalty_weight	float	Weight on active-loss reward shaping.	0.1
resources	tuple	Tuple of ResourceBundle instances (e.g. BatteryBundle). Use make_battery_bundle(case, bus_ids=[...]) to create bundles.	()
include_der	bool	If False, action space contains only bundle actions (no legacy DER injection at each bus). Use this when the agent should only control attached resources (e.g. batteries).	True
cost_mode	str	Deprecated legacy field kept for backward-compatible config loading. Core env semantics no longer depend on it; task-level constraint selection decides which channels a benchmark uses.	'aggregate'

Returns:

Type	Description
DistGridParams	DistGridParams ready for DistGridEnv.

DistGrid3PhaseEnv

from powerzoojax.envs import DistGrid3PhaseEnv, make_dist_3phase_params

Step contract

• reward = -loss_penalty_weight * p_loss_MW.
• costs = [cost_voltage_violation, cost_thermal_overload, cost_vuf_violation, cost_resource].
• info["cost_sum"] = sum(costs) is the aggregate diagnostic.
• info["max_vuf_percent"] keeps the worst per-node VUF (Fortescue voltage unbalance factor in %).

powerzoojax.envs.grid.dist_3phase.DistGrid3PhaseEnv

Bases: Environment

Three-phase distribution grid environment — BFS 3-phase power flow.

Observation layout

V_A (n_nodes,) — (v - 1) / 0.1 V_B (n_nodes,) V_C (n_nodes,) P_A (n_lines,) — P_branch / p_load_ref P_B (n_lines,) P_C (n_lines,) Q_A (n_lines,) — Q_branch / q_load_ref Q_B (n_lines,) Q_C (n_lines,) p_load_A (n_nodes,) — p_load / p_load_ref p_load_B (n_nodes,) p_load_C (n_nodes,) q_load_A (n_nodes,) — q_load / q_load_ref q_load_B (n_nodes,) q_load_C (n_nodes,) [time_sin, time_cos] (2,) battery_soc (n_batteries,) — omitted if no resources

reset(key, params)

step(key, state, action, params)

Run one environment step.

Parameters:

Name	Type	Description	Default
key	PRNGKey	PRNG key.	required
state	DistGrid3PhState	Current grid state.	required
action	Array	When include_der=True: (3*(n_nodes-1) + sum(bundle.action_dim),) in [-1, 1]. When include_der=False: (sum(bundle.action_dim),) — bundle actions only.	required
params	DistGrid3PhParams	Environment parameters.	required

observation_space(params)

action_space(params)

powerzoojax.envs.grid.dist_3phase.DistGrid3PhState

3-phase distribution grid state — pure grid quantities only.

Resource state (battery SOC, etc.) is NOT stored here.

powerzoojax.envs.grid.dist_3phase.DistGrid3PhParams

3-phase distribution grid parameters.

Scalar configuration fields are marked pytree_node=False so they are treated as static compile-time constants by jax.jit.

Action scaling note: der_low and der_high are in per-unit (p.u.) relative to base_mva. With the default base_mva = 10, der_low = -0.1 corresponds to −1 MW.

powerzoojax.envs.grid.dist_3phase.make_dist_3phase_params(n_nodes, from_nodes, to_nodes, Z_3ph_pu, load_P_3ph, load_Q_3ph, *, ref_bus=0, v_ref_mag=1.0, v_min=0.9, v_max=1.1, vuf_max=2.0, max_steps=48, steps_per_day=48, base_mva=10.0, loss_penalty_weight=0.1, line_cap=None, resources=(), include_der=True)

Create DistGrid3PhParams from raw 3-phase network data.

Parameters:

Name	Type	Description	Default
n_nodes	int	Number of buses.	required
from_nodes		(n_lines,) from-bus indices.	required
to_nodes		(n_lines,) to-bus indices.	required
Z_3ph_pu		(n_lines, 3, 3) complex impedance matrices per branch [p.u.].	required
load_P_3ph	ndarray	(T, 3*(n_nodes-1)) per-phase active load [p.u.].	required
load_Q_3ph	ndarray	(T, 3*(n_nodes-1)) per-phase reactive load [p.u.].	required
ref_bus	int	Reference (slack) bus index.	0
v_ref_mag	float	Reference voltage magnitude.	1.0
v_min	float	Voltage lower bound.	0.9
v_max	float	Voltage upper bound.	1.1
vuf_max	float	Maximum voltage unbalance factor (%).	2.0
max_steps	int	Maximum steps per episode.	48
steps_per_day	int	Steps per day for time encoding.	48
base_mva	float	System base power (MVA).	10.0
loss_penalty_weight	float	Weight on active-loss reward shaping.	0.1
line_cap		(n_lines,) aggregate apparent-power capacity per line [MVA]. 0 means no limit (thermal check skipped for that line). None defaults to all zeros (no thermal limits).	None
resources	tuple	Tuple of ResourceBundle instances (e.g. BatteryBundle). Bundle bus_idx should be 0-based bus indices (not ref_bus). Bundle injection is balanced three-phase (1/3 per phase).	()

Returns:

Type	Description
DistGrid3PhParams	DistGrid3PhParams ready for DistGrid3PhaseEnv.

Balanced BFS solver

powerzoojax.envs.grid.bfs_power_flow.prepare_bfs(case)

Build BFS topology data from a CaseData with distribution parameters.

Inactive lines (status=0, e.g. normally-open tie switches) are excluded so that the BFS spanning tree reflects the actual operating topology.

Parameters:

Name	Type	Description	Default
case	CaseData	CaseData with line_r, line_x, line_from_idx, line_to_idx.	required

Returns:

Type	Description
BFSTopoData	BFSTopoData ready for JIT-compiled BFS solver.

powerzoojax.envs.grid.bfs_power_flow.bfs_power_flow(topo, p_load_pu, q_load_pu, v_slack=1.0, max_iter=100, tol=1e-06)

Run BFS power flow via lax.while_loop (JIT-compilable).

Parameters:

Name	Type	Description	Default
topo	BFSTopoData	BFSTopoData from build_radial_topology or prepare_bfs.	required
p_load_pu	Array	(n_nodes,) active load [p.u.].	required
q_load_pu	Array	(n_nodes,) reactive load [p.u.].	required
v_slack	float	Slack bus voltage [p.u.].	1.0
max_iter	int	Maximum BFS iterations.	100
tol	float	Convergence tolerance on max|ΔV²|.	1e-06

Returns:

Type	Description
BFSResult	BFSResult with voltages, branch flows, losses, convergence info.

powerzoojax.envs.grid.bfs_power_flow.BFSTopoData

Precomputed radial topology data (built once at env construction).

All matrices are dense JAX arrays for JIT/vmap compatibility.

Three-phase BFS solver

powerzoojax.envs.grid.bfs_3phase_power_flow.build_3phase_topology(n_nodes, from_nodes, to_nodes, Z_3ph_pu, ref_bus=0, v_ref_mag=1.0)

Build 3-phase topology with BIBC/BCBV/DLF matrices.

Parameters:

Name	Type	Description	Default
n_nodes	int	Number of buses.	required
from_nodes	ndarray	(n_lines,) from-bus indices.	required
to_nodes	ndarray	(n_lines,) to-bus indices.	required
Z_3ph_pu	ndarray	(n_lines, 3, 3) complex impedance matrices per branch [p.u.].	required
ref_bus	int	Reference (slack) bus index.	0
v_ref_mag	float	Reference voltage magnitude.	1.0

Returns:

Type	Description
ThreePhaseTopoData	ThreePhaseTopoData with dense JAX arrays.

powerzoojax.envs.grid.bfs_3phase_power_flow.bfs_3phase_power_flow(topo, P_3ph_pu, Q_3ph_pu, max_iter=100, tol=1e-06)

Run 3-phase BFS power flow via lax.while_loop.

Conventions

P_3ph_pu, Q_3ph_pu are LOADS per bus-phase (positive = consumption). Internally: S_inj = -(P + jQ), I_inj = conj(S_inj / V).

Parameters:

Name	Type	Description	Default
topo	ThreePhaseTopoData	ThreePhaseTopoData from build_3phase_topology.	required
P_3ph_pu	Array	(3*n_lines,) or (n_lines, 3) active load [p.u.].	required
Q_3ph_pu	Array	(3*n_lines,) or (n_lines, 3) reactive load [p.u.].	required
max_iter	int	Max iterations.	100
tol	float	Convergence tolerance on max |ΔV|.	1e-06

Returns:

Type	Description
BFS3PhResult	BFS3PhResult with voltages, currents, branch powers, convergence info.

powerzoojax.envs.grid.bfs_3phase_power_flow.ThreePhaseTopoData

Precomputed 3-phase radial topology.