jax-envelope 0.1.1__tar.gz → 0.2.0__tar.gz

{jax_envelope-0.1.1 → jax_envelope-0.2.0}/.gitignore

@@ -166,7 +166,8 @@ wandb/
 # ruff
 .ruff_cache/
 
-# Cursor
+# Vibecoding
 .cursor
 .cursorignore
 AGENTS.md
+CLAUDE.md

{jax_envelope-0.1.1 → jax_envelope-0.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: jax-envelope
-Version: 0.1.1
+Version: 0.2.0
 Summary: A JAX-native environment interface with powerful wrappers and adapters for popular RL environment suites
 Project-URL: Homepage, https://github.com/keraJLi/envelope
 Project-URL: Repository, https://github.com/keraJLi/envelope
@@ -82,6 +82,6 @@ pip install jax-envelope
 ```
 
 ## 💞 Related projects
-* [stoax](https://github.com/EdanToledo/Stoa) is a very similar project that provides adapters and wrappers for the jumanji-like interface.
+* [stoa](https://github.com/EdanToledo/Stoa) is a very similar project that provides adapters and wrappers for the jumanji-like interface.
 * Check out all the great suites we have adapters for! [gymnax](https://github.com/RobertTLange/gymnax), [brax](https://github.com/google/brax), [jumanji](https://github.com/instadeepai/jumanji), [kinetix](https://github.com/flairox/kinetix), [craftax](https://github.com/MichaelTMatthews/craftax), [mujoco_playground](https://github.com/google-deepmind/mujoco_playground).
 * We will be adding support for [jaxmarl](https://github.com/flairox/jaxmarl) and [pgx](https://github.com/sotetsuk/pgx) in the future, as soon as we figured out the best ever MARL interface for JAX!

{jax_envelope-0.1.1 → jax_envelope-0.2.0}/README.md

@@ -56,6 +56,6 @@ pip install jax-envelope
 ```
 
 ## 💞 Related projects
-* [stoax](https://github.com/EdanToledo/Stoa) is a very similar project that provides adapters and wrappers for the jumanji-like interface.
+* [stoa](https://github.com/EdanToledo/Stoa) is a very similar project that provides adapters and wrappers for the jumanji-like interface.
 * Check out all the great suites we have adapters for! [gymnax](https://github.com/RobertTLange/gymnax), [brax](https://github.com/google/brax), [jumanji](https://github.com/instadeepai/jumanji), [kinetix](https://github.com/flairox/kinetix), [craftax](https://github.com/MichaelTMatthews/craftax), [mujoco_playground](https://github.com/google-deepmind/mujoco_playground).
 * We will be adding support for [jaxmarl](https://github.com/flairox/jaxmarl) and [pgx](https://github.com/sotetsuk/pgx) in the future, as soon as we figured out the best ever MARL interface for JAX!

{jax_envelope-0.1.1 → jax_envelope-0.2.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "jax-envelope"
-version = "0.1.1"
+version = "0.2.0"
 description = "A JAX-native environment interface with powerful wrappers and adapters for popular RL environment suites"
 readme = "README.md"
 requires-python = ">=3.12"
@@ -51,18 +51,15 @@ allow-direct-references = true
 [tool.hatch.build.targets.wheel]
 packages = ["src/envelope"]
 
-[tool.uv.sources]
-gymnax = { git = "https://github.com/RobertTLange/gymnax" }
-
 [dependency-groups]
 compat = [
-    "gymnax @ git+https://github.com/RobertTLange/gymnax@main",
     "brax>=0.13.0",
     "craftax>=1.4.3",
     "navix>=0.7.0",
     "jumanji>=1.0.1",
-    "kinetix-env>=2.0.0",
     "playground>=0.1.0",
+    "gymnax",
+    "kinetix-env",
 ]
 dev = [
     "hypothesis>=6.148.1",
@@ -71,7 +68,17 @@ dev = [
     "ruff>=0.14.2",
 ]
 
+[tool.uv]
+override-dependencies = [
+    "tensorflow-probability>=0.26.0.dev20260116",
+]
+
+[tool.uv.sources]
+gymnax = { git = "https://github.com/RobertTLange/gymnax.git" }
+kinetix-env = { git = "https://github.com/FLAIROx/Kinetix.git" }
+
 [tool.pytest.ini_options]
+testpaths = ["tests"]
 markers = [
   "compat: tests requiring optional compat dependencies",
 ]

{jax_envelope-0.1.1 → jax_envelope-0.2.0}/src/envelope/__init__.py

@@ -1,16 +1,22 @@
 from envelope.compat import create
 from envelope.environment import Environment, Info, InfoContainer
 from envelope.spaces import BatchedSpace, Continuous, Discrete, PyTreeSpace, Space
-from envelope.struct import field, static_field, FrozenPyTreeNode, Container
+from envelope.struct import Container, FrozenPyTreeNode, field, static_field
 from envelope.wrappers import (
-    Wrapper,
-    WrappedState,
     AutoResetWrapper,
+    ClipActionWrapper,
+    ContinuousObservationWrapper,
+    EpisodeStatisticsWrapper,
+    FlattenActionWrapper,
+    FlattenObservationWrapper,
     ObservationNormalizationWrapper,
+    PooledInitVmapWrapper,
     StateInjectionWrapper,
     TruncationWrapper,
-    VmapWrapper,
     VmapEnvsWrapper,
+    VmapWrapper,
+    WrappedState,
+    Wrapper,
 )
 
 __all__ = [
@@ -34,7 +40,13 @@ __all__ = [
     "Wrapper",
     "WrappedState",
     "AutoResetWrapper",
+    "ClipActionWrapper",
+    "ContinuousObservationWrapper",
+    "EpisodeStatisticsWrapper",
+    "FlattenActionWrapper",
+    "FlattenObservationWrapper",
     "ObservationNormalizationWrapper",
+    "PooledInitVmapWrapper",
     "StateInjectionWrapper",
     "TruncationWrapper",
     "VmapWrapper",

{jax_envelope-0.1.1 → jax_envelope-0.2.0}/src/envelope/compat/brax_envelope.py

@@ -48,7 +48,7 @@ class BraxEnvelope(Environment):
     def default_max_steps(self) -> int:
         return _BRAX_DEFAULT_EPISODE_LENGTH
 
-    def __post_init__(self) -> "BraxEnvelope":
+    def __post_init__(self):
         if isinstance(self.brax_env, BraxWrapper):
             warnings.warn(
                 "Environment wrapping should be handled by envelope. "
@@ -57,7 +57,7 @@ class BraxEnvelope(Environment):
             object.__setattr__(self, "brax_env", self.brax_env.unwrapped)
 
     @override
-    def reset(self, key: Key) -> tuple[State, Info]:
+    def init(self, key: Key) -> tuple[State, Info]:
         brax_state = self.brax_env.reset(key)
         info = InfoContainer(obs=brax_state.obs, reward=0.0, terminated=False)
         info = info.update(**dataclasses.asdict(brax_state))
@@ -67,7 +67,9 @@ class BraxEnvelope(Environment):
     def step(self, state: State, action: PyTree) -> tuple[State, Info]:
         brax_state = self.brax_env.step(state, action)
         info = InfoContainer(
-            obs=brax_state.obs, reward=brax_state.reward, terminated=brax_state.done
+            obs=brax_state.obs,
+            reward=brax_state.reward,
+            terminated=jnp.asarry(brax_state.done, dtype=bool).item(),
         )
         info = info.update(**dataclasses.asdict(brax_state))
         return brax_state, info

{jax_envelope-0.1.1 → jax_envelope-0.2.0}/src/envelope/compat/craftax_envelope.py

@@ -22,7 +22,7 @@ class CraftaxEnvelope(Environment):
     """Wrapper to convert a Craftax environment to a envelope environment."""
 
     craftax_env: Any = static_field()
-    env_params: PyTree
+    env_params: PyTree = static_field()  # TODO: remove static marker as soon as craftax merges https://github.com/MichaelTMatthews/Craftax/pull/48
 
     @classmethod
     def from_name(
@@ -54,12 +54,27 @@ class CraftaxEnvelope(Environment):
     def default_max_steps(self) -> int:
         return int(self.craftax_env.default_params.max_timesteps)
 
+    @cached_property
+    def _craftax_info_placeholder(self) -> PyTree:
+        key = jax.random.PRNGKey(0)
+        _, state = self.craftax_env.reset(key, self.env_params)
+        _, _, _, _, info = self.craftax_env.step(
+            key,
+            state,
+            self.craftax_env.action_space(self.env_params).sample(key),
+            self.env_params,
+        )
+        return jax.tree.map(lambda x: jnp.full_like(x, jnp.nan), info)
+
     @override
-    def reset(self, key: Key) -> tuple[State, Info]:
+    def init(self, key: Key) -> tuple[State, Info]:
+        # TODO: this function does not add env_info (or comparable) to the info
+        # container. We should add tests for this (and all other envelopes) and fix it.
         key, subkey = jax.random.split(key)
         obs, env_state = self.craftax_env.reset(subkey, self.env_params)
         state = Container().update(key=key, env_state=env_state)
         info = InfoContainer(obs=obs, reward=0.0, terminated=False)
+        info = info.update(info=self._craftax_info_placeholder)
         return state, info
 
     @override

{jax_envelope-0.1.1 → jax_envelope-0.2.0}/src/envelope/compat/gymnax_envelope.py

@@ -1,5 +1,5 @@
 from functools import cached_property
-from typing import Any, override
+from typing import Any, Callable, cast, override
 
 import jax
 import jax.numpy as jnp
@@ -10,15 +10,24 @@ from gymnax.environments.environment import EnvParams as GymnaxEnvParams
 
 from envelope import spaces as envelope_spaces
 from envelope.environment import Environment, Info, InfoContainer, State
-from envelope.struct import Container, static_field
+from envelope.struct import Container, field, static_field
 from envelope.typing import Key, PyTree
 
+_GymnaxReset = Callable[
+    [Key, GymnaxEnvParams],
+    tuple[PyTree, Any],
+]
+_GymnaxStep = Callable[
+    [Key, Any, PyTree, GymnaxEnvParams],
+    tuple[PyTree, Any, jnp.ndarray, jnp.ndarray, PyTree],
+]
+
 
 class GymnaxEnvelope(Environment):
     """Wrapper to convert a Gymnax environment to a envelope environment."""
 
     gymnax_env: GymnaxEnv = static_field()
-    env_params: PyTree
+    env_params: PyTree = field()
 
     @classmethod
     def from_name(
@@ -43,19 +52,37 @@ class GymnaxEnvelope(Environment):
     def default_max_steps(self) -> int:
         return int(self.gymnax_env.default_params.max_steps_in_episode)
 
+    @cached_property
+    def _gymnax_info_placeholder(self) -> PyTree:
+        reset_fn = cast(_GymnaxReset, self.gymnax_env.reset)
+        step_fn = cast(_GymnaxStep, self.gymnax_env.step)
+
+        key = jax.random.PRNGKey(0)
+        _, state = reset_fn(key, self.env_params)
+        _, _, _, _, info = step_fn(
+            key,
+            state,
+            self.gymnax_env.action_space(self.env_params).sample(key),
+            self.env_params,
+        )
+        return jax.tree.map(lambda x: jnp.full_like(x, jnp.nan, dtype=float), info)
+
     @override
-    def reset(self, key: Key) -> tuple[State, Info]:
+    def init(self, key: Key) -> tuple[State, Info]:
+        reset_fn = cast(_GymnaxReset, self.gymnax_env.reset)
+
         key, subkey = jax.random.split(key)
-        obs, env_state = self.gymnax_env.reset(subkey, self.env_params)
+        obs, env_state = reset_fn(subkey, self.env_params)
         state = Container().update(key=key, env_state=env_state)
         info = InfoContainer(obs=obs, reward=0.0, terminated=False)
-        info = info.update(info=None)
+        info = info.update(info=self._gymnax_info_placeholder)
         return state, info
 
     @override
     def step(self, state: State, action: PyTree) -> tuple[State, Info]:
         key, subkey = jax.random.split(state.key)
-        obs, env_state, reward, done, env_info = self.gymnax_env.step(
+        step_fn = cast(_GymnaxStep, self.gymnax_env.step)
+        obs, env_state, reward, done, env_info = step_fn(
             subkey, state.env_state, action, self.env_params
         )
         state = state.update(key=key, env_state=env_state)

{jax_envelope-0.1.1 → jax_envelope-0.2.0}/src/envelope/compat/jumanji_envelope.py

@@ -48,7 +48,7 @@ class JumanjiEnvelope(Environment):
         return self._default_time_limit
 
     @override
-    def reset(self, key: Key) -> tuple[State, Info]:
+    def init(self, key: Key) -> tuple[State, Info]:
         env_state, timestep = self.jumanji_env.reset(key)
         info = convert_jumanji_to_envelope_info(timestep)
         return env_state, info
@@ -81,8 +81,9 @@ class JumanjiEnvelope(Environment):
 
 
 def convert_jumanji_to_envelope_info(timestep: JumanjiTimeStep) -> InfoContainer:
+    term = jnp.asarray(timestep.last(), dtype=bool).item()
     info = InfoContainer(
-        obs=timestep.observation, reward=timestep.reward, terminated=timestep.last()
+        obs=timestep.observation, reward=timestep.reward, terminated=term
     ).update(**timestep.extras)
     return info
 

{jax_envelope-0.1.1 → jax_envelope-0.2.0}/src/envelope/compat/kinetix_envelope.py

@@ -28,6 +28,7 @@ from kinetix.environment import (
 from kinetix.environment.ued.ued import make_reset_fn_sample_kinetix_level
 from kinetix.util.saving import load_from_json_file
 
+from envelope import field
 from envelope import spaces as envelope_spaces
 from envelope.compat.gymnax_envelope import _convert_space as _convert_gymnax_space
 from envelope.environment import Environment, Info, InfoContainer, State
@@ -67,7 +68,7 @@ class KinetixEnvelope(Environment):
     """Wrapper to convert a Kinetix environment to a envelope environment."""
 
     kinetix_env: Any = static_field()
-    env_params: Any
+    env_params: Any = field()
 
     @property
     def default_max_steps(self) -> int:
@@ -162,7 +163,7 @@ class KinetixEnvelope(Environment):
         return cls(kinetix_env=kinetix_env, env_params=env_params)
 
     @override
-    def reset(self, key: Key) -> tuple[State, Info]:
+    def init(self, key: Key) -> tuple[State, Info]:
         key, subkey = jax.random.split(key)
         obs, env_state = self.kinetix_env.reset(subkey, self.env_params)
         state_out = Container().update(key=key, env_state=env_state)

{jax_envelope-0.1.1 → jax_envelope-0.2.0}/src/envelope/compat/mujoco_playground_envelope.py

@@ -56,7 +56,7 @@ class MujocoPlaygroundEnvelope(Environment):
         return self._default_max_steps
 
     @override
-    def reset(self, key: Key) -> tuple[State, Info]:
+    def init(self, key: Key) -> tuple[State, Info]:
         env_state = self.mujoco_playground_env.reset(key)
         info = InfoContainer(obs=env_state.obs, reward=0.0, terminated=False)
         info = info.update(**dataclasses.asdict(env_state))

{jax_envelope-0.1.1 → jax_envelope-0.2.0}/src/envelope/compat/navix_envelope.py

@@ -38,7 +38,7 @@ class NavixEnvelope(Environment):
         return _NAVIX_DEFAULT_MAX_STEPS
 
     @override
-    def reset(self, key: Key) -> tuple[State, Info]:
+    def init(self, key: Key) -> tuple[State, Info]:
         timestep = self.navix_env.reset(key)
         return timestep, convert_navix_to_envelope_info(timestep)
 

{jax_envelope-0.1.1 → jax_envelope-0.2.0}/src/envelope/environment.py

@@ -5,7 +5,7 @@ from typing import Protocol, runtime_checkable
 
 from envelope import spaces
 from envelope.struct import Container, FrozenPyTreeNode
-from envelope.typing import Key, PyTree
+from envelope.typing import Array, Key, PyTree
 
 __all__ = ["Environment", "State", "Info", "InfoContainer"]
 
@@ -23,7 +23,7 @@ class Info(Protocol):
 
 class InfoContainer(Container):
     obs: PyTree
-    reward: float
+    reward: float | Array
     terminated: bool
     truncated: bool = field(default=False)
 
@@ -38,18 +38,25 @@ class Environment(ABC, FrozenPyTreeNode):
 
     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. `reset` may optionally receive a prior
-    state (for cross-episode persistence) and arbitrary **kwargs that wrappers
-    or environments can use.
+    adding any wrapper-specific fields.
+
+    Two distinct lifecycle methods:
+        init(key) - Initialize environment and all state from scratch.
+        reset(key, state) - Reset the inner environment while preserving
+            episode-persistent state.
     """
 
     @abstractmethod
-    def reset(
-        self, key: Key, state: State | None = None, **kwargs
-    ) -> tuple[State, Info]: ...
+    def init(self, key: Key) -> tuple[State, Info]:
+        """Initialize environment and all state from scratch."""
+        ...
+
+    def reset(self, key: Key, state: State) -> tuple[State, Info]:
+        """Reset the inner environment while preserving episode-persistent state."""
+        return self.init(key)
 
     @abstractmethod
-    def step(self, state: State, action: PyTree, **kwargs) -> tuple[State, Info]: ...
+    def step(self, state: State, action: PyTree) -> tuple[State, Info]: ...
 
     @abstractmethod
     @cached_property

{jax_envelope-0.1.1 → jax_envelope-0.2.0}/src/envelope/spaces.py

@@ -1,6 +1,6 @@
 from abc import ABC, abstractmethod
 from functools import cached_property
-from typing import override
+from typing import cast, override
 
 import jax
 from jax import numpy as jnp
@@ -65,7 +65,9 @@ class Continuous(Space):
     high: float | jax.Array
 
     @classmethod
-    def from_shape(cls, low: float, high: float, shape: tuple[int]) -> "Continuous":
+    def from_shape(
+        cls, low: float, high: float, shape: tuple[int, ...]
+    ) -> "Continuous":
         return cls(
             low=jnp.full(shape, low, dtype=jnp.asarray(low).dtype),
             high=jnp.full(shape, high, dtype=jnp.asarray(high).dtype),
@@ -106,17 +108,25 @@ class PyTreeSpace(Space):
     """A Space defined by a PyTree structure of other Spaces.
 
     Args:
-        tree: A PyTree with Space objects leaves.
+        tree: A PyTree with Discrete or Continuous leaves.
 
     Usage:
         space = PyTreeSpace({
-            "action": Discrete(n=4, dtype=jnp.int32),
-            "obs": Continuous(low=0.0, high=1.0, shape=(2,), dtype=jnp.float32)
+            "action": Discrete(n=4),
+            "obs": Continuous(low=0.0, high=1.0, shape=(2,))
         })
     """
 
     tree: PyTree
 
+    def __post_init__(self):
+        leaves = jax.tree.leaves(self.tree, is_leaf=lambda x: isinstance(x, Space))
+        for leaf in leaves:
+            if not isinstance(leaf, (Discrete, Continuous)):
+                raise TypeError(
+                    f"PyTreeSpace leaves must be Discrete or Continuous, got {type(leaf).__name__}"
+                )
+
     @override
     def sample(self, key: Key) -> PyTree:
         leaves, treedef = jax.tree.flatten(
@@ -149,16 +159,23 @@ class PyTreeSpace(Space):
             is_leaf=lambda node: isinstance(node, Space),
         )
 
-
-def batch_space(space: Space, batch_size: int) -> Space:
-    if isinstance(space, PyTreeSpace):
-        batched_tree = jax.tree.map(
-            lambda sp: batch_space(sp, batch_size),
-            space.tree,
+    @property
+    def dtype(self) -> PyTree:
+        return jax.tree.map(
+            lambda space: space.dtype,
+            self.tree,
             is_leaf=lambda node: isinstance(node, Space),
         )
-        return PyTreeSpace(batched_tree)
-    return BatchedSpace(space=space, batch_size=batch_size)
+
+
+def peel_batched(space: Space) -> tuple[tuple[int, ...], Space]:
+    """Collect batch dimensions and return (batch_dims_tuple, base_space)."""
+    dims: list[int] = []
+    s: Space = space
+    while isinstance(s, BatchedSpace):
+        dims.append(s.batch_size)
+        s = s.space
+    return tuple(dims), s
 
 
 class BatchedSpace(Space):
@@ -190,16 +207,19 @@ class BatchedSpace(Space):
 
     @cached_property
     def shape(self) -> PyTree:
-        inner_shape = self.space.shape
-        # For tuple shapes (leaf spaces), prepend batch dimension.
-        # PyTree shapes are handled by wrapping leaves with BatchedSpace via batch_space.
-        if isinstance(inner_shape, tuple):
-            return (self.batch_size,) + inner_shape
-        return inner_shape
+        batch_dims, base = peel_batched(self)
+        if isinstance(base, PyTreeSpace):
+            return jax.tree.map(
+                lambda space: batch_dims + space.shape,
+                base.tree,
+                is_leaf=lambda node: isinstance(node, Space),
+            )
+        return batch_dims + base.shape
 
     @property
-    def dtype(self):
-        return getattr(self.space, "dtype", None)
+    def dtype(self) -> PyTree:
+        _, base = peel_batched(self)
+        return base.dtype
 
     def __repr__(self) -> str:
         return f"BatchedSpace(space={self.space!r}, batch_size={self.batch_size})"

{jax_envelope-0.1.1 → jax_envelope-0.2.0}/src/envelope/struct.py

@@ -1,6 +1,6 @@
 import dataclasses
 from dataclasses import KW_ONLY
-from typing import Any, Iterable, Iterator, Mapping, Self, Tuple
+from typing import Any, Iterable, Iterator, Mapping, Self, Tuple, dataclass_transform
 
 import jax
 
@@ -24,6 +24,7 @@ def static_field(**kwargs):
     return field(pytree_node=False, **kwargs)
 
 
+@dataclass_transform()
 class FrozenPyTreeNode:
     """
     Frozen dataclass base that is a JAX pytree node.
@@ -64,6 +65,7 @@ class FrozenPyTreeNode:
         return dataclasses.replace(self, **changes)
 
 
+@dataclass_transform()
 @jax.tree_util.register_pytree_node_class
 @dataclasses.dataclass(frozen=True, eq=True, repr=True, slots=False)
 class Container:
@@ -104,6 +106,13 @@ class Container:
         for k, v in self._extras.items():
             yield (k, v)
 
+    def __str__(self) -> str:
+        core_str = super().__str__()
+        if not self._extras:
+            return core_str
+        extras_str = f", {', '.join(f'{k}={v!r}' for k, v in self._extras.items())}"
+        return f"{core_str[:-1]}{extras_str})"  # remove closing parenthesis from core
+
     def update(self, **changes: PyTree) -> Self:
         core_names = {f.name for f in dataclasses.fields(self) if f.name != "_extras"}
         core_updates: dict[str, PyTree] = {}

jax_envelope-0.2.0/src/envelope/wrappers/__init__.py

@@ -0,0 +1,36 @@
+from envelope.wrappers.autoreset_wrapper import AutoResetWrapper
+from envelope.wrappers.clip_action_wrapper import ClipActionWrapper
+from envelope.wrappers.continuous_observation_wrapper import (
+    ContinuousObservationWrapper,
+)
+from envelope.wrappers.episode_statistics_wrapper import EpisodeStatisticsWrapper
+from envelope.wrappers.flatten_action_wrapper import FlattenActionWrapper
+from envelope.wrappers.flatten_observation_wrapper import FlattenObservationWrapper
+from envelope.wrappers.observation_normalization_wrapper import (
+    ObservationNormalizationWrapper,
+)
+from envelope.wrappers.pooled_init_vmap_wrapper import PooledInitVmapWrapper
+from envelope.wrappers.state_injection_wrapper import StateInjectionWrapper
+from envelope.wrappers.truncation_wrapper import TruncationWrapper
+from envelope.wrappers.vmap_envs_wrapper import VmapEnvsWrapper
+from envelope.wrappers.vmap_wrapper import VmapWrapper
+from envelope.wrappers.wrapper import WrappedState, Wrapper
+
+__all__ = [
+    # Basic functionality
+    "Wrapper",
+    "WrappedState",
+    # Wrappers
+    "AutoResetWrapper",
+    "ClipActionWrapper",
+    "ContinuousObservationWrapper",
+    "EpisodeStatisticsWrapper",
+    "FlattenActionWrapper",
+    "FlattenObservationWrapper",
+    "ObservationNormalizationWrapper",
+    "PooledInitVmapWrapper",
+    "StateInjectionWrapper",
+    "TruncationWrapper",
+    "VmapWrapper",
+    "VmapEnvsWrapper",
+]

jax_envelope-0.2.0/src/envelope/wrappers/autoreset_wrapper.py

@@ -0,0 +1,80 @@
+from typing import override
+
+import jax
+import jax.numpy as jnp
+
+from envelope.environment import Info
+from envelope.struct import field
+from envelope.typing import Key, PyTree
+from envelope.wrappers.wrapper import WrappedState, Wrapper
+
+
+class AutoResetWrapper(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)
+    """
+
+    class AutoResetState(WrappedState):
+        reset_key: jax.Array = field()
+        last_final: Info = field()
+
+    @override
+    def init(self, key: Key) -> tuple[WrappedState, Info]:
+        key, subkey = jax.random.split(key)
+        inner_state, info = self.env.init(key)
+        # Initialize last_final with the reset info (no previous episode yet)
+        last_final = jax.tree.map(lambda x: jnp.full_like(x, jnp.nan), info)
+        state = self.AutoResetState(
+            inner_state=inner_state, reset_key=subkey, last_final=last_final
+        )
+        return state, info.update(final=state.last_final)
+
+    @override
+    def reset(self, key: Key, state: WrappedState) -> tuple[WrappedState, Info]:
+        raise NotImplementedError("Reset is not implemented for AutoResetWrapper")
+
+    @override
+    def step(self, state: WrappedState, action: PyTree) -> tuple[WrappedState, Info]:
+        key, key_reset = jax.random.split(state.reset_key)
+        state = state.replace(reset_key=key)
+
+        inner_state, info = self.env.step(state.inner_state, action)
+        reset_inner_state, reset_info = self.env.reset(key_reset, inner_state)
+
+        # Select next state and info based on done
+        done = info.terminated | info.truncated
+        state = jax.tree.map(
+            lambda reset, next: jax.lax.select(done, reset, next),
+            state.replace(inner_state=reset_inner_state),
+            state.replace(inner_state=inner_state),
+        )
+        info = jax.tree.map(
+            lambda reset, next: jax.lax.select(done, reset, next),
+            reset_info.update(final=info),
+            info.update(final=state.last_final),
+        )
+        return state, info