jax-envelope 0.1.0__py3-none-any.whl

envelope/compat/mujoco_playground_envelope.py

@@ -0,0 +1,101 @@
+import dataclasses
+from functools import cached_property
+from typing import Any, override
+
+from jax import numpy as jnp
+from mujoco_playground import registry
+
+from envelope import spaces as envelope_spaces
+from envelope.environment import Environment, Info, InfoContainer, State
+from envelope.struct import static_field
+from envelope.typing import Key, PyTree
+
+_MAX_INT = int(jnp.iinfo(jnp.int32).max)
+
+
+_MUJOCO_PLAYGROUND_DEFAULT_EPISODE_LENGTH = 1000
+
+
+class MujocoPlaygroundEnvelope(Environment):
+    """Wrapper to convert a mujoco_playground environment to a envelope environment."""
+
+    mujoco_playground_env: Any = static_field()
+    _default_max_steps: int = static_field(
+        default=_MUJOCO_PLAYGROUND_DEFAULT_EPISODE_LENGTH
+    )
+
+    @classmethod
+    def from_name(
+        cls, env_name: str, env_kwargs: dict[str, Any] | None = None
+    ) -> "MujocoPlaygroundEnvelope":
+        """Creates a MujocoPlaygroundEnvelope from a name and keyword arguments.
+        env_kwargs are passed to config_overrides of mujoco_playground.registry.load."""
+        env_kwargs = env_kwargs or {}
+        if "episode_length" in env_kwargs:
+            raise ValueError(
+                "Cannot override 'episode_length' directly. "
+                "Use TruncationWrapper for episode length control."
+            )
+
+        # Get default episode_length from registry config
+        default_config = registry.get_default_config(env_name)
+        default_max_steps = default_config.episode_length
+
+        # Set episode_length to a very large value
+        # (mujoco_playground uses int for episode_length, so we use max int instead of inf)
+        env_kwargs["episode_length"] = _MAX_INT
+
+        # Pass all env_kwargs as config_overrides
+        env = registry.load(
+            env_name, config_overrides=env_kwargs if env_kwargs else None
+        )
+        return cls(mujoco_playground_env=env, _default_max_steps=default_max_steps)
+
+    @property
+    def default_max_steps(self) -> int:
+        return self._default_max_steps
+
+    @override
+    def reset(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))
+        return env_state, info
+
+    @override
+    def step(self, state: State, action: PyTree) -> tuple[State, Info]:
+        state = self.mujoco_playground_env.step(state, action)
+        info = InfoContainer(obs=state.obs, reward=state.reward, terminated=state.done)
+        info = info.update(**dataclasses.asdict(state))
+        return state, info
+
+    @override
+    @cached_property
+    def action_space(self) -> envelope_spaces.Space:
+        # MuJoCo Playground actions are typically bounded [-1, 1]
+        return envelope_spaces.Continuous.from_shape(
+            low=-1.0, high=1.0, shape=(self.mujoco_playground_env.action_size,)
+        )
+
+    @override
+    @cached_property
+    def observation_space(self) -> envelope_spaces.Space:
+        import jax
+
+        def to_space(size):
+            shape = (size,) if isinstance(size, int) else size
+            return envelope_spaces.Continuous.from_shape(
+                low=-jnp.inf, high=jnp.inf, shape=shape
+            )
+
+        def is_leaf(x):
+            return isinstance(x, int) or (
+                isinstance(x, tuple) and all(isinstance(i, int) for i in x)
+            )
+
+        space_tree = jax.tree.map(
+            to_space, self.mujoco_playground_env.observation_size, is_leaf=is_leaf
+        )
+        if isinstance(space_tree, envelope_spaces.Space):
+            return space_tree
+        return envelope_spaces.PyTreeSpace(space_tree)

envelope/compat/navix_envelope.py

