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

{jax_envelope-0.2.0 → jax_envelope-0.4.0}/.github/workflows/publish.yml

@@ -28,7 +28,7 @@ jobs:
           uv sync --group dev --locked
 
       - name: Run tests
-        run: uv run pytest -m "not compat"
+        run: uv run pytest -m "not adapters"
 
   build:
     name: Build distribution

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: jax-envelope
-Version: 0.2.0
+Version: 0.4.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
@@ -25,12 +25,13 @@ Requires-Dist: jax>=0.5.0
 Description-Content-Type: text/markdown
 
 # 💌 Envelope: a JAX-native environment interface
+
 ```python
 # Create environments from JAX-native suites you have installed, ...
 env = envelope.create("gymnax::CartPole-v1")
 
 # ... interact with the environments using a simple interface, ...
-state, info = env.reset(key)
+state, info = env.init(key)
 states, infos = jax.lax.scan(env.step, state, actions)
 plt.plot(infos.reward.cumsum())
 
@@ -41,47 +42,56 @@ env = envelope.wrappers.ObservationNormalizationWrapper(env)
 ```
 
 ## 🌍 Simple, expressive interaction!
-* **Environments are pytrees**. Squish them through JAX transformations and trace their parameters.
-* **Idiomatic jax-y interface** of `reset(key: Key) -> State, Info` and `step(state: State, action: PyTree) -> State, Info`. You can directly `jax.scan` over a `step(...)`!
-* **Spaces are super simple**. No `Tuple`, `Dict` nonsense! There are two spaces: `Continuous` and `Discrete`, which you can compose into a `PyTreeSpace`.
-* **Explicit episode truncation** supports correctly handling bootstrapping for value-function targets.
-* **No auto-reset** by default. Resetting every step can be expensive!
+
+- **Environments are pytrees**. Squish them through JAX transformations and trace their parameters.
+- **Idiomatic jax-y interface** of `init(key: Key) -> State, Info` and `step(state: State, action: PyTree) -> State, Info`. You can directly `jax.scan` over a `step(...)`!
+- **Spaces are super simple**. No `Tuple`, `Dict` nonsense! There are two spaces: `Continuous` and `Discrete`, which you can compose into a `PyTreeSpace`.
+- **Explicit episode truncation** supports correctly handling bootstrapping for value-function targets.
+- **No auto-reset** by default. Resetting every step can be expensive!
 
 ## 💪 Powerful, composable wrappers!
-* **Carry state across episodes** to track running statistics, for example to normalize observations.
-* **Composable wrappers** can be stacked in any order. For example, `ObservationNormalizationWrapper` before vs. after `VmapWrapper` gives per-env vs. global normalization.
-<!-- TODO: Add auto-reset behavior (including state injection) and optimistic resets once I implement them. -->
+
+- **Carry state across episodes** to track running statistics, for example to normalize observations.
+- **Composable wrappers** can be stacked in any order. For example, `ObservationNormalizationWrapper` before vs. after `VmapWrapper` gives per-env vs. global normalization.
 
 ## 🔌 Adapters for existing suites
