Wrappers

Wrapper Synopsis

Wrappers compose around environments via nesting: Wrapper2(env=Wrapper1(env=base_env)). Each wrapper may transform observations, actions, or spaces, and may add its own fields to the state or info.

Wrappers that need to track data across steps (e.g. a step counter, running statistics) extend WrappedState, which nests the inner environment's state as inner_state. Wrappers that only transform observations or actions (like ClipActionWrapper or ContinuousObservationWrapper) pass state through without wrapping. The unwrapped property traverses the full nesting to return the base environment's state.

Wrappers communicate additional data to user code by adding fields to the info via info.update(...). For example, EpisodeStatisticsWrapper adds stats, AutoResetWrapper adds final (a snapshot of the terminal step's info, enabling value bootstrapping), and ObservationNormalizationWrapper adds unnormalized_obs.

Vectorization

Three wrappers add batch dimensions:

• VmapWrapper vmaps a single environment with batch_size parallel instances.
• VmapEnvsWrapper vmaps over a batched pytree of environment instances, for example created via jax.vmap(make_env)(params). This is useful when different instances have different configurations.
• PooledInitVmapWrapper vmaps like VmapWrapper, but pre-computes a pool of initial states and samples from them on reset. It also includes built-in autoreset logic, making it an alternative to AutoResetWrapper + VmapWrapper.

Wrapper Ordering

The key constraint is that AutoResetWrapper calls reset() on its inner wrapper chain when an episode ends. Wrappers that need their reset() triggered on episode boundaries (e.g. TruncationWrapper resetting its step counter) must therefore be inside AutoResetWrapper. Vectorization wrappers must be outside, since autoreset operates per-element.

From innermost to outermost:

base env → Observation/action transforms → Episode logic → AutoReset → Vectorization

A concrete example with all layers:

VmapWrapper                              # outermost: adds batch dim
└─ AutoResetWrapper                      # resets on done, adds `final`
   └─ StateInjectionWrapper              # (optional) overrides reset target
      └─ EpisodeStatisticsWrapper        # tracks reward/length
         └─ TruncationWrapper            # caps episode length
            └─ ObservationNormalizationWrapper
               └─ ContinuousObservationWrapper
                  └─ ClipActionWrapper
                     └─ base env         # innermost

Not all wrappers are needed in every pipeline. The ordering between wrappers in the same layer (e.g. the observation/action transforms) is flexible.

API Reference

envelope.wrappers.Wrapper

Bases: Environment

Wrapper for environments.

envelope.wrappers.WrappedState

Bases: FrozenPyTreeNode

envelope.wrappers.AutoResetWrapper

Bases: Wrapper

Wrapper that automatically resets the environment when an episode ends.

When a step results in termination or truncation, this wrapper immediately resets the environment. The returned info preserves critical information from the terminal step while providing the new episode's initial observation.

Info fields after a terminal step (terminated=True or truncated=True): obs: Initial observation from the new episode (after reset). final: Full info snapshot from the terminal step (before reset). terminated: True if the episode ended due to termination. truncated: True if the episode ended due to truncation. reward: Reward from the terminal step.

Info fields during normal steps (terminated=False and truncated=False): obs: Current observation. final: Info snapshot from the last completed episode (persisted). terminated: False. truncated: False. reward: Reward from the step.

This design enables correct value bootstrapping

• Use final.obs for value estimation of the true next state
• On termination: V(s_final) = 0 (episode truly ended)
• On truncation: bootstrap from V(final.obs) (episode cut off artificially)
• final persists until the next episode completes, giving easy access to last episode's aggregated stats (e.g., final.episode_return)

envelope.wrappers.ClipActionWrapper

Bases: Wrapper

envelope.wrappers.ContinuousObservationWrapper

Bases: Wrapper

envelope.wrappers.EpisodeStatisticsWrapper

Bases: Wrapper

envelope.wrappers.FlattenActionWrapper

Bases: Wrapper

envelope.wrappers.FlattenObservationWrapper

Bases: Wrapper

envelope.wrappers.ObservationNormalizationWrapper

Bases: Wrapper

Attributes

stats_spec = static_field(default=None) class-attribute instance-attribute

Per-leaf normalization statistics spec as a pytree of jax.ShapeDtypeStruct. Shapes must be broadcastable to the observation leaves. If None, inferred from the observation_space with BatchedSpace ignored; each leaf must have a floating dtype.

envelope.wrappers.PooledInitVmapWrapper

Bases: Wrapper

envelope.wrappers.StateInjectionWrapper

Bases: Wrapper

Stores a state that all resets return to.

For UED: use set_reset_state() to update the injected state, then all resets (including auto-reset) return to that state until it's changed again.

Usage

env = AutoResetWrapper(StateInjectionWrapper(env=base_env))
state, info = env.init(key)

for outer_iter in range(num_outer_iters):
    # Sample a new task and set it as the reset state
    task_state, task_obs = sample_task(key)
    state = env.set_reset_state(state, task_state, task_obs)

    # Run episode - auto-resets return to task_state
    for inner_step in range(num_inner_steps):
        state, info = env.step(state, policy(info.obs))

Functions

set_reset_state(state, reset_state, reset_obs)

Update the state that resets will return to.

This method traverses the wrapped state hierarchy to find and update the InjectedState, then reconstructs the full state tree.

Parameters:

Name	Type	Description	Default
state	WrappedState	Current state (can be from any outer wrapper)	required
reset_state	PyTree	The state to reset to (inner environment state)	required
reset_obs	PyTree	The observation to return on reset	required

Returns:

Type	Description
WrappedState	New state with updated reset fields at the appropriate level

envelope.wrappers.TruncationWrapper

Bases: Wrapper

envelope.wrappers.VmapWrapper

Bases: Wrapper

Does not wrap the state.

envelope.wrappers.VmapEnvsWrapper

Bases: Wrapper

Vectorizes over a batched collection of environment instances (vmapping over 'self').

Usage

envs = jax.vmap(make_env)(params_batch)  # env pytree batched on leading axis
wrapped = VmapEnvsWrapper(env=envs, batch_size=B)
state, info = wrapped.init(keys)  # keys shape (B, 2) or single key
next_state, info = wrapped.step(state, action)