@@ -0,0 +1,86 @@
+import dataclasses
+from functools import cached_property
+from typing import Any, override
+
+import jax.numpy as jnp
+import navix
+from navix import spaces as navix_spaces
+from navix.environments.environment import Environment as NavixEnv
+
+from envelope import spaces as envelope_spaces
+from envelope.environment import Environment, Info, InfoContainer, State
+from envelope.typing import Key, PyTree
+
+_NAVIX_DEFAULT_MAX_STEPS = 100
+
+
+class NavixEnvelope(Environment):
+    """Wrapper to convert a Navix environment to a envelope environment."""
+
+    navix_env: NavixEnv
+
+    @classmethod
+    def from_name(
+        cls, env_name: str, env_kwargs: dict[str, Any] | None = None
+    ) -> "NavixEnvelope":
+        env_kwargs = env_kwargs or {}
+        if "max_steps" in env_kwargs:
+            raise ValueError(
+                "Cannot override 'max_steps' directly. "
+                "Use TruncationWrapper for episode length control."
+            )
+        env_kwargs["max_steps"] = jnp.inf
+        navix_env = navix.make(env_name, **env_kwargs)
+        return cls(navix_env=navix_env)
+
+    @property
+    def default_max_steps(self) -> int:
+        return _NAVIX_DEFAULT_MAX_STEPS
+
+    @override
+    def reset(self, key: Key) -> tuple[State, Info]:
+        timestep = self.navix_env.reset(key)
+        return timestep, convert_navix_to_envelope_info(timestep)
+
+    @override
+    def step(self, state: State, action: PyTree) -> tuple[State, Info]:
+        timestep = self.navix_env.step(state, action)
+        return timestep, convert_navix_to_envelope_info(timestep)
+
+    @override
+    @cached_property
+    def action_space(self) -> envelope_spaces.Space:
+        return convert_navix_to_envelope_space(self.navix_env.action_space)
+
+    @override
+    @cached_property
+    def observation_space(self) -> envelope_spaces.Space:
+        return convert_navix_to_envelope_space(self.navix_env.observation_space)
+
+
+def convert_navix_to_envelope_info(nvx_timestep: navix.Timestep) -> InfoContainer:
+    timestep_dict = dataclasses.asdict(nvx_timestep)
+    step_type = timestep_dict.pop("step_type")
+    info = InfoContainer(
+        obs=timestep_dict.pop("observation"),
+        reward=timestep_dict.pop("reward"),
+        terminated=step_type == navix.StepType.TERMINATION,
+        truncated=step_type == navix.StepType.TRUNCATION,
+    )
+    info = info.update(**timestep_dict)
+    return info
+
+
+def convert_navix_to_envelope_space(
+    nvx_space: navix_spaces.Space,
+) -> envelope_spaces.Space:
+    if isinstance(nvx_space, navix_spaces.Discrete):
+        n = jnp.asarray(nvx_space.n).astype(nvx_space.dtype)
+        return envelope_spaces.Discrete.from_shape(n, shape=nvx_space.shape)
+
+    elif isinstance(nvx_space, navix_spaces.Continuous):
+        low = jnp.asarray(nvx_space.minimum).astype(nvx_space.dtype)
+        high = jnp.asarray(nvx_space.maximum).astype(nvx_space.dtype)
+        return envelope_spaces.Continuous.from_shape(low, high, shape=nvx_space.shape)
+
+    raise ValueError(f"Unsupported space type: {type(nvx_space)}")

envelope/environment.py

@@ -0,0 +1,64 @@
+from abc import ABC, abstractmethod
+from dataclasses import field
+from functools import cached_property
+from typing import Protocol, runtime_checkable
+
+from envelope import spaces
+from envelope.struct import Container, FrozenPyTreeNode
+from envelope.typing import Key, PyTree
+
+__all__ = ["Environment", "State", "Info", "InfoContainer"]
+
+
+@runtime_checkable
+class Info(Protocol):
+    obs: PyTree
+    reward: float
+    terminated: bool
+    truncated: bool
+
+    def update(self, **changes: PyTree) -> "Info": ...
+    def __getattr__(self, name: str) -> PyTree: ...
+
+
+class InfoContainer(Container):
+    obs: PyTree
+    reward: float
+    terminated: bool
+    truncated: bool = field(default=False)
+
+
+# State remains a general PyTree alias; environments are not forced to WrappedState
+State = PyTree
+
+
+class Environment(ABC, FrozenPyTreeNode):
+    """
+    Base class for all environments.
+
+    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.
+    """
+
+    @abstractmethod
+    def reset(
+        self, key: Key, state: State | None = None, **kwargs
+    ) -> tuple[State, Info]: ...
+
+    @abstractmethod
+    def step(self, state: State, action: PyTree, **kwargs) -> tuple[State, Info]: ...
+
+    @abstractmethod
+    @cached_property
+    def observation_space(self) -> spaces.Space: ...
+
+    @abstractmethod
+    @cached_property
+    def action_space(self) -> spaces.Space: ...
+
+    @property
+    def unwrapped(self) -> "Environment":
+        return self

