Resources¶

Public API for standalone resource environments and the bundle protocol used to attach devices to grid, market, and microgrid envs. For physics, see Physics → Resources. For the composed data-center microgrid, see API → Microgrid.

Quick orientation¶

Use the resource surface in two modes:

standalone *Env: run one device by itself
*Bundle: attach many devices to a parent env through params.resources

Not every resource provides both forms:

Resource	Standalone env	Bundle
Battery	`BatteryEnv`	`BatteryBundle`
Renewable	`RenewableEnv` / `SolarEnv` / `WindEnv`	`RenewableBundle`
Vehicle	`VehicleEnv`	no public bundle
FlexLoad	`FlexLoadEnv`	`FlexLoadBundle`
Diesel	no standalone env	`DieselBundle`
Data center	`DataCenterEnv`	no public bundle

Cost semantics by resource¶

Resource	`reward`	`costs`
`BatteryEnv`	`0.0`	cycle-throughput cost
`RenewableEnv`	`0.0`	always `0.0`
`VehicleEnv`	`0.0`	departure SOC shortfall
`FlexLoadEnv`	`0.0`	curtailment + deferred-demand discomfort + simultaneous activation penalty
`DataCenterEnv`	`0.0`	expired-task SLA density
`DieselBundle`	no standalone reward	`cost_info["cost"] = 0.0`; fuel and carbon reported separately

Diagnostics not folded into costs:

BatteryEnv: cost_action_clip
DataCenterEnv: cost_overtemp
DieselBundle: fuel_cost, carbon_kg

`BatteryEnv` and `BatteryBundle`¶

from powerzoojax.envs import BatteryEnv
from powerzoojax.envs.resource.battery import BatteryBundle, make_battery_bundle

Sign convention: P > 0 discharge (inject), P < 0 charge (draw). Feasible power is clipped by converter rating and SOC headroom; SOC update uses one-way charge / discharge efficiency.

`powerzoojax.envs.resource.battery.BatteryEnv` ¶

Bases: Environment

Battery Energy Storage System environment.

scalar in [-1, 1], denormalized internally to [-power_mw, power_mw] MW.

Positive = discharge (inject to grid), negative = charge (draw from grid).

Observation: 6-D [soc, p_norm, headroom_norm, charge_room_norm, time_sin, time_cos].

`reset(key, params)` ¶

`step(key, state, action, params)` ¶

`observation_space(params)` ¶

`action_space(params)` ¶

`powerzoojax.envs.resource.battery.BatteryState` ¶

Bases: ResourceState

Battery state extending ResourceState.

Attributes:

Name	Type	Description
`soc`	`Array`	State of charge [0, 1].
`done`	`Array`	Episode termination flag.

`powerzoojax.envs.resource.battery.BatteryParams` ¶

Bases: ResourceParams

Battery parameters extending ResourceParams.

Attributes:

Name	Type	Description
`capacity_mwh`	`float`	Energy capacity [MWh].
`power_mw`	`float`	Maximum charge/discharge power [MW].
`eta_charge`	`float`	One-way charging efficiency [0, 1].
`eta_discharge`	`float`	One-way discharging efficiency [0, 1].
`soc_min`	`float`	Minimum allowed SOC.
`soc_max`	`float`	Maximum allowed SOC.
`initial_soc`	`float`	SOC at reset.

All scalar fields are pytree_node=False — static under JIT.

`powerzoojax.envs.resource.battery.make_battery_params(bus_id=-1, capacity_mwh=50.0, power_mw=20.0, eta_charge=None, eta_discharge=None, eta_roundtrip=None, soc_min=0.1, soc_max=0.9, initial_soc=0.5, delta_t_hours=0.5, steps_per_day=48, max_steps=48, cycle_cost_per_mwh=0.0)` ¶

Build BatteryParams with the same efficiency rules as PowerZoo BatteryEnv.

Parameters:

Name	Type	Description	Default
`bus_id`	`int`	Connection bus ID for grid attachment. -1 = unregistered sentinel. Set this when passing the params to `make_dist_params(resources=[...])`.	`-1`
`capacity_mwh`	`float`	Energy capacity [MWh].	`50.0`
`power_mw`	`float`	Maximum charge/discharge power [MW].	`20.0`

Omitted one-way efficiencies default to 0.95 unless eta_roundtrip is set, in which case missing sides use sqrt(eta_roundtrip).

`powerzoojax.envs.resource.battery.compute_feasible_power(soc, desired_power, params)` ¶

Compute feasible grid-side power respecting SOC and power limits.

Discharge (P > 0): max_discharge = (soc - soc_min) * capacity * η_d / dt Charge (P < 0): max_charge = (soc_max - soc) * capacity / (η_c * dt)

`powerzoojax.envs.resource.battery.update_soc(soc, power, params)` ¶

Update SOC based on grid-side power (Coulomb counting with losses).

Discharge (P > 0): Δsoc = -P * dt / (η_d * capacity) Charge (P < 0): Δsoc = -P * dt * η_c / capacity (P is negative, so -P is positive)

`powerzoojax.envs.resource.battery.BatteryBundle` ¶

Struct-of-arrays battery bundle for grid env attachment.

Groups N batteries with identical action/obs dimensions. Works in MW throughout; the grid env handles unit conversion.

When enable_q_control=False (default, P-only mode): Action per device: [P_norm] Obs per device: [soc, p_norm] When enable_q_control=True (P+Q mode): Action per device: [P_norm, Q_norm] Obs per device: [soc, p_norm, q_norm] PQ circle constraint: P² + Q² ≤ S_rated² (P-priority projection).

`reset(key)` ¶

Return initial bundle state (deterministic or uniform-random SOC).

`step(state, action, ctx)` ¶

Run one step for all N devices.

Parameters:

Name	Type	Description	Default
`state`	`BatteryBundleState`	Current bundle state.	required
`action`	`Array`	(action_dim,) in [-1, 1]. P-only: one scalar per device. P+Q: [P_norm, Q_norm] per device (device-major, interleaved).	required
`ctx`	`Dict[str, Any]`	Grid-provided context (unused by battery, reserved).	required

Returns:

Type	Description
`Tuple[BatteryBundleState, Array, Array, Array, Dict[str, Array]]`	(new_state, p_inject_mw, q_inject_mvar, obs_slice, cost_info)

`observe(state, ctx)` ¶