-| 📦 | # 🤖 | # 🌍 |
-|------|------|------|
-| [gymnax](https://github.com/RobertTLange/gymnax) | 🕺 | 24 |
-| [brax](https://github.com/google/brax) | 🕺 | 12 |
-| [jumanji](https://github.com/instadeepai/jumanji) | 🕺 / 👯 | 25 / 1 |
-| [kinetix](https://github.com/flairox/kinetix) | 🕺 | 74 |
-| [craftax](https://github.com/MichaelTMatthews/craftax) | 🕺 | 4 |
-| [mujoco_playground](https://github.com/google-deepmind/mujoco_playground) | 🕺 | 54 |
-| | |
-| Total | 🕺 / 👯 | 193 / 1 |
+
+
+| 📦                                                                        | # 🤖    | # 🌍    |
+| ------------------------------------------------------------------------- | ------- | ------- |
+| [gymnax](https://github.com/RobertTLange/gymnax)                          | 🕺      | 24      |
+| [brax](https://github.com/google/brax)                                    | 🕺      | 12      |
+| [jumanji](https://github.com/instadeepai/jumanji)                         | 🕺 / 👯 | 25 / 1  |
+| [kinetix](https://github.com/flairox/kinetix)                             | 🕺      | 74      |
+| [craftax](https://github.com/MichaelTMatthews/craftax)                    | 🕺      | 4       |
+| [mujoco_playground](https://github.com/google-deepmind/mujoco_playground) | 🕺      | 54      |
+|                                                                           |         |         |
+| Total                                                                     | 🕺 / 👯 | 193 / 1 |
+
 
 ```python
 envelope.create("📦::🌍")
 ```
+
 let's you create environments from any of the above!
 
 ## 📝 Testing
-- **Default (no optional compat deps required)**: `uv run pytest -m "not compat"`
-- **Compat suite (requires full compat dependency group)**:
-  - `uv sync --group compat`
-  - `uv run pytest -m compat`
-  - If any compat dependency is missing/broken, the run will fail fast with an error telling you what to install.
+
+- **Default (no optional adapters deps required)**: `uv run pytest -m "not adapters"`
+- **Adapters suite (requires full adapters dependency group)**:
+  - `uv sync --group adapters`
+  - `uv run pytest -m adapters`
+  - If any adapter dependency is missing/broken, the run will fail fast with an error telling you what to install.
 
 ## 🏗️ Installation
+
 ```bash
 pip install jax-envelope
 ```
 
 ## 💞 Related projects
-* [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!
+
+- [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.4.0/README.md

@@ -0,0 +1,71 @@
+# 💌 Envelope: a JAX-native environment interface
+
+```python
+# Create environments from JAX-native suites you have installed, ...
+env = envelope.create("gymnax::CartPole-v1")
+
+# ... interact with the environments using a simple interface, ...
+state, info = env.init(key)
+states, infos = jax.lax.scan(env.step, state, actions)
+plt.plot(infos.reward.cumsum())
+
+# ... and enjoy a powerful ecosystem of wrappers.
+env = envelope.wrappers.AutoResetWrapper(env)
+env = envelope.wrappers.VmapWrapper(env)
+env = envelope.wrappers.ObservationNormalizationWrapper(env)
+```
+
+## 🌍 Simple, expressive interaction!
+
+- **Environments are pytrees**. Squish them through JAX transformations and trace their parameters.
+- **Idiomatic jax-y interface** of `init(key: Key) -> State, Info` and `step(state: State, action: PyTree) -> State, Info`. You can directly `jax.scan` over a `step(...)`!
+- **Spaces are super simple**. No `Tuple`, `Dict` nonsense! There are two spaces: `Continuous` and `Discrete`, which you can compose into a `PyTreeSpace`.
+- **Explicit episode truncation** supports correctly handling bootstrapping for value-function targets.
+- **No auto-reset** by default. Resetting every step can be expensive!
+
+## 💪 Powerful, composable wrappers!
+
+- **Carry state across episodes** to track running statistics, for example to normalize observations.
+- **Composable wrappers** can be stacked in any order. For example, `ObservationNormalizationWrapper` before vs. after `VmapWrapper` gives per-env vs. global normalization.
+
+## 🔌 Adapters for existing suites
+
+
+| 📦                                                                        | # 🤖    | # 🌍    |
+| ------------------------------------------------------------------------- | ------- | ------- |
+| [gymnax](https://github.com/RobertTLange/gymnax)                          | 🕺      | 24      |
+| [brax](https://github.com/google/brax)                                    | 🕺      | 12      |
+| [jumanji](https://github.com/instadeepai/jumanji)                         | 🕺 / 👯 | 25 / 1  |
+| [kinetix](https://github.com/flairox/kinetix)                             | 🕺      | 74      |
+| [craftax](https://github.com/MichaelTMatthews/craftax)                    | 🕺      | 4       |
+| [mujoco_playground](https://github.com/google-deepmind/mujoco_playground) | 🕺      | 54      |
+|                                                                           |         |         |
+| Total                                                                     | 🕺 / 👯 | 193 / 1 |
+
+
+```python
+envelope.create("📦::🌍")
+```
+
+let's you create environments from any of the above!
+
+## 📝 Testing
+
+- **Default (no optional adapters deps required)**: `uv run pytest -m "not adapters"`
+- **Adapters suite (requires full adapters dependency group)**:
+  - `uv sync --group adapters`
+  - `uv run pytest -m adapters`
+  - If any adapter dependency is missing/broken, the run will fail fast with an error telling you what to install.
+
+## 🏗️ Installation
+
+```bash
+pip install jax-envelope
+```
+
+## 💞 Related projects
+
+- [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.2.0 → jax_envelope-0.4.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "jax-envelope"
-version = "0.2.0"
+version = "0.4.0"
 description = "A JAX-native environment interface with powerful wrappers and adapters for popular RL environment suites"
 readme = "README.md"
 requires-python = ">=3.12"
@@ -52,7 +52,7 @@ allow-direct-references = true
 packages = ["src/envelope"]
 
 [dependency-groups]
-compat = [
+adapters = [
     "brax>=0.13.0",
     "craftax>=1.4.3",
     "navix>=0.7.0",
@@ -80,5 +80,5 @@ kinetix-env = { git = "https://github.com/FLAIROx/Kinetix.git" }
 [tool.pytest.ini_options]
 testpaths = ["tests"]
 markers = [
-  "compat: tests requiring optional compat dependencies",
+  "adapters: tests requiring optional adapters dependencies",
 ]

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

@@ -1,4 +1,4 @@
-from envelope.compat import create
+from envelope.adapters 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

{jax_envelope-0.2.0/src/envelope/compat → jax_envelope-0.4.0/src/envelope/adapters}/__init__.py

@@ -4,14 +4,14 @@ 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"),
+    "gymnax": ("envelope.adapters.gymnax_envelope", "GymnaxEnvelope"),
+    "brax": ("envelope.adapters.brax_envelope", "BraxEnvelope"),
+    "navix": ("envelope.adapters.navix_envelope", "NavixEnvelope"),
+    "jumanji": ("envelope.adapters.jumanji_envelope", "JumanjiEnvelope"),
+    "kinetix": ("envelope.adapters.kinetix_envelope", "KinetixEnvelope"),
+    "craftax": ("envelope.adapters.craftax_envelope", "CraftaxEnvelope"),
     "mujoco_playground": (
-        "envelope.compat.mujoco_playground_envelope",
+        "envelope.adapters.mujoco_playground_envelope",
         "MujocoPlaygroundEnvelope",
     ),
 }

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

@@ -69,7 +69,7 @@ class BraxEnvelope(Environment):
         info = InfoContainer(
             obs=brax_state.obs,
             reward=brax_state.reward,
-            terminated=jnp.asarry(brax_state.done, dtype=bool).item(),
+            terminated=jnp.asarray(brax_state.done, dtype=bool),
         )
         info = info.update(**dataclasses.asdict(brax_state))
         return brax_state, info

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

@@ -10,7 +10,7 @@ from craftax.craftax_classic.envs.craftax_state import (
 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.adapters.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

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

@@ -81,7 +81,7 @@ class JumanjiEnvelope(Environment):
 
 
 def convert_jumanji_to_envelope_info(timestep: JumanjiTimeStep) -> InfoContainer:
-    term = jnp.asarray(timestep.last(), dtype=bool).item()
+    term = jnp.asarray(timestep.last(), dtype=bool)
     info = InfoContainer(
         obs=timestep.observation, reward=timestep.reward, terminated=term
     ).update(**timestep.extras)

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

@@ -1,7 +1,7 @@
 """Kinetix compatibility wrapper.
 
 This module exposes Kinetix environments through the `envelope.environment.Environment`
-API. It mirrors envelope's compat philosophy:
+API. It mirrors envelope's adapters philosophy:
 - prefer *no* environment-side auto-reset (use `AutoResetWrapper` in envelope)
 - prefer *no* fixed episode time-limits (use `TruncationWrapper` in envelope)
 
@@ -30,7 +30,7 @@ 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.adapters.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

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

@@ -42,7 +42,7 @@ class Environment(ABC, FrozenPyTreeNode):
 
     Two distinct lifecycle methods:
         init(key) - Initialize environment and all state from scratch.
-        reset(key, state) - Reset the inner environment while preserving
+        reset(state, key) - Reset the inner environment while preserving
             episode-persistent state.
     """
 
@@ -51,7 +51,7 @@ class Environment(ABC, FrozenPyTreeNode):
         """Initialize environment and all state from scratch."""
         ...
 
-    def reset(self, key: Key, state: State) -> tuple[State, Info]:
+    def reset(self, state: State, key: Key) -> tuple[State, Info]:
         """Reset the inner environment while preserving episode-persistent state."""
         return self.init(key)
 

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

@@ -1,6 +1,6 @@
 from abc import ABC, abstractmethod
 from functools import cached_property
-from typing import cast, override
+from typing import override
 
 import jax
 from jax import numpy as jnp

{jax_envelope-0.2.0 → jax_envelope-0.4.0}/src/envelope/wrappers/autoreset_wrapper.py

@@ -54,7 +54,7 @@ class AutoResetWrapper(Wrapper):
         return state, info.update(final=state.last_final)
 
     @override
-    def reset(self, key: Key, state: WrappedState) -> tuple[WrappedState, Info]:
+    def reset(self, state: WrappedState, key: Key) -> tuple[WrappedState, Info]:
         raise NotImplementedError("Reset is not implemented for AutoResetWrapper")
 
     @override
@@ -63,7 +63,7 @@ class AutoResetWrapper(Wrapper):
         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)
+        reset_inner_state, reset_info = self.env.reset(inner_state, key_reset)
 
         # Select next state and info based on done
         done = info.terminated | info.truncated

{jax_envelope-0.2.0 → jax_envelope-0.4.0}/src/envelope/wrappers/continuous_observation_wrapper.py

@@ -35,8 +35,8 @@ class ContinuousObservationWrapper(Wrapper):
         return state, info
 
     @override
-    def reset(self, key: Key, state: State) -> tuple[State, Info]:
-        state, info = self.env.reset(key, state)
+    def reset(self, state: State, key: Key) -> tuple[State, Info]:
+        state, info = self.env.reset(state, key)
         info = info.update(obs=to_float(info.obs))
         return state, info
 

{jax_envelope-0.2.0 → jax_envelope-0.4.0}/src/envelope/wrappers/episode_statistics_wrapper.py

@@ -24,8 +24,8 @@ class EpisodeStatisticsWrapper(Wrapper):
         return state, info.update(stats=state.stats)
 
     @override
-    def reset(self, key: Key, state: State) -> tuple[State, Info]:
-        inner_state, info = self.env.reset(key, state.inner_state)
+    def reset(self, state: State, key: Key) -> tuple[State, Info]:
+        inner_state, info = self.env.reset(state.inner_state, key)
         state = state.replace(inner_state=inner_state)
         return state, info.update(stats=state.stats)
 

{jax_envelope-0.2.0 → jax_envelope-0.4.0}/src/envelope/wrappers/flatten_observation_wrapper.py

@@ -37,10 +37,9 @@ class FlattenObservationWrapper(Wrapper):
         return state, info
 
     @override
-    def reset(self, key: Key, state: State) -> tuple[State, Info]:
-        state, info = self.env.reset(key, state)
-        info = info.update(obs=flatten_x(info.obs))
-        return state, info
+    def reset(self, state: State, key: Key) -> tuple[State, Info]:
+        next_state, info = self.env.reset(state, key)
+        return next_state, info.update(obs=flatten_x(info.obs))
 
     @override
     def step(self, state: State, action: PyTree) -> tuple[State, Info]:

{jax_envelope-0.2.0 → jax_envelope-0.4.0}/src/envelope/wrappers/observation_normalization_wrapper.py

@@ -76,8 +76,8 @@ class ObservationNormalizationWrapper(Wrapper):
         return self._normalize_and_update(next_state, info)
 
     @override
-    def reset(self, key: Key, state: WrappedState) -> tuple[WrappedState, Info]:
-        inner_state, info = self.env.reset(key, state.inner_state)
+    def reset(self, state: WrappedState, key: Key) -> tuple[WrappedState, Info]:
+        inner_state, info = self.env.reset(state.inner_state, key)
         # Preserve running statistics across resets
         next_state = self.ObservationNormalizationState(
             inner_state=inner_state, rmv_state=state.rmv_state

{jax_envelope-0.2.0 → jax_envelope-0.4.0}/src/envelope/wrappers/pooled_init_vmap_wrapper.py

@@ -36,7 +36,7 @@ class PooledInitVmapWrapper(Wrapper):
         return state, info.update(final=pholder_info)
 
     @override
-    def reset(self, key: Key, state: WrappedState) -> tuple[WrappedState, Info]:
+    def reset(self, state: WrappedState, key: Key) -> tuple[WrappedState, Info]:
         # It's hard to support reset for this wrapper.
         # We would have to init the state of a pool of unwrapped environments, and then
         # somehow inject this into the stack of wrapped states. The current data
@@ -48,7 +48,7 @@ class PooledInitVmapWrapper(Wrapper):
         # episodes before vmapping, we will implement this later.
         keys = _split_or_keep_key(key, self.batch_size + 1)
         key_next, keys_pool = keys[0], keys[1:]
-        inner_state, info = jax.vmap(self.env.reset)(keys_pool, state.inner_state)
+        inner_state, info = jax.vmap(self.env.reset)(state.inner_state, keys_pool)
         state = state.replace(inner_state=inner_state, init_key=key_next)
         return state, info.update(final=state.last_final)
 

{jax_envelope-0.2.0 → jax_envelope-0.4.0}/src/envelope/wrappers/state_injection_wrapper.py

@@ -69,13 +69,13 @@ class StateInjectionWrapper(Wrapper):
         return state, info
 
     @override
-    def reset(self, key: Key, state: WrappedState) -> tuple[WrappedState, Info]:
+    def reset(self, state: WrappedState, key: Key) -> tuple[WrappedState, Info]:
         # If reset state is set, use it instead of resetting inner env
         if state.reset_state is not None and state.reset_obs is not None:
             inner_state = state.reset_state
             info = InfoContainer(obs=state.reset_obs, reward=0.0, terminated=False)
         elif state.reset_state is None and state.reset_obs is None:
-            inner_state, info = self.env.reset(key, state.inner_state)
+            inner_state, info = self.env.reset(state.inner_state, key)
         else:
             raise ValueError("State must set both reset_state and reset_obs or neither")
 

{jax_envelope-0.2.0 → jax_envelope-0.4.0}/src/envelope/wrappers/truncation_wrapper.py

@@ -21,8 +21,8 @@ class TruncationWrapper(Wrapper):
         return state, info.update(truncated=self.max_steps <= 0)
 
     @override
-    def reset(self, key: Key, state: WrappedState) -> tuple[WrappedState, Info]:
-        inner_state, info = self.env.reset(key, state.inner_state)
+    def reset(self, state: WrappedState, key: Key) -> tuple[WrappedState, Info]:
+        inner_state, info = self.env.reset(state.inner_state, key)
         state = state.replace(inner_state=inner_state, steps=0)
         return state, info.update(truncated=self.max_steps <= 0)
 

{jax_envelope-0.2.0 → jax_envelope-0.4.0}/src/envelope/wrappers/vmap_envs_wrapper.py

@@ -40,11 +40,9 @@ class VmapEnvsWrapper(Wrapper):
         return state, info
 
     @override
-    def reset(self, key: Key, state: PyTree) -> tuple[WrappedState, Info]:
+    def reset(self, state: PyTree, key: Key) -> tuple[WrappedState, Info]:
         keys = self._split_keys(key)
-        state, info = jax.vmap(lambda e, k, s: e.reset(k, s))(
-            self.env, keys, state
-        )
+        state, info = jax.vmap(lambda e, s, k: e.reset(s, k))(self.env, state, keys)
         return state, info
 
     @override
@@ -58,7 +56,9 @@ class VmapEnvsWrapper(Wrapper):
     @property
     def observation_space(self) -> spaces.Space:
         env0 = _index_env(self.env, 0, self.batch_size)
-        return spaces.BatchedSpace(space=env0.observation_space, batch_size=self.batch_size)
+        return spaces.BatchedSpace(
+            space=env0.observation_space, batch_size=self.batch_size
+        )
 
     @override
     @cached_property

{jax_envelope-0.2.0 → jax_envelope-0.4.0}/src/envelope/wrappers/vmap_wrapper.py

@@ -41,9 +41,9 @@ class VmapWrapper(Wrapper):
         return state, info
 
     @override
-    def reset(self, key: Key, state: PyTree) -> tuple[WrappedState, Info]:
+    def reset(self, state: PyTree, key: Key) -> tuple[WrappedState, Info]:
         keys = _split_or_keep_key(key, self.batch_size)
-        state, info = jax.vmap(self.env.reset)(keys, state)
+        state, info = jax.vmap(self.env.reset)(state, keys)
         return state, info
 
     @override

{jax_envelope-0.2.0 → jax_envelope-0.4.0}/src/envelope/wrappers/wrapper.py

@@ -29,11 +29,11 @@ class Wrapper(Environment):
         return self.env.init(key)
 
     @override
-    def reset(self, key: Key, state: State) -> tuple[State, Info]:
-        return self.env.reset(key, state)
+    def reset(self, state: State, key: Key) -> tuple[State, Info]:
+        return self.env.reset(state, key)
 
     @override
-    def step(self, state: WrappedState, action: PyTree) -> tuple[WrappedState, Info]:
+    def step(self, state: State, action: PyTree) -> tuple[State, Info]:
         return self.env.step(state, action)
 
     @override

{jax_envelope-0.2.0/tests/compat → jax_envelope-0.4.0/tests/adapters}/conftest.py

@@ -1,6 +1,6 @@
 import pytest
 
-pytestmark = pytest.mark.compat
+pytestmark = pytest.mark.adapters
 
 
 @pytest.fixture(scope="module")

{jax_envelope-0.2.0/tests/compat → jax_envelope-0.4.0/tests/adapters}/contract.py

@@ -1,6 +1,6 @@
-"""Shared contract helpers for compat wrappers.
+"""Shared contract helpers for adapters.
 
-These functions enforce a consistent baseline across all compat wrappers:
+These functions enforce a consistent baseline across all adapters:
 - reset/step return (state, info) with Info fields present
 - reward is scalar-ish and finite
 - action sampling is valid for action_space

{jax_envelope-0.2.0/tests/compat → jax_envelope-0.4.0/tests/adapters}/test_brax_compat.py

@@ -1,4 +1,4 @@
-"""Tests for envelope.compat.brax_envelope module."""
+"""Tests for envelope.adapters.brax_envelope module."""
 
 # ruff: noqa: E402
 
@@ -7,14 +7,14 @@ from copy import deepcopy
 import jax
 import pytest
 
-pytestmark = pytest.mark.compat
+pytestmark = pytest.mark.adapters
 
 pytest.importorskip("brax")
 
 from brax.envs import Wrapper as BraxWrapper
 
-from envelope.compat.brax_envelope import BraxEnvelope
-from tests.compat.contract import (
+from envelope.adapters.brax_envelope import BraxEnvelope
+from tests.adapters.contract import (
     assert_jitted_rollout_contract,
     assert_reset_step_contract,
 )
@@ -93,8 +93,8 @@ def test_wrapper_unwrapping():
 
     # Create a simple wrapper
     class SimpleWrapper(BraxWrapper):
-        def reset(self, rng):
-            return self.env.reset(rng)
+        def init(self, rng):
+            return self.env.init(rng)
 
         def step(self, state, action):
             return self.env.step(state, action)

{jax_envelope-0.2.0/tests/compat → jax_envelope-0.4.0/tests/adapters}/test_craftax_compat.py

@@ -1,4 +1,4 @@
-"""Tests for envelope.compat.craftax_envelope module."""
+"""Tests for envelope.adapters.craftax_envelope module."""
 
 # ruff: noqa: E402
 
@@ -8,12 +8,12 @@ import jax
 import jax.numpy as jnp
 import pytest
 
-pytestmark = pytest.mark.compat
+pytestmark = pytest.mark.adapters
 
 pytest.importorskip("craftax")
 
 from envelope.spaces import Continuous, Discrete
-from tests.compat.contract import (
+from tests.adapters.contract import (
     assert_jitted_rollout_contract,
     assert_reset_step_contract,
 )
@@ -35,7 +35,7 @@ def craftax_env_id(request: pytest.FixtureRequest) -> str:
 
 @pytest.fixture(scope="module")
 def craftax_env(craftax_env_id: str):
-    from envelope.compat.craftax_envelope import CraftaxEnvelope
+    from envelope.adapters.craftax_envelope import CraftaxEnvelope
 
     return CraftaxEnvelope.from_name(craftax_env_id)
 
@@ -114,7 +114,7 @@ class _DummyEnv:
 
 
 def test_from_name_errors_on_auto_reset():
-    from envelope.compat.craftax_envelope import CraftaxEnvelope
+    from envelope.adapters.craftax_envelope import CraftaxEnvelope
 
     with pytest.raises(ValueError, match="Cannot override 'auto_reset' directly"):
         CraftaxEnvelope.from_name("AnyEnv", env_kwargs={"auto_reset": True})