envelope/spaces.py

@@ -0,0 +1,205 @@
+from abc import ABC, abstractmethod
+from functools import cached_property
+from typing import override
+
+import jax
+from jax import numpy as jnp
+
+from envelope.struct import FrozenPyTreeNode, static_field
+from envelope.typing import Key, PyTree
+
+
+class Space(ABC, FrozenPyTreeNode):
+    @abstractmethod
+    def sample(self, key: Key) -> PyTree: ...
+
+    @abstractmethod
+    def contains(self, x: PyTree) -> bool: ...
+
+
+class Discrete(Space):
+    """
+    A discrete space with a given number of elements. `n` can be a scalar or an array.
+    The shape and dtype of the space are inferred from `n`.
+
+    Args:
+        n: The number of elements in the space.
+    """
+
+    n: int | jax.Array
+
+    @classmethod
+    def from_shape(cls, n: int, shape: tuple[int]) -> "Discrete":
+        return cls(n=jnp.full(shape, n, dtype=jnp.asarray(n).dtype))
+
+    @property
+    def shape(self) -> tuple[int, ...]:
+        return jnp.asarray(self.n).shape
+
+    @property
+    def dtype(self):
+        return jnp.asarray(self.n).dtype
+
+    def sample(self, key: Key) -> jax.Array:
+        return jax.random.randint(key, self.shape, 0, self.n, dtype=self.dtype)
+
+    def contains(self, x: int | jax.Array) -> bool:
+        return jnp.all(x >= 0) & jnp.all(x < self.n)
+
+    def __repr__(self) -> str:
+        return f"Discrete(shape={self.shape}, dtype={self.dtype}, n={self.n})"
+
+
+class Continuous(Space):
+    """
+    A continuous space with a given lower and upper bound. `low` and `high` can be
+    scalars or arrays. The shape and dtype of the space are inferred from `low` and
+    `high`.
+
+    Args:
+        low: The lower bound of the space.
+        high: The upper bound of the space.
+    """
+
+    low: float | jax.Array
+    high: float | jax.Array
+
+    @classmethod
+    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),
+        )
+
+    @property
+    def dtype(self):
+        if jnp.asarray(self.low).dtype != jnp.asarray(self.high).dtype:
+            raise ValueError("low and high must have the same dtype")
+
+        return jnp.asarray(self.low).dtype
+
+    @property
+    def shape(self) -> tuple[int, ...]:
+        if jnp.asarray(self.low).shape != jnp.asarray(self.high).shape:
+            raise ValueError("low and high must have the same shape")
+
+        return jnp.asarray(self.low).shape
+
+    @override
+    def sample(self, key: Key) -> jax.Array:
+        uniform_sample = jax.random.uniform(key, self.shape, self.dtype)
+        return self.low + uniform_sample * (self.high - self.low)
+
+    @override
+    def contains(self, x: jax.Array) -> bool:
+        return jnp.all((x >= jnp.asarray(self.low)) & (x <= jnp.asarray(self.high)))
+
+    def __repr__(self) -> str:
+        dtype_str = getattr(self.dtype, "__name__", str(self.dtype))
+        return (
+            f"Continuous(shape={self.shape}, dtype={dtype_str}, "
+            f"low={self.low}, high={self.high})"
+        )
+
+
+class PyTreeSpace(Space):
+    """A Space defined by a PyTree structure of other Spaces.
+
+    Args:
+        tree: A PyTree with Space objects leaves.
+
+    Usage:
+        space = PyTreeSpace({
+            "action": Discrete(n=4, dtype=jnp.int32),
+            "obs": Continuous(low=0.0, high=1.0, shape=(2,), dtype=jnp.float32)
+        })
+    """
+
+    tree: PyTree
+
+    @override
+    def sample(self, key: Key) -> PyTree:
+        leaves, treedef = jax.tree.flatten(
+            self.tree, is_leaf=lambda x: isinstance(x, Space)
+        )
+        keys = jax.random.split(key, len(leaves))
+        samples = [space.sample(key) for key, space in zip(keys, leaves)]
+        return jax.tree.unflatten(treedef, samples)
+
+    @override
+    def contains(self, x: PyTree) -> bool:
+        # Use tree.map to check containment for each space-value pair
+        contains = jax.tree.map(
+            lambda space, xi: space.contains(xi),
+            self.tree,
+            x,
+            is_leaf=lambda node: isinstance(node, Space),
+        )
+        return jnp.all(jnp.array(jax.tree.leaves(contains)))
+
+    def __repr__(self) -> str:
+        """Return a string representation showing the tree structure."""
+        return f"{self.__class__.__name__}({self.tree!r})"
+
+    @property
+    def shape(self) -> PyTree:
+        return jax.tree.map(
+            lambda space: space.shape,
+            self.tree,
+            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,
+            is_leaf=lambda node: isinstance(node, Space),
+        )
+        return PyTreeSpace(batched_tree)
+    return BatchedSpace(space=space, batch_size=batch_size)
+
+
+class BatchedSpace(Space):
+    """
+    A view that adds a leading batch dimension to a base Space without
+    materializing or broadcasting its parameters.
+    """
+
+    space: Space
+    batch_size: int = static_field()
+
+    def sample(self, key: Key) -> PyTree:
+        # Accept single PRNGKey or a batch of keys shaped (batch_size, 2)
+        if getattr(key, "shape", ()) == (2,):
+            keys = jax.random.split(key, self.batch_size)
+        else:
+            if key.shape[0] != self.batch_size:
+                raise ValueError(
+                    f"sample key's leading dimension ({key.shape[0]}) must match "
+                    f"batch_size ({self.batch_size})."
+                )
+            keys = key
+        return jax.vmap(self.space.sample)(keys)
+
+    def contains(self, x: PyTree) -> bool:
+        # x is expected to be batched on the leading dimension
+        result = jax.vmap(self.space.contains)(x)
+        return jnp.all(jnp.asarray(result))
+
+    @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
+
+    @property
+    def dtype(self):
+        return getattr(self.space, "dtype", None)
+
+    def __repr__(self) -> str:
+        return f"BatchedSpace(space={self.space!r}, batch_size={self.batch_size})"