Rebuild obs_slice from state without stepping (for _get_obs / reset).

`powerzoojax.envs.resource.battery.BatteryBundleState` ¶

State for a BatteryBundle — SOC of all N devices.

`powerzoojax.envs.resource.battery.make_battery_bundle(case=None, bus_ids=None, *, n_devices=None, power_mw=20.0, capacity_mwh=50.0, s_rated_mva=None, enable_q_control=False, soc_min=0.1, soc_max=0.9, eta_charge=0.95, eta_discharge=0.95, initial_soc=0.5, dt_hours=0.5, randomize_initial_soc=False, soc_init_low=0.3, soc_init_high=0.7, cycle_cost_per_mwh=0.0)` ¶

Create a BatteryBundle.

Two construction modes:

Grid mode (case and bus_ids both provided): Resolves bus_idx through case.node_ids. n_devices equals len(bus_ids). Used by DistGridEnv/TransGridEnv.

Busless mode (case=None, bus_ids=None): For behind-the-meter envs (e.g. DataCenterMicrogridEnv) where no power-flow exists. n_devices must be supplied (default 1). bus_idx = arange(n_devices) and the host env sums injections directly rather than scattering by index.

All broadcasting and index lookup is done in NumPy at init time (CPU), then converted to jnp arrays. NOT called inside JIT.

Parameters:

Name	Type	Description	Default
`case`		CaseData instance with `node_ids` attribute (optional).	`None`
`bus_ids`		External node IDs (optional).	`None`
`n_devices`		Number of batteries (busless mode; ignored if `bus_ids`).	`None`
`power_mw`		Max power per device [MW]. Scalar or Sequence[float].	`20.0`
`capacity_mwh`		Energy capacity per device [MWh].	`50.0`
`s_rated_mva`		Inverter apparent power rating [MVA] per device. Defaults to `power_mw` (inverter sized to active power rating). Set to e.g. `1.1 * power_mw` for extra reactive headroom.	`None`
`enable_q_control`	`bool`	If True, action space includes reactive power (per_device_action_dim=2, per_device_obs_dim=3) with PQ circle constraint. Default False preserves P-only behavior.	`False`
`soc_min / soc_max`		SOC bounds.	required
`eta_charge / eta_discharge`		One-way efficiencies.	required
`initial_soc`		Initial SOC (used when `randomize_initial_soc` is False).	`0.5`
`dt_hours`		Time-step duration [h].	`0.5`
`randomize_initial_soc`	`bool`	If True, sample initial SOC per device from `Uniform(soc_init_low, soc_init_high)` at each reset (training).	`False`
`soc_init_low / soc_init_high`		Bounds for random initial SOC.	required

Returns:

Type	Description
`BatteryBundle`	BatteryBundle with (n_devices,) arrays.

`RenewableEnv`, `SolarEnv`, `WindEnv`, and `RenewableBundle`¶

from powerzoojax.envs import RenewableEnv, SolarEnv, WindEnv
from powerzoojax.envs.resource.renewable import RenewableBundle, make_renewable_bundle

Action a in [-1, 1] maps to curtailment (1 - a) / 2. Output equals capacity_mw * capacity_factor * (1 - curtailment). With enable_q_control=True, reactive power is clipped by inverter PQ-circle.

`powerzoojax.envs.resource.renewable.RenewableEnv` ¶

Bases: Environment

Renewable (Solar/Wind) environment.

Action: scalar in [-1, 1], mapped internally to curtailment [0, 1]. Observation: 4-D [capacity_factor, p_norm, time_sin, time_cos].

`powerzoojax.envs.resource.renewable.RenewableState` ¶

Bases: ResourceState

Renewable resource state.

Attributes:

Name	Type	Description
`capacity_factor`	`Array`	Current capacity factor [0, 1] before curtailment.
`done`	`Array`	Episode termination flag.

`powerzoojax.envs.resource.renewable.RenewableParams` ¶

Bases: ResourceParams

Renewable resource parameters.

Attributes:

Name	Type	Description
`capacity_mw`	`float`	Installed nameplate capacity [MW].
`profiles`	`Array`	Pre-loaded capacity factor time series, shape (n_steps,), normalized [0, 1].
`allow_curtailment`	`int`	Whether curtailment action is honored (bool-like int).
`enable_q_control`	`bool`	If True, action is 2-D [curtailment_norm, q_norm] and obs includes q_norm. PQ circle constraint is applied (P-priority).
`s_rated_mva`	`float`	Inverter apparent power rating [MVA]. Only used when `enable_q_control=True`. Defaults to `capacity_mw`.

Scalar fields are pytree_node=False — static under JIT. profiles remains a traced pytree leaf (array data).

`powerzoojax.envs.resource.renewable.SolarEnv` ¶

Bases: RenewableEnv

Solar PV — subclass with bell-curve default profile.

`powerzoojax.envs.resource.renewable.WindEnv` ¶

Bases: RenewableEnv

Wind turbine — subclass with variable default profile.

`powerzoojax.envs.resource.renewable.RenewableBundle` ¶

Struct-of-arrays renewable bundle for grid env attachment.

Groups N inverter-interfaced renewable generators (PV or wind) driven by pre-loaded capacity-factor profiles. Supports active curtailment and reactive power support with an inverter PQ-circle constraint (P-priority projection, matching BatteryBundle).

Action / observation dimensions are controlled by two static flags:

enable_curtailment (default True) When True, the action vector includes a per-device curtailment scalar in [-1, 1] (+1 = no curtailment, -1 = full curtailment). The observation vector includes a curtail_norm channel. enable_q_control (default True) When True, the action vector includes a per-device reactive scalar in [-1, 1] (scaled by s_rated then PQ-circle clipped with P-priority). The observation vector includes a q_norm channel.

Resulting per-device dimensions:

+--------------------+----------------+-------------+-------------------+ | (curtail, Q) | action_dim | obs_dim | obs layout | +--------------------+----------------+-------------+-------------------+ | (True, True) | 2 (default) | 4 (default) | cf, p_norm, q, c | | (True, False) | 1 | 3 | cf, p_norm, c | | (False, True) | 1 | 3 | cf, p_norm, q | | (False, False) | 0 | 1 | cf | +--------------------+----------------+-------------+-------------------+

The (False, False) mode is used by behind-the-meter envs (such as DataCenterMicrogridEnv) where PV is purely exogenous.

