Skip to content

v0.2.0 - init/reset, space ordering, new wrappers

Choose a tag to compare

View all tags
@keraJLi keraJLi released this 06 Feb 00:22
· 78 commits to main since this release
v0.2.0
3888356

Release description

Summary

This release introduces a clearer environment lifecycle (init vs reset), richer auto-reset and info semantics, a strict batch → tree → continuous / discrete space composition model, and several new wrappers. The public API has breaking changes; see below.

Breaking changes

Environment lifecycle: init and reset

  • reset(key, state=None, **kwargs) has been replaced by two methods:
    • init(key) — Initialize the environment from scratch. Use for the first episode or when no prior state exists. Returns (State, Info).
    • reset(key, state) — Reset the inner environment while preserving wrapper-level state (e.g. running stats). Takes the current state; default implementation is return self.init(key).
  • step(state, action, **kwargs) — The **kwargs argument has been removed. Use step(state, action) only.

AutoResetWrapper: info.final instead of next_obs

  • When a step ends an episode (terminated or truncated), the wrapper now returns:
    • info.obs — The initial observation of the new episode (after reset).
    • info.final — A snapshot of the terminal step’s info (before reset), for correct value bootstrapping and episode statistics.
  • The previous “next observation” / next_obs-style behavior is replaced by info.final.obs for the true next state when done.
  • AutoResetWrapper.reset is not supported and raises NotImplementedError.

Spaces: batch_space removed, composition order enforced

  • batch_space(space, batch_size) has been removed. Use BatchedSpace(space=..., batch_size=...) directly.
  • Space composition is now explicit:
    • PyTreeSpace may only have Discrete or Continuous as leaves (the “base” layer).
    • BatchedSpace may wrap any Space (base, PyTreeSpace, or another BatchedSpace).
    • Ordering: batch → tree → base.
  • BatchedSpace.shape for a PyTreeSpace now prepends batch dimension(s) to each leaf shape (previously the tree shape was returned without a batch dimension).

EpisodeStatisticsWrapper

  • Episode stats are no longer a dict in info. They are exposed as an EpisodeStatistics struct with reward and length (arrays), and are attached to both state and info as stats.

Info and compat

  • InfoContainer.reward may be float | Array (e.g. batched rewards).
  • Compat adapters (Brax, Jumanji, Gymnax, etc.) now implement init(key) instead of reset(key); terminated is normalized to a scalar bool where needed (e.g. Jumanji, Brax).

New features

New wrappers

  • ClipActionWrapper — Clips actions to the environment’s action space (BatchedSpace, PyTreeSpace, Continuous, Discrete).
  • FlattenActionWrapper — Flattens actions to a single vector; action_space becomes one Continuous or Discrete over the flattened leaves.
  • FlattenObservationWrapper — Flattens observations to a single vector; observation_space is one Continuous or Discrete.
  • ContinuousObservationWrapper — Converts observations to float and maps Discrete observation spaces to Continuous (e.g. for learning on integer obs).
  • PooledInitVmapWrapper — Batched env with a pool of initial states; on done, each env is reset to a randomly chosen pre-initialized state from the pool, with info.final semantics aligned with AutoResetWrapper.

Spaces

  • peel_batched(space) — Returns (batch_dims_tuple, base_space) by stripping BatchedSpace layers.
  • PyTreeSpace.dtype — Property returning a PyTree of dtypes for each leaf.
  • BatchedSpace.dtype — Now defined via the base space after peeling batch dimensions.

Struct and container

  • Container — __str__ now includes extra fields (extras) in the string representation.
  • FrozenPyTreeNode and Container — @dataclass_transform() applied for better type-checker support.

AutoResetWrapper

  • Docstring and behavior now describe info.final for terminal step data and value bootstrapping (use final.obs for V(s′) on truncation, and final for last-episode stats).

ObservationNormalizationWrapper

  • observation_space is now derived from the inner space by mapping each leaf to Continuous(-inf, inf).
  • reset preserves running statistics (rmv_state) across resets.

Vmap wrappers

  • VmapWrapper and VmapEnvsWrapper support both init and reset; reset passes state into the inner vmap. Key handling supports both a single key (split into batch) and a batch of keys.

Config and tests

  • pyproject.toml: testpaths = ["tests"]; gymnax and kinetix-env sourced from git where needed; tensorflow-probability override for compatibility.
  • create() — When default_max_steps is set, the created env is wrapped with TruncationWrapper (tested).

Fixes and small changes

  • README: “stoax” → “stoa” (typo).
  • .gitignore: CLAUDE.md added; comment updated.
  • Compat suites: env_params and info placeholders adjusted (Gymnax, Craftax, Kinetix) so info has a consistent structure from the first step.
  • TruncationWrapper: reset correctly passes state.inner_state to the inner env and resets steps to 0.
  • StateInjectionWrapper: init always calls env.init; reset uses injected state when reset_state / reset_obs are set, otherwise env.reset(key, state.inner_state).

Full Changelog: v0.1.1...v0.2.0

Assets 2
Loading