v0.4.2 - Change kinetix adapter to sample premade levels from size

The Kinetix adapter now resets to a random level from a size category ("s", "m", "l") instead of loading individual levels. On each reset, a level is uniformly sampled from all packaged levels in the chosen category (10 for s, 24 for m, 40 for l). The create_premade classmethod has been replaced by create_from_size(size), and from_name now accepts "s", "m", "l", or "random". Individual level IDs like "s/h4_thrust_aim" are no longer supported. The "random" mode for fully procedural UED levels remains unchanged. Usage: KinetixEnvelope.create_from_size("m") or create("kinetix::s").

v0.4.1 - new-style keys, better contracts

v0.4.1 - new-style keys, better contracts

Improvements

• New-style JAX PRNG keys - BatchedSpace and VmapEnvsWrapper now expect jax.random.key() instead of the old-style jax.random.PRNGKey().
• Jumanji obs structure preserved - Spec conversion now uses Spec._constructor to faithfully reconstruct observation pytrees (namedtuples, etc.) instead of flattening them into dicts.
• PyTreeSpace can now be empty - Removed the restriction requiring at least one leaf, which was too strict for some use cases.
• Navix obs space workaround - Clamped observation space n to a minimum of 10 to work around navix#109.

Testing

• Shared contract helpers - Moved contract.py to tests/ and added a unified assert_obs_matches_space checker. All adapter tests now use the same obs validation instead of ad-hoc callbacks.
• Wrapper contract smoke tests - Every wrapper now gets a contract smoke test (reset/step shapes, info fields, obs containment) via a single parametrized test.
• TestInfo uses Container - Switched from FrozenPyTreeNode for adherence to Info protocol.

Full Changelog: v0.4.0...v0.4.1

v0.4.0 - Rename compat -> adapters

Release v0.4.0

Breaking changes

• envelope.compat → envelope.adapters
  We renamed the submodule compat to adapters as this is more fitting.

---

Full Changelog: v0.3.0...v0.4.0

v0.3.0 reset(key, state) -> reset(state, key)

Release v0.3.0

Breaking changes

• reset(key, state) → reset(state, key)
  In-place reset (preserving episode-persistent wrapper state) now takes state first, then key:
  state, info = env.reset(state, key). This is the same argument order as step.

---

Bug fixes

• Brax compat: Fixed typo jnp.asarry → jnp.asarray.
• Brax compat: Fixed bug where .item() was called on a traced value of done.
• Jumanji compat: Same fix.

Full Changelog: v0.2.0...v0.3.0

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

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

v0.1.1 - updated envelope namespace

All of envelope's main features are now available under the envelope namespace, allowing you to import them easily:

from envelope import create, Environment, Container, Wrapper, VmapWrapper, ...

Initial release!

💌 jax-envelope v0.1.0

The first release of envelope — a JAX-native environment interface for reinforcement learning.

---

✨ Highlights

Simple, expressive interaction

• Environments are pytrees — trace parameters, vmap, jit, all just work
• Idiomatic interface:
  • reset(key) → (State, Info)
  • step(state, action) → (State, Info)

Powerful, composable wrappers

• VmapWrapper / VmapEnvsWrapper for batched environments
• AutoResetWrapper with state injection for UED
• ObservationNormalizationWrapper with running statistics across episodes
• TruncationWrapper and explicit truncation handling for proper value bootstrapping

Adapters for 190+ environments

• gymnax, brax, jumanji, kinetix, craftax, mujoco_playground

---

📦 Installation

pip install jax-envelope

import envelope

env = envelope.create("brax::ant")
state, info = env.reset(key)
state, info = env.step(state, action)