Action layout (when present): [curtail, q] per device — that is, [curtail_0, q_0, curtail_1, q_1, ...]; if only one of curtail / Q is enabled the corresponding scalar fills slot 0.

`reset(key)` ¶

Return the initial state — t=0, outputs at profile[0].

`step(state, action, ctx)` ¶

Run one step for all N devices.

Parameters:

Name	Type	Description	Default
`state`	`RenewableBundleState`	Current bundle state.	required
`action`	`Array`	`(action_dim,)` in `[-1, 1]`. Layout depends on the static `enable_curtailment` / `enable_q_control` flags (see class docstring). Empty array when both flags False.	required
`ctx`	`Dict[str, Any]`	Grid-provided context (unused; reserved).	required

Returns:

Type	Description
`Tuple[RenewableBundleState, Array, Array, Array, Dict[str, Array]]`	`(new_state, p_inject_mw, q_inject_mvar, obs_slice, cost_info)`

`observe(state, ctx)` ¶

Flat obs slice of shape (obs_dim,).

Per-device channels included (in order): cf always; p_norm if any control enabled; q_norm if enable_q_control; curtail_norm if enable_curtailment. See class docstring for the resulting dimensions.

`powerzoojax.envs.resource.renewable.RenewableBundleState` ¶

State for a RenewableBundle — time index + last-step outputs.

Attributes:

Name	Type	Description
`t`	`Array`	Internal time step (int32 scalar). Used to index `profiles`.
`cf`	`Array`	Last-step capacity factor, shape `(n_devices,)`.
`p_mw`	`Array`	Last-step active output [MW], shape `(n_devices,)`.
`q_mvar`	`Array`	Last-step reactive output [MVAr], shape `(n_devices,)`.
`curtail_frac`	`Array`	Last-step curtailment fraction ∈ [0,1], shape `(n_devices,)`.

`powerzoojax.envs.resource.renewable.make_renewable_bundle(case=None, bus_ids=None, *, n_devices=None, capacity_mw=0.2, s_rated_mva=None, profiles=None, max_steps=48, allow_curtailment=True, enable_curtailment=True, enable_q_control=True, curtail_cost_per_mwh=0.0, dt_hours=0.5)` ¶

Create a RenewableBundle.

Two construction modes:

Grid mode (case and bus_ids both provided): bus_idx resolved through case.node_ids; n_devices = len(bus_ids). Used by DistGridEnv/TransGridEnv.

Busless mode (case=None, bus_ids=None): For behind-the-meter envs (e.g. DataCenterMicrogridEnv). n_devices must be supplied (default 1). bus_idx = arange(n_devices) and the host env sums injections.

Parameters:

Name	Type	Description	Default
`case`		`CaseData` instance with a `node_ids` attribute (optional).	`None`
`bus_ids`	`Optional[Sequence[int]]`	External node IDs (optional).	`None`
`n_devices`	`Optional[int]`	Number of devices (busless mode; ignored if `bus_ids`).	`None`
`capacity_mw`		Nameplate capacity [MW]. Scalar or length-`n` sequence.	`0.2`
`s_rated_mva`	`Optional[float]`	Inverter apparent power rating [MVA]. Defaults to `1.1 * capacity_mw` to leave reactive headroom at full output.	`None`
`profiles`	`Optional[Array]`	Capacity-factor time series, shape `(T, n)` or `(T,)`. If `None`, a synthetic bell-curve PV profile of length `max_steps` is broadcast to every device.	`None`
`max_steps`	`int`	Used only when `profiles is None` to build the default.	`48`
`allow_curtailment`	`bool`	Soft flag. If False, the curtailment action (if any) is ignored and every device runs at its MPPT point.	`True`
`enable_curtailment`	`bool`	Hard flag. When False, the curtailment scalar is removed from the action / observation vectors entirely.	`True`
`enable_q_control`	`bool`	Hard flag. When False, the reactive-power scalar is removed from the action / observation vectors entirely.	`True`
`curtail_cost_per_mwh`		Opportunity cost for curtailed energy `[$/MWh]`.	`0.0`
`dt_hours`	`float`	Time-step duration.	`0.5`

Returns:

Type	Description
`RenewableBundle`	`RenewableBundle` with `(n_devices,)` SoA arrays.

`VehicleEnv`¶

from powerzoojax.envs import VehicleEnv

Schedule-driven SOC dynamics. Trip energy is subtracted at departure; charge / V2G discharge is allowed only when is_home=1. info["cost"] penalizes leaving below soc_departure_min.

`powerzoojax.envs.resource.vehicle.VehicleEnv` ¶

Bases: Environment

Electric Vehicle environment.

scalar in [-1, 1], denormalized internally to [-p_charge_max, p_discharge_max] MW.

Positive = V2G discharge, negative = G2V charge. Uses piecewise linear mapping: 0 → 0 (idle), +1 → +p_discharge_max, -1 → -p_charge_max.

Observation: 9-D [soc, p_norm, is_home, departure_ready, time_to_dep_norm, time_to_arr_norm, time_sin, time_cos, soc_dep_min].

`powerzoojax.envs.resource.vehicle.VehicleState` ¶

Bases: ResourceState

Vehicle state.

Attributes:

Name	Type	Description
`soc`	`Array`	State of charge [0, 1].
`is_home`	`Array`	1.0 if at home (available for charge/discharge), 0.0 if away.
`time_of_day`	`Array`	Current time in hours [0, 24).
`done`	`Array`	Episode termination flag.

`powerzoojax.envs.resource.vehicle.VehicleParams` ¶

Bases: ResourceParams

Vehicle parameters.

Attributes:

Name	Type	Description
`capacity_mwh`	`float`	Battery capacity [MWh].
`soc_init`	`float`	Initial SOC.
`soc_min`	`/ soc_max`	SOC bounds.
`soc_departure_min`	`float`	Required SOC at departure.
`p_charge_max_mw`	`float`	Max charging power [MW].
`p_discharge_max_mw`	`float`	Max V2G discharge power [MW].
`eta_charge`	`/ eta_discharge`	One-way efficiencies.
`trip_departures`	`Array`	Departure times [hours], shape (MAX_TRIPS,). Unused = inf.
`trip_arrivals`	`Array`	Arrival times [hours], shape (MAX_TRIPS,). Unused = inf.
`trip_energies`	`Array`	Energy consumed per trip [MWh], shape (MAX_TRIPS,). Unused = 0.
`n_trips`	`int`	Number of actual trips (int32).

