Wrapper States

The state containers used by Envrax's stateful wrappers.

envrax.wrappers.frame_stack.FrameStackState

Bases: EnvState, Generic[InnerStateT]

State for FrameStackObservation.

Generic over the inner env's state type so env_state is precisely typed when the wrapper is parameterised. The base rng/step/done fields are forwarded copies from the inner state so that the framework (e.g. VecEnv auto-reset) can read them from the outer state.

Parameters:

Name	Type	Description	Default
env_state	InnerStateT	Underlying environment state (may itself be a wrapped state).	required
obs_stack	Array	uint8[H, W, n_stack] — Ring buffer of the last n_stack processed observations, oldest frame at channel index 0.	required

Source code in envrax/wrappers/frame_stack.py

@chex.dataclass
class FrameStackState(EnvState, Generic[InnerStateT]):
    """
    State for `FrameStackObservation`.

    Generic over the inner env's state type so `env_state` is precisely
    typed when the wrapper is parameterised. The base `rng`/`step`/`done`
    fields are forwarded copies from the inner state so that the framework
    (e.g. `VecEnv` auto-reset) can read them from the outer state.

    Parameters
    ----------
    env_state : InnerStateT
        Underlying environment state (may itself be a wrapped state).
    obs_stack : jax.Array
        uint8[H, W, n_stack] — Ring buffer of the last `n_stack` processed
        observations, oldest frame at channel index 0.
    """

    env_state: InnerStateT
    obs_stack: jax.Array

envrax.wrappers.record_episode_statistics.EpisodeStatisticsState

Bases: EnvState, Generic[InnerStateT]

State for RecordEpisodeStatistics.

Generic over the inner env's state type so env_state is precisely typed when the wrapper is parameterised. The base rng/step/done fields are forwarded copies from the inner state.

Parameters:

Name	Type	Description	Default
env_state	InnerStateT	Inner environment state.	required
episode_return	Array	Cumulative reward for the current episode. float32 scalar.	required
episode_length	Array	Number of steps taken in the current episode. int32 scalar.	required

Source code in envrax/wrappers/record_episode_statistics.py

@chex.dataclass
class EpisodeStatisticsState(EnvState, Generic[InnerStateT]):
    """
    State for `RecordEpisodeStatistics`.

    Generic over the inner env's state type so `env_state` is precisely
    typed when the wrapper is parameterised. The base `rng`/`step`/`done`
    fields are forwarded copies from the inner state.

    Parameters
    ----------
    env_state : InnerStateT
        Inner environment state.
    episode_return : chex.Array
        Cumulative reward for the current episode. float32 scalar.
    episode_length : chex.Array
        Number of steps taken in the current episode. int32 scalar.
    """

    env_state: InnerStateT
    episode_return: chex.Array
    episode_length: chex.Array