Environment

Basic interface: init, reset and step

Many RL libraries have step and reset. Envelope introduces an additional init method, splitting the traditional reset into two:

• init(key) initializes environment state from scratch.
• reset(state, key) resets the current episode while preserving persistent state. For example, ObservationNormalizationWrapper keeps its running statistics across resets, while only resetting the inner environment.
• step(state, action) steps the environment.

By default, reset simply calls init, which is correct for most environments where the starting state is sampled from a fixed distribution.

Environment methods return State and Info tuples. Info is a structural protocol; InfoContainer is its default implementation, built on Container. Wrappers extend the info with extra fields via update() — for example, EpisodeStatisticsWrapper adds a stats field, which consumers can then access as info.stats.

Pytree Structure Contract

JAX transformations such as jax.lax.scan, jax.vmap, and jax.jit require that the pytree structure (treedef) of inputs and outputs remains consistent across iterations. This imposes a critical requirement on envelope environments:

init, reset and step should return State and Info objects with identical pytree structures and leaf shapes.

This guarantees that you can map jnp.where on the Info they produce, or emit the Info as the output of jax.lax.scan.

API Reference

envelope.environment.Environment

Bases: ABC, FrozenPyTreeNode

Base class for all environments.

State is an opaque PyTree owned by each environment; wrappers that stack environments should expose their wrapped env state as inner_state while adding any wrapper-specific fields.

Attributes

action_space abstractmethod cached property

The space of actions.

observation_space abstractmethod cached property

The space of observations.

Functions

init(key) abstractmethod

Initialize environment state and sample the initial info.

This method closely resembles the reset method of gymnasium or gymnax. However, it is normally called only once per environment lifecycle, as subsequent resets should be performed using the Environment.reset method, which may preserve episode-persistent state.

reset(state, key)

Reset the environment while preserving episode-persistent state.

By default, state is ignored and init is called, as this most closely matches the standard reinforcement learning setting in which the starting state is sampled from a fixed distribution.

step(state, action) abstractmethod

Step the environment given an action, returning the next state and info.

envelope.typing.Info

Bases: Protocol

Info is a runtime-checkable Protocol that defines required fields and methods for environment emissions, including observation, reward, and termination/truncation flags.

Attributes:

Name	Type	Description
obs	PyTree	The observation from the environment.
reward	float	The reward from the environment.
terminated	bool	Whether the episode has terminated.
truncated	bool	Whether the episode has truncated.

Functions

update(**changes)

Update the info container with new values. This method should return a new instance with updated and potentially new values.

envelope.environment.InfoContainer dataclass

Bases: Container

Info container for environment emissions, including observation, reward, and termination/truncation flags. This container implements the Info protocol.

Attributes:

Name	Type	Description
obs	PyTree	The observation from the environment.
reward	float | Array	The reward from the environment.
terminated	bool	Whether the episode has terminated.
truncated	bool	Whether the episode has truncated.