Scalar fields are pytree_node=False — static under JIT. Trip schedule arrays remain traced pytree leaves.

`powerzoojax.envs.resource.vehicle.make_vehicle_params(E_max_kWh=60.0, soc_init=0.8, soc_min=0.1, soc_max=0.95, soc_departure_min=0.8, p_charge_max_kW=7.0, p_discharge_max_kW=7.0, eta_charge=0.95, eta_discharge=0.95, commute_schedule=None, delta_t_minutes=15.0, steps_per_day=None, max_steps=None)` ¶

Create VehicleParams with kWh/kW inputs (PowerZoo-compatible).

Parameters:

Name	Type	Description	Default
`commute_schedule`	`list`	List of dicts with 'departure', 'arrival', 'energy_kWh'. Default: single trip 8am-6pm, 15kWh.	`None`

`FlexLoadEnv` and `FlexLoadBundle`¶

from powerzoojax.envs import FlexLoadEnv
from powerzoojax.envs.resource.flexload import FlexLoadBundle, make_flexload_bundle

Two actions: curtail now and shift demand out now. Buffered demand is released over the next shift_horizon steps. Net injection follows curtail + shift_out - shift_in. Optional lmp= makes price visible in the observation but does not change physics.

`powerzoojax.envs.resource.flexload.FlexLoadEnv` ¶

Bases: Environment

Flexible Load demand-response environment.

2-D in [0, 1] × [0, 1] (unit-scaled):

action[0] : curtailment fraction × curtail_cap_mw action[1] : shift-out fraction × shift_cap_mw

8-D [curtail_norm, shift_out_norm, shift_in_norm,

buffer_fill_ratio, buffer_energy_norm, time_sin, time_cos, price_norm].

`step(key, state, action, params, lmp=0.0)` ¶

Run one environment step.

Parameters:

Name	Type	Description	Default
`lmp`	`Array`	Locational marginal price for this step [$/MWh]. Pass the market-cleared price (market participation) or the current value from an exogenous RTP/ToU profile. Defaults to 0.0.	`0.0`

`powerzoojax.envs.resource.flexload.FlexLoadState` ¶

Bases: ResourceState

FlexLoad state.

Attributes:

Name	Type	Description
`curtailed_mw`	`Array`	Current step curtailment [MW].
`shift_out_mw`	`Array`	Current step shift-out amount [MW].
`shift_in_mw`	`Array`	Current step shift-in (released deferred demand) [MW].
`deferred_buffer`	`Array`	Ring buffer of per-step deferred release amounts.
`buffer_head`	`Array`	Read pointer into deferred_buffer.
`buffer_size`	`Array`	Number of valid entries in buffer.
`current_lmp`	`Array`	Current locational marginal price [$/MWh].
`done`	`Array`	Episode termination flag.

`powerzoojax.envs.resource.flexload.FlexLoadParams` ¶

Bases: ResourceParams

FlexLoad parameters.

Attributes:

Name	Type	Description
`curtail_cap_mw`	`float`	Maximum curtailment capacity [MW].
`shift_cap_mw`	`float`	Maximum shift capacity per step [MW].
`shift_horizon`	`int`	Number of future steps over which deferred demand is released.
`curtail_cost_per_mwh`	`float`	Discomfort/compensation cost for curtailment [$/MWh].
`shift_cost_per_mwh`	`float`	Holding cost rate for buffered deferred demand [$/MWh].
`complementarity_penalty`	`float`	Penalty for co-activating curtailment and shift [$/MW].
`price_ref`	`float`	Reference LMP for price normalization in observation [$/MWh].

All scalar fields are pytree_node=False — static under JIT. Ring buffer capacity is fixed at DEFAULT_MAX_BUFFER (compile-time constant required by lax.scan). Do not change at runtime.

LMP is NOT stored in params. Pass it as lmp to step() each call — either from a market clearing result or from an exogenous price profile.

`powerzoojax.envs.resource.flexload.FlexLoadBundle` ¶

Struct-of-arrays FlexLoad bundle for grid env attachment.

Groups N FlexLoad devices with identical action/obs dimensions. Implements the ResourceBundle duck-typed protocol used by DistGridEnv / DistGridMARLEnv.

Action per device : [curtail_frac, shift_frac] ∈ [0, 1]² Obs per device : [curtail_norm, shift_out_norm, shift_in_norm, buf_fill_ratio, buf_energy_norm] (5 dims)

`reset(key)` ¶

Return zero-initialised bundle state.

`step(state, action, ctx)` ¶

Run one step for all N FlexLoad devices.

Parameters:

Name	Type	Description	Default
`state`	`FlexLoadBundleState`	Current bundle state.	required
`action`	`Array`	(action_dim,) in [0, 1], device-major layout `[curtail_0, shift_0, curtail_1, shift_1, ...]`.	required
`ctx`	`Dict[str, Any]`	Grid-provided context (unused; reserved for future use).	required

Returns:

Type	Description
`Tuple[FlexLoadBundleState, Array, Array, Array, Dict[str, Array]]`	(new_state, p_inject_mw, q_inject_mvar, obs_slice, cost_info)

`observe(state, ctx)` ¶

Return flat obs slice of shape (obs_dim,) = (n_devices * 5,).

Per-device obs layout

[curtail_norm, shift_out_norm, shift_in_norm, buf_fill_ratio, buf_energy_norm]

`powerzoojax.envs.resource.flexload.FlexLoadBundleState` ¶

State for a FlexLoadBundle — one ring buffer per device.

Attributes:

Name	Type	Description
`curtailed_mw`	`Array`	Current-step curtailment [MW], shape (n_devices,).
`shift_out_mw`	`Array`	Current-step shift-out [MW], shape (n_devices,).
`shift_in_mw`	`Array`	Current-step released deferred demand [MW], shape (n_devices,).
`deferred_buffer`	`Array`	Per-device ring buffers, shape (n_devices, DEFAULT_MAX_BUFFER).
`buffer_head`	`Array`	Read pointer per device, shape (n_devices,) int32.
`buffer_size`	`Array`	Valid entries per device, shape (n_devices,) int32.

`powerzoojax.envs.resource.flexload.make_flexload_bundle(case, bus_ids, *, curtail_cap_mw=0.1, shift_cap_mw=0.1, shift_horizon=4, curtail_cost_per_mwh=50.0, shift_cost_per_mwh=10.0, dt_hours=0.5)` ¶

