jax-envelope 0.1.0__py3-none-any.whl → 0.2.0__py3-none-any.whl

envelope/__init__.py

@@ -0,0 +1,54 @@
+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 Container, FrozenPyTreeNode, field, static_field
+from envelope.wrappers import (
+    AutoResetWrapper,
+    ClipActionWrapper,
+    ContinuousObservationWrapper,
+    EpisodeStatisticsWrapper,
+    FlattenActionWrapper,
+    FlattenObservationWrapper,
+    ObservationNormalizationWrapper,
+    PooledInitVmapWrapper,
+    StateInjectionWrapper,
+    TruncationWrapper,
+    VmapEnvsWrapper,
+    VmapWrapper,
+    WrappedState,
+    Wrapper,
+)
+
+__all__ = [
+    # Basic functionality
+    "create",
+    "Environment",
+    "Info",
+    "InfoContainer",
+    # Spaces
+    "Space",
+    "BatchedSpace",
+    "Continuous",
+    "Discrete",
+    "PyTreeSpace",
+    # Struct
+    "field",
+    "static_field",
+    "FrozenPyTreeNode",
+    "Container",
+    # Wrappers
+    "Wrapper",
+    "WrappedState",
+    "AutoResetWrapper",
+    "ClipActionWrapper",
+    "ContinuousObservationWrapper",
+    "EpisodeStatisticsWrapper",
+    "FlattenActionWrapper",
+    "FlattenObservationWrapper",
+    "ObservationNormalizationWrapper",
+    "PooledInitVmapWrapper",
+    "StateInjectionWrapper",
+    "TruncationWrapper",
+    "VmapWrapper",
+    "VmapEnvsWrapper",
+]

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

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

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)

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
 

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)

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))

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)
 

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

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})"

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] = {}

envelope/typing.py

@@ -1,4 +1,3 @@
-from enum import Enum
 from typing import Any, TypeAlias
 
 import jax
@@ -6,18 +5,3 @@ import jax
 PyTree: TypeAlias = Any
 Key: TypeAlias = jax.Array
 Array: TypeAlias = jax.Array
-
-
-class BatchKind(Enum):
-    """
-    Batch semantics for environments.
-    - VMAP: environment represents instances compatible with `jax.vmap`, for example by
-      wrapping it in a `VmapWrapper`.
-    - NATIVE_POOL: environment represents a batch of instances via a native pool. This
-      is the case when it is wrapping a non-jax-based environment that supports native
-      batching, for example those provided by envpool. Environments in this mode cannot
-      be vmapped, as it would break the native batching semantics.
-    """
-
-    VMAP = "vmap"
-    NATIVE_POOL = "native_pool"

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",
+]

envelope/wrappers/autoreset_wrapper.py

@@ -1,4 +1,7 @@
+from typing import override
+
 import jax
+import jax.numpy as jnp
 
 from envelope.environment import Info
 from envelope.struct import field
@@ -7,30 +10,71 @@ 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()
 
-    def reset(
-        self, key: Key, state: PyTree | None = None, **kwargs
-    ) -> tuple[WrappedState, Info]:
+    @override
+    def init(self, key: Key) -> tuple[WrappedState, Info]:
         key, subkey = jax.random.split(key)
-        inner_state = state.inner_state if state else None
-        inner_state, info = self.env.reset(key, inner_state, **kwargs)
-        state = self.AutoResetState(inner_state=inner_state, reset_key=subkey)
-        return state, info.update(next_obs=info.obs)
-
-    def step(
-        self, state: WrappedState, action: PyTree, **kwargs
-    ) -> tuple[WrappedState, Info]:
-        inner_state, info_step = self.env.step(state.inner_state, action, **kwargs)
-        done = info_step.terminated | info_step.truncated
-
-        state = self.AutoResetState(inner_state=inner_state, reset_key=state.reset_key)
-        info = info_step.update(next_obs=info_step.obs)
-
-        state, info = jax.lax.cond(
-            done,
-            lambda: self.reset(state.reset_key, state),
-            lambda: (state, info),
+        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

envelope/wrappers/clip_action_wrapper.py

@@ -0,0 +1,27 @@
+from typing import override
+
+import jax
+import jax.numpy as jnp
+
+from envelope.environment import Info, State
+from envelope.spaces import BatchedSpace, Continuous, Discrete, PyTreeSpace, Space
+from envelope.typing import PyTree
+from envelope.wrappers.wrapper import Wrapper
+
+
+def clip_action(action: PyTree, space: Space) -> PyTree:
+    if isinstance(space, BatchedSpace):
+        return jax.vmap(clip_action, in_axes=(0, None))(action, space.space)
+    elif isinstance(space, PyTreeSpace):
+        return jax.tree.map(clip_action, action, space.tree)
+    elif isinstance(space, Continuous):
+        return jnp.clip(action, space.low, space.high)
+    elif isinstance(space, Discrete):
+        return jnp.clip(action, 0, space.n - 1)
+
+
+class ClipActionWrapper(Wrapper):
+    @override
+    def step(self, state: State, action: PyTree) -> tuple[State, Info]:
+        action = clip_action(action, self.action_space)
+        return self.env.step(state, action)