envelope/struct.py

@@ -0,0 +1,148 @@
+import dataclasses
+from dataclasses import KW_ONLY
+from typing import Any, Iterable, Iterator, Mapping, Self, Tuple
+
+import jax
+
+from envelope.typing import PyTree
+
+__all__ = ["FrozenPyTreeNode", "field", "static_field", "Container"]
+
+
+def field(*, pytree_node: bool = True, **kwargs):
+    """
+    Dataclass field helper.
+    Set pytree_node=False for static (non-transformed) fields.
+    """
+    meta = dict(kwargs.pop("metadata", {}) or {})
+    meta["pytree_node"] = pytree_node
+    return dataclasses.field(metadata=meta, **kwargs)
+
+
+def static_field(**kwargs):
+    """Shorthand for field(pytree_node=False, ...)."""
+    return field(pytree_node=False, **kwargs)
+
+
+class FrozenPyTreeNode:
+    """
+    Frozen dataclass base that is a JAX pytree node.
+
+    Usage:
+        class Foo(FrozenPyTreeNode):
+            a: Any                      # pytree leaf
+            b: int = static_field()     # static, not a leaf
+
+        x = Foo(a={"w": 1.0}, b=0)
+        y = x.replace(b=1)
+    """
+
+    # Turn subclasses into frozen dataclasses and register with JAX.
+    def __init_subclass__(cls, *, dataclass_kwargs: dict[str, Any] | None = None, **kw):
+        super().__init_subclass__(**kw)
+        # Check if this specific class (not parent) has already been processed
+        if "__is_envelope_pytreenode__" in cls.__dict__:
+            return
+        opts = dict(frozen=True, eq=True, repr=True, slots=False)
+        if dataclass_kwargs:
+            opts.update(dataclass_kwargs)
+        dataclasses.dataclass(cls, **opts)  # modify in place
+        cls.__is_envelope_pytreenode__ = True
+
+        data = []
+        static = []
+        for f in dataclasses.fields(cls):
+            if f.metadata.get("pytree_node", True):
+                data.append(f.name)
+            else:
+                static.append(f.name)
+
+        jax.tree_util.register_dataclass(cls, data, static)
+
+    # convenience
+    def replace(self, **changes):
+        return dataclasses.replace(self, **changes)
+
+
+@jax.tree_util.register_pytree_node_class
+@dataclasses.dataclass(frozen=True, eq=True, repr=True, slots=False)
+class Container:
+    _: KW_ONLY
+    _extras: Mapping[str, PyTree] = dataclasses.field(default_factory=dict, repr=False)
+
+    def __init_subclass__(cls, *, dataclass_kwargs: dict[str, Any] | None = None, **kw):
+        super().__init_subclass__(**kw)
+        if "__is_container_dataclass__" in cls.__dict__:
+            return
+
+        opts = dict(frozen=True, eq=True, repr=True, slots=False)
+        if dataclass_kwargs:
+            opts.update(dataclass_kwargs)
+
+        dataclasses.dataclass(cls, **opts)
+        cls.__is_container_dataclass__ = True
+        jax.tree_util.register_pytree_node_class(cls)
+
+    def __getattr__(self, name: str) -> PyTree:
+        # bypass __getattr__ when accessing _extras to avoid recursion
+        extras = object.__getattribute__(self, "_extras")
+        if name in extras:
+            return extras[name]
+        self_name = type(self).__name__
+        raise AttributeError(f"'{self_name}' object has no attribute '{name}'")
+
+    def __dir__(self) -> Iterable[str]:
+        core_names = {f.name for f in dataclasses.fields(self) if f.name != "_extras"}
+        return sorted(set(super().__dir__()) | core_names | set(self._extras.keys()))
+
+    def __iter__(self) -> Iterator[Tuple[str, PyTree]]:
+        for f in dataclasses.fields(self):
+            if f.name == "_extras":
+                continue
+            yield (f.name, getattr(self, f.name))
+        # extras
+        for k, v in self._extras.items():
+            yield (k, v)
+
+    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] = {}
+        extras_updates: dict[str, PyTree] = {}
+
+        for k, v in changes.items():
+            if k in core_names:
+                core_updates[k] = v
+            else:
+                extras_updates[k] = v
+
+        new = dataclasses.replace(self, **core_updates)
+        new_extras = {**self._extras, **extras_updates}
+        object.__setattr__(new, "_extras", new_extras)
+        return new
+
+    def tree_flatten(self) -> Tuple[Tuple[PyTree, ...], Tuple[Any, ...]]:
+        core_fields = [f for f in dataclasses.fields(self) if f.name != "_extras"]
+        core_keys = tuple(f.name for f in core_fields)
+        core_vals = tuple(getattr(self, name) for name in core_keys)
+
+        extras_keys = tuple(self._extras.keys())
+        extras_vals = tuple(self._extras[k] for k in extras_keys)
+
+        children = core_vals + extras_vals
+        aux_data = (self.__class__, core_keys, extras_keys)
+        return children, aux_data
+
+    @classmethod
+    def tree_unflatten(cls, aux_data, children: Tuple[PyTree, ...]) -> Self:
+        actual_cls, core_keys, extras_keys = aux_data
+        n_core = len(core_keys)
+
+        core_vals = children[:n_core]
+        extras_vals = children[n_core:]
+
+        core_kwargs = dict(zip(core_keys, core_vals))
+        extras = dict(zip(extras_keys, extras_vals))
+
+        obj = actual_cls(**core_kwargs)
+        object.__setattr__(obj, "_extras", extras)
+        return obj

envelope/typing.py

@@ -0,0 +1,23 @@
+from enum import Enum
+from typing import Any, TypeAlias
+
+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/autoreset_wrapper.py

@@ -0,0 +1,36 @@
+import jax
+
+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):
+    class AutoResetState(WrappedState):
+        reset_key: jax.Array = field()
+
+    def reset(
+        self, key: Key, state: PyTree | None = None, **kwargs
+    ) -> 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),
+        )
+        return state, info