Create a FlexLoadBundle by resolving bus_ids through case.node_ids.

All index lookup is done in NumPy at init time (CPU), then converted to jnp arrays. NOT called inside JIT.

Parameters:

Name	Type	Description	Default
`case`		CaseData instance with `node_ids` attribute.	required
`bus_ids`		External node IDs (Sequence[int]).	required
`curtail_cap_mw`		Max curtailment per device [MW]. Scalar or Sequence.	`0.1`
`shift_cap_mw`		Max shift-out per device [MW]. Scalar or Sequence.	`0.1`
`shift_horizon`		Deferred release horizon in steps. Scalar or Sequence.	`4`
`curtail_cost_per_mwh`		Curtailment discomfort cost [$/MWh]. Scalar or Seq.	`50.0`
`shift_cost_per_mwh`		Shift holding cost [$/MWh]. Scalar or Seq.	`10.0`
`dt_hours`	`float`	Time-step duration [h].	`0.5`

Returns:

Type	Description
`FlexLoadBundle`	FlexLoadBundle with (n_devices,) SoA arrays.

`DieselBundle`¶

from powerzoojax.envs.resource.diesel import (
    DieselParams,
    DieselBundle,
    DieselBundleState,
    make_diesel_bundle,
    compute_dg_power,
    compute_dg_fuel_cost,
    compute_dg_emissions,
)

Diesel is currently exposed as pure helpers plus a grid-attachable bundle, not as a standalone Environment. Action per device is a scalar in [0, 1] mapped to active power in [0, p_max]. cost_info["cost"] stays zero; fuel and carbon are reported separately.

`powerzoojax.envs.resource.diesel.DieselParams` ¶

Static parameters for a single diesel generator.

Attributes:

Name	Type	Description
`p_dg_max_mw`	`float`	Nameplate capacity [MW].
`fuel_cost_per_mwh`	`float`	Variable fuel cost [$/MWh].
`emission_factor`	`float`	Carbon intensity [kgCO2 / kWh_e].

All fields are pytree_node=False — compile-time constants under JIT.

`powerzoojax.envs.resource.diesel.compute_dg_power(desired_norm, p_dg_max_mw)` ¶

Compute DG active power from normalised setpoint in [0, 1].

`powerzoojax.envs.resource.diesel.compute_dg_fuel_cost(p_dg_mw, dt_h, fuel_cost_per_mwh)` ¶

Fuel cost [$] for one time step.

`powerzoojax.envs.resource.diesel.compute_dg_emissions(p_dg_mw, dt_h, emission_factor)` ¶

CO2 emissions [kgCO2] for one time step.

MW × h = MWh; × 1e3 → kWh_e; × emission_factor[kg/kWh] → kgCO2.

`powerzoojax.envs.resource.diesel.DieselBundle` ¶

Struct-of-arrays diesel-generator bundle for grid env attachment.

Groups N DGs with identical action/obs dimensions. All quantities are in MW; the host env is responsible for unit conversion. Action per device is a single scalar in [0, 1] mapped to [0, p_max].

Attributes:

Name	Type	Description
`n_devices`	`int`	Number of devices.
`per_device_action_dim`	`int`	1.
`per_device_obs_dim`	`int`	2 (p_norm, dg_margin_norm).
`dt_hours`	`float`	Time-step duration [h].
`bus_idx`	`Array`	Internal bus index per device, shape (n_devices,) int32. For busless mode (microgrid), defaults to `arange(n_devices)`.
`p_max`	`Array`	Rated power per device [MW], shape (n_devices,).
`fuel_cost_per_mwh`	`Array`	Fuel cost [$/MWh], shape (n_devices,).
`emission_factor`	`Array`	kgCO2 / kWh_e, shape (n_devices,).
`p_min_norm`	`Array`	Minimum loading fraction in `[0, 1)` per device, shape (n_devices,). Default 0.0 reproduces the unconstrained (linear-from-0) behaviour. When > 0 a soft deadband at `p_min_norm / 2` is applied: requested setpoints below the deadband shut the generator OFF (p = 0); setpoints above the deadband are clamped to `[p_min_norm, 1]`. Real diesel gensets should not run below ~30% of rated load (wet stacking damages the engine), so a typical opt-in value is `p_min_norm = 0.3`.

`reset(key)` ¶

Initial bundle state — DG starts idle (p=0).

`step(state, action, ctx)` ¶

Run one step for all N DGs.

Parameters:

Name	Type	Description	Default
`state`	`DieselBundleState`	Current bundle state.	required
`action`	`Array`	(n_devices,) in [0, 1] — DG normalised setpoints.	required
`ctx`	`Dict[str, Any]`	Grid-provided context (unused, reserved for future use).	required

Returns:

Type	Description
`Tuple[DieselBundleState, Array, Array, Array, Dict[str, Array]]`	(new_state, p_inject_mw, q_inject_mvar=0, obs_slice, cost_info)

`observe(state, ctx)` ¶

