jax-envelope 0.1.0__py3-none-any.whl

envelope/compat/__init__.py

@@ -0,0 +1,97 @@
+"""Compatibility wrappers for various RL environment libraries."""
+
+from typing import Any, Protocol, Self
+
+# Lazy imports to avoid requiring all dependencies at once
+_env_module_map = {
+    "gymnax": ("envelope.compat.gymnax_envelope", "GymnaxEnvelope"),
+    "brax": ("envelope.compat.brax_envelope", "BraxEnvelope"),
+    "navix": ("envelope.compat.navix_envelope", "NavixEnvelope"),
+    "jumanji": ("envelope.compat.jumanji_envelope", "JumanjiEnvelope"),
+    "kinetix": ("envelope.compat.kinetix_envelope", "KinetixEnvelope"),
+    "craftax": ("envelope.compat.craftax_envelope", "CraftaxEnvelope"),
+    "mujoco_playground": (
+        "envelope.compat.mujoco_playground_envelope",
+        "MujocoPlaygroundEnvelope",
+    ),
+}
+
+
+class HasFromNameInit(Protocol):
+    @classmethod
+    def from_name(
+        cls, env_name: str, env_kwargs: dict[str, Any] | None = None, **kwargs
+    ) -> Self: ...
+
+    """Creates an environment from a name and keyword arguments. Unless otherwise noted,
+    the created environment will have it's default parameters, with truncation and auto
+    reset disabled.
+    
+    Args:
+        env_name: Environment name
+        env_kwargs: Keyword arguments passed to the environment constructor
+        **kwargs: Additional keyword arguments passed to the environment wrapper
+    """
+
+
+def create(env_name: str, env_kwargs: dict[str, Any] | None = None, **kwargs):
+    """Create an environment from a prefixed environment ID.
+
+    Args:
+        env_name: Environment ID in the format "suite::env_name" (e.g., "brax::ant")
+        env_kwargs: Keyword arguments passed to the suite's environment constructor
+        **kwargs: Additional keyword arguments passed to the environment wrapper
+
+    Returns:
+        An instance of the wrapped environment
+
+    Examples:
+        >>> env = create("jumanji::snake")
+        >>> env = create("brax::ant", env_kwargs={"backend": "spring"})
+        >>> env = create("gymnax::CartPole-v1", env_params=...)
+    """
+    original_env_id = env_name
+    if "::" not in env_name:
+        raise ValueError(
+            f"Environment ID must be in format 'suite::env_name', got: {original_env_id}"
+        )
+
+    suite, env_name = env_name.split("::", 1)
+    if not suite or not env_name:
+        raise ValueError(
+            f"Environment ID must be in format 'suite::env_name', got: {original_env_id}"
+        )
+
+    if suite not in _env_module_map:
+        raise ValueError(
+            f"Unknown environment suite: {suite}. "
+            f"Available suites: {list(_env_module_map.keys())}"
+        )
+
+    # Lazy import the wrapper class
+    module_name, class_name = _env_module_map[suite]
+    try:
+        import importlib
+
+        module = importlib.import_module(module_name)
+        env_class: HasFromNameInit = getattr(module, class_name)
+    except ImportError as e:
+        raise ImportError(
+            f"Failed to import {suite} wrapper. "
+            f"Make sure you have installed the '{suite}' dependencies. "
+            f"Original error: {e}"
+        ) from e
+
+    env = env_class.from_name(env_name, env_kwargs=env_kwargs, **kwargs)
+
+    # Wrap with TruncationWrapper using adapter's default
+    default_max_steps = getattr(env, "default_max_steps", None)
+    if default_max_steps is not None:
+        from envelope.wrappers.truncation_wrapper import TruncationWrapper
+
+        env = TruncationWrapper(env=env, max_steps=int(default_max_steps))
+
+    return env
+
+
+__all__ = ["create"]

envelope/compat/brax_envelope.py

