Spaces¶
PowerZooJax uses lightweight JAX-friendly space descriptors for observations and actions.
Space types¶
powerzoojax.envs.spaces.Space
¶
powerzoojax.envs.spaces.MultiDiscrete
¶
Helper constructors¶
powerzoojax.envs.spaces.make_box(low, high, shape=None, dtype=jnp.float32)
¶
Create a Box space, broadcasting scalar low/high to shape if needed.
NOTE: This is a setup-time helper — do NOT call inside JIT-compiled code
(it uses jnp.broadcast_to and Python shape inference).
powerzoojax.envs.spaces.make_discrete(n, dtype=jnp.int32)
¶
Create a Discrete space with validation.
powerzoojax.envs.spaces.make_multi_discrete(nvec, dtype=jnp.int32)
¶
Create a MultiDiscrete space with validation.
powerzoojax.envs.spaces.make_multi_binary(n, dtype=jnp.int32)
¶
Create a MultiBinary space with validation.
Environment contract¶
All environments implement the same high-level interface:
powerzoojax.envs.base.Environment
¶
Bases: ABC
Abstract base class for JAX-based environments.
Follows Gymnax-style API: pure functions, explicit state, JIT-compatible.
reset(key, params)
abstractmethod
¶
Reset environment to initial state.
step(key, state, action, params)
abstractmethod
¶
Execute one environment step -> (obs, state, reward, costs, done, info).
info values must be JAX arrays or Python scalars so that the dict remains a valid pytree under vmap / pmap batching.
costs is the explicit CMDP constraint vector for this env. Its
ordering must match self.constraint_names(params) and remain
static under JIT / vmap. Zero-constraint envs should return an empty
float32 vector with shape (0,).
PRNG contract for auto-reset: subclass implementations that embed
auto-reset logic (jnp.where(done, reset_state, next_state)) MUST
call jax.random.split(key) to obtain a separate key for
self.reset(). Reusing the same key for both transition noise
and reset would correlate the new episode's initial state with the
last step's randomness, breaking MDP independence.
step_auto_reset(key, state, action, params)
¶
Step with automatic reset on done -- enables lax.scan rollouts.
All env subclasses already embed auto-reset inside step().
This wrapper applies stop_gradient to the returned observation
and state so gradients do not flow across episode boundaries
(correct for model-free RL).
Note: the underlying step() must split key before using it
for reset (see step() docstring). Do NOT pass the same key
to both transition randomness and reset().
observation_space(params)
abstractmethod
¶
Return observation space specification.
action_space(params)
abstractmethod
¶
Return action space specification.
powerzoojax.envs.base.EnvState
¶
Base environment state (immutable PyTree).
All subclasses should use @struct.dataclass decorator.
powerzoojax.envs.base.EnvParams
¶
Base environment parameters (immutable PyTree).
Static configuration that doesn't change during an episode.
max_steps and delta_t_hours are marked as non-leaf so they
are part of the pytree structure (static under JIT).