Rebuild obs_slice from state (used by host env's _get_obs / reset).

`powerzoojax.envs.resource.diesel.DieselBundleState` ¶

State for a DieselBundle — last-step DG output.

Attributes:

Name	Type	Description
`p_dg_mw`	`Array`	Last-step active output [MW], shape `(n_devices,)`.

`powerzoojax.envs.resource.diesel.make_diesel_bundle(case=None, bus_ids=None, *, n_devices=None, p_max_mw=0.6, fuel_cost_per_mwh=300.0, emission_factor=0.8, p_min_norm=0.0, dt_hours=0.5)` ¶

Create a DieselBundle.

Two modes:

Grid mode (case and bus_ids provided): bus_idx is resolved through case.node_ids exactly like make_battery_bundle. n_devices = len(bus_ids).

Busless mode (case=None, bus_ids=None): Used by behind-the-meter envs (e.g. DataCenterMicrogridEnv). n_devices must be supplied explicitly (default 1). bus_idx = arange(n_devices) and the host env sums injections rather than scattering by index.

Parameters:

Name	Type	Description	Default
`case`	`Optional[Any]`	Grid case data (optional, for grid mode).	`None`
`bus_ids`	`Optional[Sequence[int]]`	External node IDs (optional, for grid mode).	`None`
`n_devices`	`Optional[int]`	Number of DGs (busless mode; ignored if `bus_ids` given).	`None`
`p_max_mw`	`float`	Nameplate capacity [MW]. Scalar or per-device sequence.	`0.6`
`fuel_cost_per_mwh`	`float`	Variable fuel cost [$/MWh].	`300.0`
`emission_factor`	`float`	Carbon intensity [kgCO2 / kWh_e].	`0.8`
`p_min_norm`	`float`	Minimum loading fraction in `[0, 1)`. Default 0.0 keeps the legacy unconstrained linear behaviour. Set to e.g. 0.3 to model real diesel gensets that should not run below ~30 % of rated load (wet stacking).	`0.0`
`dt_hours`	`float`	Time-step duration [h].	`0.5`

Returns:

Type	Description
`DieselBundle`	DieselBundle ready to be attached to grid/microgrid params.

`DataCenterEnv`¶

from powerzoojax.envs import DataCenterEnv, make_datacenter_params

Three coupled layers: IT power, cooling power, and zone thermal dynamics. Poisson task arrivals go into a fixed-size masked buffer; an EDF-style greedy scheduler allocates GPUs. Active power is always load-like (current_p_mw <= 0).

`powerzoojax.envs.resource.datacenter.DataCenterEnv` ¶

Bases: Environment

AI Data Center resource environment (pure-JAX).

3-D in [0, 1]:

action[0] r_train — GPU fraction for training action[1] r_finetune — GPU fraction for finetuning (of remaining capacity) action[2] cool_norm — cooling setpoint fraction

Observation: 11-D (see module docstring). Info keys: cost (SLA density), cost_overtemp, n_expired, p_dc_mw.

`powerzoojax.envs.resource.datacenter.DataCenterState` ¶

Bases: ResourceState

DataCenter state extending ResourceState.

Attributes:

Name	Type	Description
`t_zone`	`Array`	Zone temperature [°C].
`t_setpoint`	`Array`	Current cooling setpoint [°C].
`t_outdoor`	`Array`	Outdoor temperature [°C].
`p_it_mw`	`Array`	IT (compute) power consumption [MW].
`p_cool_mw`	`Array`	Cooling power consumption [MW].
`gpus_infer`	`Array`	GPUs occupied by inference (exogenous diurnal load).
`gpus_active`	`Array`	Total active GPUs (inference + scheduled tasks).
`task_gpus`	`Array`	GPU count for each task slot [_MAX_TASKS].
`task_deadline`	`Array`	Deadline (absolute step) for each slot [_MAX_TASKS].
`task_duration`	`Array`	Total required duration (steps) for each slot [_MAX_TASKS].
`task_remaining`	`Array`	Steps remaining (only meaningful when running) [_MAX_TASKS].
`task_type`	`Array`	Task type — 0=training, 1=finetuning [_MAX_TASKS].
`task_eta`	`Array`	GPU utilisation efficiency for each slot [_MAX_TASKS].
`task_status`	`Array`	Slot status — 0=inactive, 1=waiting, 2=running [_MAX_TASKS].
`sla_violations`	`Array`	Cumulative count of expired tasks (SLA violations).
`done`	`Array`	Episode termination flag.

`powerzoojax.envs.resource.datacenter.DataCenterParams` ¶

Bases: ResourceParams

DataCenter parameters.

Attributes:

Name	Type	Description
`n_gpus`	`int`	Total GPU count in the data center.
`gpu_idle_w`	`float`	Per-GPU idle power [W] (H100 measured ~55 W).
`gpu_active_w`	`float`	Per-GPU active power [W] (H100 training ~580 W).
`p_base_mw`	`float`	Non-GPU baseline IT power [MW] (networking, storage).
`infer_gpu_peak`	`int`	Peak GPU count consumed by inference (diurnal profile).
`cop_ref`	`float`	Reference COP at t_ref outdoor temperature.
`cop_decay`	`float`	COP fractional decay per °C above t_ref.
`t_ref`	`float`	Reference outdoor temperature for COP [°C].
`c_thermal`	`float`	Thermal capacitance [kWh/°C].
`ua_cooling`	`float`	Cooling heat-transfer coefficient [kW/°C].
`h_wall`	`float`	Building envelope heat-transfer coefficient [kW/°C].
`t_set_min`	`float`	Minimum cooling setpoint [°C].
`t_set_max`	`float`	Maximum cooling setpoint [°C] (ASHRAE upper bound).
`t_initial`	`float`	Initial zone temperature at reset [°C].
`t_critical`	`float`	Over-temperature safety threshold [°C].
`p_aux_frac`	`float`	Auxiliary power as fraction of IT power.
`train_arrival_interval`	`int`	Mean steps between training task arrivals (Poisson rate = 1/interval).
`train_gpu_{lo,hi}`	`int`	Uniform GPU count range for training tasks.
`train_dur_{lo,hi}`	`int`	Uniform duration range [steps] for training tasks.
`train_deadline_slack`	`float`	deadline = arrive + duration × slack (training).
`train_gpu_eta`	`float`	GPU utilisation efficiency for training.
`ft_*`	`float`	Same parameters for finetuning tasks.

All fields are pytree_node=False — static compile-time constants under JIT.

powerzoojax.envs.resource.datacenter.make_datacenter_params(n_gpus=1000, gpu_idle_w=55.0, gpu_active_w=580.0, p_base_mw=0.5, infer_gpu_peak=400, cop_ref=5.0, cop_decay=0.04, t_ref=20.0, c_thermal=500.0, ua_cooling=200.0, h_wall=5.0, t_set_min=18.0, t_set_max=27.0, t_initial=22.0, t_critical=35.0, p_aux_frac=0.05, train_arrival_interval=8, train_gpu_lo=50, train_gpu_hi=200, train_dur_lo=10, train_dur_hi=50, train_deadline_slack=2.0, train_gpu_eta=0.9, ft_arrival_interval=4, ft_gpu_lo=10, ft_gpu_hi=50, ft_dur_lo=5, ft_dur_hi=20, ft_deadline_slack=3.0, ft_gpu_eta=0.75, delta_t_hours=0.5, steps_per_day=48, max_steps=48) ¶

Build DataCenterParams with physical defaults matching PowerZoo DataCenterEnv.

Base resource protocol¶

`powerzoojax.envs.resource.base.ResourceState` ¶

Base state for all resource environments.

Attributes:

Name	Type	Description
`current_p_mw`	`Array`	Active power injection [MW]. Positive = inject to grid, negative = draw.
`current_q_mvar`	`Array`	Reactive power [MVAr]. Usually 0 for DER resources.
`time_step`	`Array`	Current simulation step (int32).

`powerzoojax.envs.resource.base.ResourceParams` ¶

Base parameters for all resource environments.

Attributes:

Name	Type	Description
`bus_id`	`int`	Connection bus ID (int32). -1 = unregistered sentinel.
`delta_t_hours`	`float`	Time step duration in hours.
`steps_per_day`	`int`	Number of steps per day (for time encoding).
`max_steps`	`int`	Maximum steps per episode.

All fields are marked pytree_node=False — they are static compile-time constants under jax.jit and will not be traced.

`powerzoojax.envs.resource.base.ResourceBundle` ¶

Protocol documentation for resource bundles (duck-typed, no ABC).

A bundle groups N homogeneous devices (batteries, PVs, ...) into a single struct-of-arrays container. It is a @struct.dataclass with:

Static fields (pytree_node=False, compile-time constants)::

n_devices            : int    — number of devices
per_device_action_dim: int    — action dim per device (battery=1)
per_device_obs_dim   : int    — obs dim per device   (battery=2)
dt_hours             : float  — time-step duration [h]

Traced leaves (jnp.ndarray, JAX-traced)::

bus_idx    : int32  (n_devices,)   — grid node index per device
<device-specific SoA arrays>       — e.g. power_max, capacity, ...

Derived properties::

action_dim = n_devices * per_device_action_dim
obs_dim    = n_devices * per_device_obs_dim

Methods::

reset(key) -> BundleState
step(state, action, ctx) -> (new_state, p_inject_mw, q_inject_mvar,
                              obs_slice, cost_info)
observe(state, ctx) -> obs_slice

obs_slice uses device-major layout::

obs_slice[i * per_device_obs_dim : (i+1) * per_device_obs_dim]

gives device i's observation. Single-agent sees the flat vector; MARL adapters reshape via per_device_*_dim — zero bundle changes.

`powerzoojax.envs.resource.base.ResourceBundleState` ¶

Marker base for bundle state pytrees (documentation only).

Each concrete bundle (BatteryBundle, PVBundle, ...) defines its own state @struct.dataclass. This empty base exists solely for type-hint documentation — it carries no fields and no runtime cost.

Resources¶

Quick orientation¶

Cost semantics by resource¶

BatteryEnv and BatteryBundle¶

powerzoojax.envs.resource.battery.BatteryEnv ¶

reset(key, params) ¶

step(key, state, action, params) ¶

observation_space(params) ¶

action_space(params) ¶

powerzoojax.envs.resource.battery.BatteryState ¶

powerzoojax.envs.resource.battery.BatteryParams ¶

powerzoojax.envs.resource.battery.make_battery_params(bus_id=-1, capacity_mwh=50.0, power_mw=20.0, eta_charge=None, eta_discharge=None, eta_roundtrip=None, soc_min=0.1, soc_max=0.9, initial_soc=0.5, delta_t_hours=0.5, steps_per_day=48, max_steps=48, cycle_cost_per_mwh=0.0) ¶

powerzoojax.envs.resource.battery.compute_feasible_power(soc, desired_power, params) ¶

powerzoojax.envs.resource.battery.update_soc(soc, power, params) ¶

powerzoojax.envs.resource.battery.BatteryBundle ¶

reset(key) ¶

step(state, action, ctx) ¶

observe(state, ctx) ¶

powerzoojax.envs.resource.battery.BatteryBundleState ¶

RenewableEnv, SolarEnv, WindEnv, and RenewableBundle¶

powerzoojax.envs.resource.renewable.RenewableEnv ¶

powerzoojax.envs.resource.renewable.RenewableState ¶

powerzoojax.envs.resource.renewable.RenewableParams ¶

powerzoojax.envs.resource.renewable.SolarEnv ¶

powerzoojax.envs.resource.renewable.WindEnv ¶

powerzoojax.envs.resource.renewable.RenewableBundle ¶

reset(key) ¶

step(state, action, ctx) ¶

observe(state, ctx) ¶

powerzoojax.envs.resource.renewable.RenewableBundleState ¶

powerzoojax.envs.resource.renewable.make_renewable_bundle(case=None, bus_ids=None, *, n_devices=None, capacity_mw=0.2, s_rated_mva=None, profiles=None, max_steps=48, allow_curtailment=True, enable_curtailment=True, enable_q_control=True, curtail_cost_per_mwh=0.0, dt_hours=0.5) ¶

VehicleEnv¶

powerzoojax.envs.resource.vehicle.VehicleEnv ¶

powerzoojax.envs.resource.vehicle.VehicleState ¶

powerzoojax.envs.resource.vehicle.VehicleParams ¶

powerzoojax.envs.resource.vehicle.make_vehicle_params(E_max_kWh=60.0, soc_init=0.8, soc_min=0.1, soc_max=0.95, soc_departure_min=0.8, p_charge_max_kW=7.0, p_discharge_max_kW=7.0, eta_charge=0.95, eta_discharge=0.95, commute_schedule=None, delta_t_minutes=15.0, steps_per_day=None, max_steps=None) ¶

FlexLoadEnv and FlexLoadBundle¶

powerzoojax.envs.resource.flexload.FlexLoadEnv ¶

step(key, state, action, params, lmp=0.0) ¶

powerzoojax.envs.resource.flexload.FlexLoadState ¶

powerzoojax.envs.resource.flexload.FlexLoadParams ¶

powerzoojax.envs.resource.flexload.FlexLoadBundle ¶

reset(key) ¶

step(state, action, ctx) ¶

observe(state, ctx) ¶

powerzoojax.envs.resource.flexload.FlexLoadBundleState ¶

powerzoojax.envs.resource.flexload.make_flexload_bundle(case, bus_ids, *, curtail_cap_mw=0.1, shift_cap_mw=0.1, shift_horizon=4, curtail_cost_per_mwh=50.0, shift_cost_per_mwh=10.0, dt_hours=0.5) ¶

DieselBundle¶

powerzoojax.envs.resource.diesel.DieselParams ¶

powerzoojax.envs.resource.diesel.compute_dg_power(desired_norm, p_dg_max_mw) ¶

powerzoojax.envs.resource.diesel.compute_dg_fuel_cost(p_dg_mw, dt_h, fuel_cost_per_mwh) ¶

powerzoojax.envs.resource.diesel.compute_dg_emissions(p_dg_mw, dt_h, emission_factor) ¶

powerzoojax.envs.resource.diesel.DieselBundle ¶

reset(key) ¶

step(state, action, ctx) ¶

observe(state, ctx) ¶

powerzoojax.envs.resource.diesel.DieselBundleState ¶

powerzoojax.envs.resource.diesel.make_diesel_bundle(case=None, bus_ids=None, *, n_devices=None, p_max_mw=0.6, fuel_cost_per_mwh=300.0, emission_factor=0.8, p_min_norm=0.0, dt_hours=0.5) ¶

DataCenterEnv¶

powerzoojax.envs.resource.datacenter.DataCenterEnv ¶

powerzoojax.envs.resource.datacenter.DataCenterState ¶

powerzoojax.envs.resource.datacenter.DataCenterParams ¶

Base resource protocol¶

powerzoojax.envs.resource.base.ResourceState ¶

powerzoojax.envs.resource.base.ResourceParams ¶

powerzoojax.envs.resource.base.ResourceBundle ¶

powerzoojax.envs.resource.base.ResourceBundleState ¶

`BatteryEnv` and `BatteryBundle`¶

`powerzoojax.envs.resource.battery.BatteryEnv` ¶

`reset(key, params)` ¶

`step(key, state, action, params)` ¶

`observation_space(params)` ¶

`action_space(params)` ¶

`powerzoojax.envs.resource.battery.BatteryState` ¶

`powerzoojax.envs.resource.battery.BatteryParams` ¶

`powerzoojax.envs.resource.battery.make_battery_params(bus_id=-1, capacity_mwh=50.0, power_mw=20.0, eta_charge=None, eta_discharge=None, eta_roundtrip=None, soc_min=0.1, soc_max=0.9, initial_soc=0.5, delta_t_hours=0.5, steps_per_day=48, max_steps=48, cycle_cost_per_mwh=0.0)` ¶

`powerzoojax.envs.resource.battery.compute_feasible_power(soc, desired_power, params)` ¶

`powerzoojax.envs.resource.battery.update_soc(soc, power, params)` ¶

`powerzoojax.envs.resource.battery.BatteryBundle` ¶

`reset(key)` ¶

`step(state, action, ctx)` ¶

`observe(state, ctx)` ¶

`powerzoojax.envs.resource.battery.BatteryBundleState` ¶

`RenewableEnv`, `SolarEnv`, `WindEnv`, and `RenewableBundle`¶

`powerzoojax.envs.resource.renewable.RenewableEnv` ¶

`powerzoojax.envs.resource.renewable.RenewableState` ¶

`powerzoojax.envs.resource.renewable.RenewableParams` ¶

`powerzoojax.envs.resource.renewable.SolarEnv` ¶

`powerzoojax.envs.resource.renewable.WindEnv` ¶

`powerzoojax.envs.resource.renewable.RenewableBundle` ¶

`reset(key)` ¶

`step(state, action, ctx)` ¶

`observe(state, ctx)` ¶

`powerzoojax.envs.resource.renewable.RenewableBundleState` ¶

`powerzoojax.envs.resource.renewable.make_renewable_bundle(case=None, bus_ids=None, *, n_devices=None, capacity_mw=0.2, s_rated_mva=None, profiles=None, max_steps=48, allow_curtailment=True, enable_curtailment=True, enable_q_control=True, curtail_cost_per_mwh=0.0, dt_hours=0.5)` ¶

`VehicleEnv`¶

`powerzoojax.envs.resource.vehicle.VehicleEnv` ¶

`powerzoojax.envs.resource.vehicle.VehicleState` ¶

`powerzoojax.envs.resource.vehicle.VehicleParams` ¶

`FlexLoadEnv` and `FlexLoadBundle`¶

`powerzoojax.envs.resource.flexload.FlexLoadEnv` ¶

`step(key, state, action, params, lmp=0.0)` ¶

`powerzoojax.envs.resource.flexload.FlexLoadState` ¶

`powerzoojax.envs.resource.flexload.FlexLoadParams` ¶

`powerzoojax.envs.resource.flexload.FlexLoadBundle` ¶

`reset(key)` ¶

`step(state, action, ctx)` ¶

`observe(state, ctx)` ¶

`powerzoojax.envs.resource.flexload.FlexLoadBundleState` ¶

`powerzoojax.envs.resource.flexload.make_flexload_bundle(case, bus_ids, *, curtail_cap_mw=0.1, shift_cap_mw=0.1, shift_horizon=4, curtail_cost_per_mwh=50.0, shift_cost_per_mwh=10.0, dt_hours=0.5)` ¶

`DieselBundle`¶

`powerzoojax.envs.resource.diesel.DieselParams` ¶

`powerzoojax.envs.resource.diesel.compute_dg_power(desired_norm, p_dg_max_mw)` ¶

`powerzoojax.envs.resource.diesel.compute_dg_fuel_cost(p_dg_mw, dt_h, fuel_cost_per_mwh)` ¶

`powerzoojax.envs.resource.diesel.compute_dg_emissions(p_dg_mw, dt_h, emission_factor)` ¶

`powerzoojax.envs.resource.diesel.DieselBundle` ¶

`reset(key)` ¶

`step(state, action, ctx)` ¶

`observe(state, ctx)` ¶

`powerzoojax.envs.resource.diesel.DieselBundleState` ¶

`powerzoojax.envs.resource.diesel.make_diesel_bundle(case=None, bus_ids=None, *, n_devices=None, p_max_mw=0.6, fuel_cost_per_mwh=300.0, emission_factor=0.8, p_min_norm=0.0, dt_hours=0.5)` ¶

`DataCenterEnv`¶

`powerzoojax.envs.resource.datacenter.DataCenterEnv` ¶

`powerzoojax.envs.resource.datacenter.DataCenterState` ¶

`powerzoojax.envs.resource.datacenter.DataCenterParams` ¶

`powerzoojax.envs.resource.base.ResourceState` ¶

`powerzoojax.envs.resource.base.ResourceParams` ¶

`powerzoojax.envs.resource.base.ResourceBundle` ¶

`powerzoojax.envs.resource.base.ResourceBundleState` ¶