@@ -0,0 +1,98 @@
+import dataclasses
+import warnings
+from copy import copy
+from functools import cached_property
+from typing import Any, override
+
+from brax.envs import Env as BraxEnv
+from brax.envs import Wrapper as BraxWrapper
+from brax.envs import create as brax_create
+from jax import numpy as jnp
+
+from envelope import spaces
+from envelope.environment import Environment, Info, InfoContainer, State
+from envelope.struct import static_field
+from envelope.typing import Key, PyTree
+
+# Default episode_length in brax.envs.create()
+_BRAX_DEFAULT_EPISODE_LENGTH = 1000
+
+
+class BraxEnvelope(Environment):
+    """Wrapper to convert a Brax environment to a envelope environment."""
+
+    brax_env: BraxEnv = static_field()
+
+    @classmethod
+    def from_name(
+        cls, env_name: str, env_kwargs: dict[str, Any] | None = None
+    ) -> "BraxEnvelope":
+        env_kwargs = env_kwargs or {}
+        if "episode_length" in env_kwargs:
+            raise ValueError(
+                "Cannot override 'episode_length' directly. "
+                "Use TruncationWrapper for episode length control."
+            )
+        if "auto_reset" in env_kwargs:
+            raise ValueError(
+                "Cannot override 'auto_reset' directly. "
+                "Use AutoResetWrapper for auto-reset behavior."
+            )
+
+        env_kwargs["episode_length"] = jnp.inf
+        env_kwargs["auto_reset"] = False
+        env = brax_create(env_name, **env_kwargs)
+        return cls(brax_env=env)
+
+    @property
+    def default_max_steps(self) -> int:
+        return _BRAX_DEFAULT_EPISODE_LENGTH
+
+    def __post_init__(self) -> "BraxEnvelope":
+        if isinstance(self.brax_env, BraxWrapper):
+            warnings.warn(
+                "Environment wrapping should be handled by envelope. "
+                "Unwrapping brax environment before converting..."
+            )
+            object.__setattr__(self, "brax_env", self.brax_env.unwrapped)
+
+    @override
+    def reset(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))
+        return brax_state, info
+
+    @override
+    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
+        )
+        info = info.update(**dataclasses.asdict(brax_state))
+        return brax_state, info
+
+    @override
+    @cached_property
+    def action_space(self) -> spaces.Space:
+        # All brax environments have action limit of -1 to 1
+        return spaces.Continuous.from_shape(
+            low=-1.0, high=1.0, shape=(self.brax_env.action_size,)
+        )
+
+    @override
+    @cached_property
+    def observation_space(self) -> spaces.Space:
+        # All brax environments have observation limit of -inf to inf
+        return spaces.Continuous.from_shape(
+            low=-jnp.inf, high=jnp.inf, shape=(self.brax_env.observation_size,)
+        )
+
+    def __deepcopy__(self, memo):
+        warnings.warn(
+            f"Trying to deepcopy {type(self).__name__}, which contains a brax env. "
+            "Brax envs throw an error when deepcopying, so a shallow copy is returned.",
+            category=RuntimeWarning,
+            stacklevel=2,
+        )
+        return copy(self)

envelope/compat/craftax_envelope.py

@@ -0,0 +1,86 @@
+from functools import cached_property
+from typing import Any, override
+
+import jax
+import jax.numpy as jnp
+from craftax.craftax.craftax_state import EnvParams as CraftaxEnvParams
+from craftax.craftax_classic.envs.craftax_state import (
+    EnvParams as CraftaxClassicEnvParams,
+)
+from craftax.craftax_env import make_craftax_env_from_name
+
+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
+from envelope.struct import Container, static_field
+from envelope.typing import Key, PyTree, TypeAlias
+
+EnvParams: TypeAlias = CraftaxEnvParams | CraftaxClassicEnvParams
+
+
+class CraftaxEnvelope(Environment):
+    """Wrapper to convert a Craftax environment to a envelope environment."""
+
+    craftax_env: Any = static_field()
+    env_params: PyTree
+
+    @classmethod
+    def from_name(
+        cls,
+        env_name: str,
+        env_params: EnvParams | None = None,
+        env_kwargs: dict[str, Any] | None = None,
+    ) -> "CraftaxEnvelope":
+        env_kwargs = env_kwargs or {}
+        if "max_timesteps" in env_kwargs:
+            raise ValueError(
+                "Cannot override 'max_timesteps' directly. "
+                "Use TruncationWrapper for episode length control."
+            )
+        if "auto_reset" in env_kwargs:
+            raise ValueError(
+                "Cannot override 'auto_reset' directly. "
+                "Use AutoResetWrapper for auto-reset behavior."
+            )
+
+        env_kwargs["auto_reset"] = False
+        env = make_craftax_env_from_name(env_name, **env_kwargs)
+        default_params = env.default_params.replace(max_timesteps=jnp.inf)
+
+        env_params = env_params or default_params
+        return cls(craftax_env=env, env_params=env_params)
+
+    @property
+    def default_max_steps(self) -> int:
+        return int(self.craftax_env.default_params.max_timesteps)
+
+    @override
+    def reset(self, key: Key) -> tuple[State, Info]:
+        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)
+        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.craftax_env.step(
+            subkey, state.env_state, action, self.env_params
+        )
+        state = state.update(key=key, env_state=env_state)
+        info = InfoContainer(obs=obs, reward=reward, terminated=done)
+        info = info.update(info=env_info)
+        return state, info
+
+    @override
+    @cached_property
+    def action_space(self) -> envelope_spaces.Space:
+        return _convert_gymnax_space(self.craftax_env.action_space(self.env_params))
+
+    @override
+    @cached_property
+    def observation_space(self) -> envelope_spaces.Space:
+        return _convert_gymnax_space(
+            self.craftax_env.observation_space(self.env_params)
+        )

envelope/compat/gymnax_envelope.py

@@ -0,0 +1,91 @@
+from functools import cached_property
+from typing import Any, override
+
+import jax
+import jax.numpy as jnp
+from gymnax import make as gymnax_create
+from gymnax.environments import spaces as gymnax_spaces
+from gymnax.environments.environment import Environment as GymnaxEnv
+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.typing import Key, PyTree
+
+
+class GymnaxEnvelope(Environment):
+    """Wrapper to convert a Gymnax environment to a envelope environment."""
+
+    gymnax_env: GymnaxEnv = static_field()
+    env_params: PyTree
+
+    @classmethod
+    def from_name(
+        cls,
+        env_name: str,
+        env_params: GymnaxEnvParams | None = None,
+        env_kwargs: dict[str, Any] | None = None,
+    ) -> "GymnaxEnvelope":
+        env_kwargs = env_kwargs or {}
+        if "max_steps_in_episode" in env_kwargs:
+            raise ValueError(
+                "Cannot override 'max_steps_in_episode' directly. "
+                "Use TruncationWrapper for episode length control."
+            )
+        gymnax_env, default_params = gymnax_create(env_name, **env_kwargs)
+        default_params = default_params.replace(max_steps_in_episode=jnp.inf)
+
+        env_params = env_params or default_params
+        return cls(gymnax_env=gymnax_env, env_params=env_params)
+
+    @property
+    def default_max_steps(self) -> int:
+        return int(self.gymnax_env.default_params.max_steps_in_episode)
+
+    @override
+    def reset(self, key: Key) -> tuple[State, Info]:
+        key, subkey = jax.random.split(key)
+        obs, env_state = self.gymnax_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=None)
+        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(
+            subkey, state.env_state, action, self.env_params
+        )
+        state = state.update(key=key, env_state=env_state)
+        info = InfoContainer(obs=obs, reward=reward, terminated=done)
+        info = info.update(info=env_info)
+        return state, info
+
+    @override
+    @cached_property
+    def action_space(self) -> envelope_spaces.Space:
+        return _convert_space(self.gymnax_env.action_space(self.env_params))
+
+    @override
+    @cached_property
+    def observation_space(self) -> envelope_spaces.Space:
+        return _convert_space(self.gymnax_env.observation_space(self.env_params))
+
+
+def _convert_space(gmx_space: gymnax_spaces.Space) -> envelope_spaces.Space:
+    if isinstance(gmx_space, gymnax_spaces.Box):
+        low = jnp.broadcast_to(gmx_space.low, gmx_space.shape).astype(gmx_space.dtype)
+        high = jnp.broadcast_to(gmx_space.high, gmx_space.shape).astype(gmx_space.dtype)
+        return envelope_spaces.Continuous(low=low, high=high)
+    elif isinstance(gmx_space, gymnax_spaces.Discrete):
+        n = jnp.broadcast_to(gmx_space.n, gmx_space.shape).astype(gmx_space.dtype)
+        return envelope_spaces.Discrete(n=n)
+    elif isinstance(gmx_space, gymnax_spaces.Tuple):
+        spaces = tuple(_convert_space(space) for space in gmx_space.spaces)
+        return envelope_spaces.PyTreeSpace(spaces)
+    elif isinstance(gmx_space, gymnax_spaces.Dict):
+        spaces = {k: _convert_space(space) for k, space in gmx_space.spaces.items()}
+        return envelope_spaces.PyTreeSpace(spaces)
+    raise ValueError(f"Unsupported space type: {type(gmx_space)}")

envelope/compat/jumanji_envelope.py

@@ -0,0 +1,127 @@
+import warnings
+from copy import copy
+from functools import cached_property
+from typing import Any, override
+
+import jax.numpy as jnp
+import jumanji
+from jumanji.specs import Array, BoundedArray, DiscreteArray, MultiDiscreteArray
+from jumanji.types import TimeStep as JumanjiTimeStep
+
+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 = jnp.iinfo(jnp.int32).max
+
+
+class JumanjiEnvelope(Environment):
+    """Wrapper to convert a Jumanji environment to a envelope environment."""
+
+    jumanji_env: Any = static_field()
+    _default_time_limit: int | None = static_field(default=None)
+
+    @classmethod
+    def from_name(
+        cls, env_name: str, env_kwargs: dict[str, Any] | None = None
+    ) -> "JumanjiEnvelope":
+        env_kwargs = env_kwargs or {}
+        if "time_limit" in env_kwargs:
+            raise ValueError(
+                "Cannot override 'time_limit' directly. "
+                "Use TruncationWrapper for episode length control."
+            )
+
+        # Create env first with defaults to capture default time_limit
+        temp_env = jumanji.make(env_name, **env_kwargs)
+        default_time_limit = getattr(temp_env, "time_limit", None)
+
+        # Now create env with time_limit=_MAX_INT (if env supports it)
+        if default_time_limit is not None:
+            env_kwargs["time_limit"] = _MAX_INT
+        env = jumanji.make(env_name, **env_kwargs)
+        return cls(jumanji_env=env, _default_time_limit=default_time_limit)
+
+    @property
+    def default_max_steps(self) -> int | None:
+        return self._default_time_limit
+
+    @override
+    def reset(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
+
+    @override
+    def step(self, state: State, action: PyTree) -> tuple[State, Info]:
+        env_state, timestep = self.jumanji_env.step(state, action)
+        info = convert_jumanji_to_envelope_info(timestep)
+        return env_state, info
+
+    @override
+    @cached_property
+    def action_space(self) -> envelope_spaces.Space:
+        return convert_jumanji_spec_to_envelope_space(self.jumanji_env.action_spec)
+
+    @override
+    @cached_property
+    def observation_space(self) -> envelope_spaces.Space:
+        return convert_jumanji_spec_to_envelope_space(self.jumanji_env.observation_spec)
+
+    def __deepcopy__(self, memo):
+        warnings.warn(
+            f"Trying to deepcopy {type(self).__name__}, which contains a jumanji env. "
+            "Jumanji envs may throw an error when deepcopying, so a shallow copy is "
+            "returned.",
+            category=RuntimeWarning,
+            stacklevel=2,
+        )
+        return copy(self)
+
+
+def convert_jumanji_to_envelope_info(timestep: JumanjiTimeStep) -> InfoContainer:
+    info = InfoContainer(
+        obs=timestep.observation, reward=timestep.reward, terminated=timestep.last()
+    ).update(**timestep.extras)
+    return info
+
+
+def convert_jumanji_spec_to_envelope_space(spec: Any) -> envelope_spaces.Space:
+    """Convert a Jumanji Spec to a envelope Space."""
+
+    if isinstance(spec, (DiscreteArray, MultiDiscreteArray)):
+        n = jnp.asarray(spec.num_values, dtype=spec.dtype)
+        if getattr(spec, "shape", ()) not in ((), n.shape):
+            n = jnp.broadcast_to(n, spec.shape)
+        return envelope_spaces.Discrete(n=n)
+
+    if isinstance(spec, BoundedArray):
+        low = jnp.broadcast_to(jnp.asarray(spec.minimum, dtype=spec.dtype), spec.shape)
+        high = jnp.broadcast_to(jnp.asarray(spec.maximum, dtype=spec.dtype), spec.shape)
+        return envelope_spaces.Continuous(low=low, high=high)
+
+    if isinstance(spec, Array):
+        dtype = jnp.dtype(spec.dtype)
+        if not jnp.issubdtype(dtype, jnp.floating):
+            raise NotImplementedError(
+                "Unbounded jumanji Array specs are only supported for floating dtypes. "
+                f"Got dtype={dtype} for spec={spec!r}."
+            )
+        low = jnp.full(spec.shape, -jnp.inf, dtype=dtype)
+        high = jnp.full(spec.shape, jnp.inf, dtype=dtype)
+        return envelope_spaces.Continuous(low=low, high=high)
+
+    # Structured specs (most Jumanji envs): access private mapping when available.
+    subspecs = getattr(spec, "_specs", None)
+    if isinstance(subspecs, dict):
+        tree = {
+            k: convert_jumanji_spec_to_envelope_space(v) for k, v in subspecs.items()
+        }
+        return envelope_spaces.PyTreeSpace(tree)
+
+    if isinstance(spec, (tuple, list)):
+        tree = tuple(convert_jumanji_spec_to_envelope_space(s) for s in spec)
+        return envelope_spaces.PyTreeSpace(tree)
+
+    raise ValueError(f"Unsupported spec type: {type(spec)}")

envelope/compat/kinetix_envelope.py

@@ -0,0 +1,194 @@
+"""Kinetix compatibility wrapper.
+
+This module exposes Kinetix environments through the `envelope.environment.Environment`
+API. It mirrors envelope's compat philosophy:
+- prefer *no* environment-side auto-reset (use `AutoResetWrapper` in envelope)
+- prefer *no* fixed episode time-limits (use `TruncationWrapper` in envelope)
+
+`from_name` supports premade level ids like `s/h4_thrust_aim` (optionally with
+`.json`). For maximum flexibility, users can bypass level handling entirely by
+passing a custom `reset_fn`.
+"""
+
+from __future__ import annotations
+
+import warnings
+from functools import cached_property
+from typing import Any, Callable, Literal, override
+
+import jax
+import jax.numpy as jnp
+from kinetix.environment import (
+    ActionType,
+    EnvParams,
+    ObservationType,
+    StaticEnvParams,
+    make_kinetix_env,
+)
+from kinetix.environment.ued.ued import make_reset_fn_sample_kinetix_level
+from kinetix.util.saving import load_from_json_file
+
+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
+from envelope.struct import Container, static_field
+from envelope.typing import Key, PyTree
+
+LevelResetFn = Callable[[Key], Any]
+
+
+def _normalize_level_id(level_id: str) -> str:
+    """Normalize a path-like level id.
+
+    Examples:
+        - ``"s/h4_thrust_aim"`` -> ``"s/h4_thrust_aim.json"``
+        - ``"/s/h4_thrust_aim.json"`` -> ``"s/h4_thrust_aim.json"``
+    """
+    level_id = level_id.strip().lstrip("/")
+    if not level_id:
+        raise ValueError("level_id must be a non-empty string")
+    if level_id.endswith("/"):
+        raise ValueError("level_id must not end with '/'")
+    if not level_id.endswith(".json"):
+        level_id = f"{level_id}.json"
+    return level_id
+
+
+def _warn_auto_reset(auto_reset: bool) -> None:
+    if auto_reset:
+        warnings.warn(
+            "Creating a KinetixEnvelope with auto_reset=True is not recommended, use "
+            "an AutoResetWrapper instead.",
+            stacklevel=2,
+        )
+
+
+class KinetixEnvelope(Environment):
+    """Wrapper to convert a Kinetix environment to a envelope environment."""
+
+    kinetix_env: Any = static_field()
+    env_params: Any
+
+    @property
+    def default_max_steps(self) -> int:
+        return int(EnvParams().max_timesteps)
+
+    @classmethod
+    def from_name(
+        cls,
+        env_name: str | Literal["random"],
+        env_params: EnvParams | None = None,
+        env_kwargs: dict[str, Any] | None = None,
+    ) -> "KinetixEnvelope":
+        env_kwargs = env_kwargs or {}
+        if "max_timesteps" in env_kwargs:
+            raise ValueError(
+                "Cannot override 'max_timesteps' directly. "
+                "Use TruncationWrapper for episode length control."
+            )
+        if "auto_reset" in env_kwargs:
+            raise ValueError(
+                "Cannot override 'auto_reset' directly. "
+                "Use AutoResetWrapper for auto-reset behavior."
+            )
+
+        env_kwargs["auto_reset"] = False
+        if env_name == "random":
+            return cls.create_random(env_params=env_params, **env_kwargs)
+
+        if (
+            env_params is not None
+            or "env_params" in env_kwargs
+            or "static_env_params" in env_kwargs
+        ):
+            raise ValueError(
+                "env_params and static_env_params cannot be passed when creating a "
+                "KinetixEnvelope from a premade level."
+            )
+        return cls.create_premade(env_name, **env_kwargs)
+
+    @classmethod
+    def create_premade(
+        cls,
+        env_name: str,
+        action_type: ActionType = ActionType.CONTINUOUS,
+        observation_type: ObservationType = ObservationType.SYMBOLIC_FLAT,
+        auto_reset: bool = False,
+    ) -> "KinetixEnvelope":
+        _warn_auto_reset(auto_reset)
+
+        # Load level.
+        level_id_json = _normalize_level_id(env_name)
+        level, static_env_params, env_params = load_from_json_file(level_id_json)
+        env_params = env_params.replace(max_timesteps=jnp.inf) if env_params else None
+
+        def reset_fn(_: Key) -> Any:
+            return level
+
+        # Create environment.
+        kinetix_env = make_kinetix_env(
+            action_type=action_type,
+            observation_type=observation_type,
+            reset_fn=reset_fn,
+            env_params=env_params,
+            static_env_params=static_env_params,
+            auto_reset=auto_reset,
+        )
+        return cls(kinetix_env=kinetix_env, env_params=env_params)
+
+    @classmethod
+    def create_random(
+        cls,
+        action_type: ActionType = ActionType.CONTINUOUS,
+        observation_type: ObservationType = ObservationType.SYMBOLIC_FLAT,
+        env_params: EnvParams | None = None,
+        static_env_params: StaticEnvParams = StaticEnvParams(),
+        auto_reset: bool = False,
+    ) -> "KinetixEnvelope":
+        _warn_auto_reset(auto_reset)
+        if env_params is None:
+            env_params = EnvParams()
+        env_params = env_params.replace(max_timesteps=jnp.inf)
+
+        reset_fn = make_reset_fn_sample_kinetix_level(env_params, static_env_params)
+        kinetix_env = make_kinetix_env(
+            action_type=action_type,
+            observation_type=observation_type,
+            reset_fn=reset_fn,
+            env_params=env_params,
+            static_env_params=static_env_params,
+            auto_reset=auto_reset,
+        )
+        return cls(kinetix_env=kinetix_env, env_params=env_params)
+
+    @override
+    def reset(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)
+        info = InfoContainer(obs=obs, reward=0.0, terminated=False)
+        info = info.update(info=None)
+        return state_out, 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.kinetix_env.step(
+            subkey, state.env_state, action, self.env_params
+        )
+        state_out = state.update(key=key, env_state=env_state)
+        info = InfoContainer(obs=obs, reward=reward, terminated=done)
+        info = info.update(info=env_info)
+        return state_out, info
+
+    @override
+    @cached_property
+    def action_space(self) -> envelope_spaces.Space:
+        return _convert_gymnax_space(self.kinetix_env.action_space(self.env_params))
+
+    @override
+    @cached_property
+    def observation_space(self) -> envelope_spaces.Space:
+        return _convert_gymnax_space(
+            self.kinetix_env.observation_space(self.env_params)
+        )