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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: jax-envelope
-Version: 0.2.0
+Version: 0.3.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,35 +42,42 @@ 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`
@@ -77,11 +85,14 @@ let's you create environments from any of the above!
   - If any compat 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.2.0 → jax_envelope-0.3.0}/README.md

@@ -1,10 +1,11 @@
 # 💌 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())
 
@@ -15,35 +16,42 @@ 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`
@@ -51,11 +59,14 @@ let's you create environments from any of the above!
   - If any compat 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.2.0 → jax_envelope-0.3.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "jax-envelope"
-version = "0.2.0"
+version = "0.3.0"
 description = "A JAX-native environment interface with powerful wrappers and adapters for popular RL environment suites"
 readme = "README.md"
 requires-python = ">=3.12"

{jax_envelope-0.2.0 → jax_envelope-0.3.0}/src/envelope/compat/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 → jax_envelope-0.3.0}/src/envelope/compat/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 → jax_envelope-0.3.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.3.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.3.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.3.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.3.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.3.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.3.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.3.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.3.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.3.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.3.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.3.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.3.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 → jax_envelope-0.3.0}/tests/compat/test_brax_compat.py

@@ -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 → jax_envelope-0.3.0}/tests/wrappers/helpers.py

@@ -98,7 +98,7 @@ class StepCounterEnv(Environment):
             truncated=truncated,
         )
 
-    def reset(self, key: Key, state: State) -> tuple[StepState, InfoContainer]:
+    def reset(self, state: State, key: Key) -> tuple[StepState, InfoContainer]:
         return self.init(key)
 
     def step(
@@ -198,9 +198,7 @@ class NoStepsEnv(Environment):
             obs=s.env_state, reward=0.0, terminated=False, truncated=False
         )
 
-    def reset(
-        self, key: Key, state: State
-    ) -> tuple[NoStepsState, InfoContainer]:
+    def reset(self, state: State, key: Key) -> tuple[NoStepsState, InfoContainer]:
         return self.init(key)
 
     def step(
@@ -230,9 +228,7 @@ class AlternatingTerminationEnv(Environment):
             obs=s.env_state, reward=0.0, terminated=False, truncated=False
         )
 
-    def reset(
-        self, key: Key, state: State
-    ) -> tuple[StepState, InfoContainer]:
+    def reset(self, state: State, key: Key) -> tuple[StepState, InfoContainer]:
         return self.init(key)
 
     def step(
@@ -266,7 +262,7 @@ class ScalarToyEnv(Environment):
         s = jnp.asarray(0.0, dtype=jnp.float32)
         return s, InfoContainer(obs=s, reward=0.0, terminated=False, truncated=False)
 
-    def reset(self, key: Key, state: State) -> tuple[State, Info]:
+    def reset(self, state: State, key: Key) -> tuple[State, Info]:
         return self.init(key)
 
     def step(self, state: State, action: jax.Array) -> tuple[State, Info]:
@@ -300,7 +296,7 @@ class VectorToyEnv(Environment):
         s = jnp.zeros((self.dim,), dtype=jnp.float32)
         return s, InfoContainer(obs=s, reward=0.0, terminated=False, truncated=False)
 
-    def reset(self, key: Key, state: State) -> tuple[State, Info]:
+    def reset(self, state: State, key: Key) -> tuple[State, Info]:
         return self.init(key)
 
     def step(self, state: State, action: jax.Array) -> tuple[State, Info]:
@@ -330,7 +326,7 @@ class FlagDoneEnv(Environment):
         z = jnp.array(0.0)
         return z, InfoContainer(obs=z, reward=0.0, terminated=False, truncated=False)
 
-    def reset(self, key: Key, state: State):
+    def reset(self, state: State, key: Key):
         return self.init(key)
 
     def step(self, state: State, action: jax.Array):
@@ -366,7 +362,7 @@ class ParamEnv(Environment):
         s = jnp.asarray([self.offset, -self.offset], dtype=jnp.float32)
         return s, InfoContainer(obs=s, reward=0.0, terminated=False, truncated=False)
 
-    def reset(self, key: Key, state: State) -> tuple[State, Info]:
+    def reset(self, state: State, key: Key) -> tuple[State, Info]:
         return self.init(key)
 
     def step(self, state: State, action: jax.Array) -> tuple[State, Info]:
@@ -401,7 +397,7 @@ class VectorObsEnv(Environment):
         s = jnp.linspace(0.0, 1.0, self.dim, dtype=jnp.float32)
         return s, InfoContainer(obs=s, reward=0.0, terminated=False, truncated=False)
 
-    def reset(self, key: Key, state: State):
+    def reset(self, state: State, key: Key):
         return self.init(key)
 
     def step(self, state: State, action: jax.Array):
@@ -444,7 +440,7 @@ class PyTreeObsEnv(Environment):
         s = obs
         return s, InfoContainer(obs=obs, reward=0.0, terminated=False, truncated=False)
 
-    def reset(self, key: Key, state: State):
+    def reset(self, state: State, key: Key):
         return self.init(key)
 
     def step(self, state: State, action: jax.Array):
@@ -476,7 +472,7 @@ class ConstantObsEnv(Environment):
         obs = jnp.asarray(self.value, self.dtype) * jnp.ones(self.shape, self.dtype)
         return 0, InfoContainer(obs=obs, reward=0.0, terminated=False, truncated=False)
 
-    def reset(self, key: Key, state: State):
+    def reset(self, state: State, key: Key):
         return self.init(key)
 
     def step(self, state: State, action: jax.Array):
@@ -500,10 +496,12 @@ class PyTreeActionEnv(Environment):
 
     @cached_property
     def action_space(self) -> PyTreeSpace:
-        return PyTreeSpace({
-            "a": Continuous.from_shape(low=-1.0, high=1.0, shape=(2,)),
-            "b": Continuous.from_shape(low=-1.0, high=1.0, shape=(3,)),
-        })
+        return PyTreeSpace(
+            {
+                "a": Continuous.from_shape(low=-1.0, high=1.0, shape=(2,)),
+                "b": Continuous.from_shape(low=-1.0, high=1.0, shape=(3,)),
+            }
+        )
 
     def _action_to_vec(self, action: PyTree) -> jax.Array:
         leaves = jax.tree.leaves(action)
@@ -513,12 +511,10 @@ class PyTreeActionEnv(Environment):
         s = jnp.zeros(5, dtype=jnp.float32)
         return s, InfoContainer(obs=s, reward=0.0, terminated=False, truncated=False)
 
-    def reset(self, key: Key, state: State) -> tuple[jax.Array, InfoContainer]:
+    def reset(self, state: State, key: Key) -> tuple[jax.Array, InfoContainer]:
         return self.init(key)
 
-    def step(
-        self, state: jax.Array, action: PyTree
-    ) -> tuple[jax.Array, InfoContainer]:
+    def step(self, state: jax.Array, action: PyTree) -> tuple[jax.Array, InfoContainer]:
         vec = self._action_to_vec(action)
         ns = state + jnp.asarray(vec, dtype=jnp.float32)
         reward = jnp.sum(vec)
@@ -545,7 +541,7 @@ class IntObsEnv(Environment):
         s = jnp.array(0, dtype=jnp.int32)
         return s, InfoContainer(obs=s, reward=0.0, terminated=False, truncated=False)
 
-    def reset(self, key: Key, state: State):
+    def reset(self, state: State, key: Key):
         return self.init(key)
 
     def step(self, state: State, action: jax.Array):
@@ -581,7 +577,7 @@ class RandomImageEnv(Environment):
             obs=obs.astype(self.dtype), reward=0.0, terminated=False, truncated=False
         )
 
-    def reset(self, key: Key, state: State):
+    def reset(self, state: State, key: Key):
         return self.init(key)
 
     def step(self, state: State, action: jax.Array):
@@ -618,9 +614,7 @@ class WrapperSimpleEnv(Environment):
         info = TestInfo(obs=state, reward=0.0, terminated=False, truncated=False)
         return state, info
 
-    def reset(
-        self, key: Key, state: State
-    ) -> tuple[jax.Array, TestInfo]:
+    def reset(self, state: State, key: Key) -> tuple[jax.Array, TestInfo]:
         return self.init(key)
 
     def step(self, state: jax.Array, action: jax.Array) -> tuple[jax.Array, TestInfo]:
@@ -650,9 +644,7 @@ class WrapperEnvWithFields(Environment):
         info = TestInfo(obs=state, reward=0.0, terminated=False, truncated=False)
         return state, info
 
-    def reset(
-        self, key: Key, state: State
-    ) -> tuple[jax.Array, TestInfo]:
+    def reset(self, state: State, key: Key) -> tuple[jax.Array, TestInfo]:
         return self.init(key)
 
     def step(self, state: jax.Array, action: jax.Array) -> tuple[jax.Array, TestInfo]:
@@ -679,9 +671,7 @@ class WrapperEnvWithMethods(Environment):
         info = TestInfo(obs=state, reward=0.0, terminated=False, truncated=False)
         return state, info
 
-    def reset(
-        self, key: Key, state: State
-    ) -> tuple[jax.Array, TestInfo]:
+    def reset(self, state: State, key: Key) -> tuple[jax.Array, TestInfo]:
         return self.init(key)
 
     def step(self, state: jax.Array, action: jax.Array) -> tuple[jax.Array, TestInfo]:
@@ -727,7 +717,7 @@ def make_wrapper_discrete_env() -> Environment:
             info = TestInfo(obs=state, reward=0.0, terminated=False, truncated=False)
             return state, info
 
-        def reset(self, key: Key, state: State):
+        def reset(self, state: State, key: Key):
             return self.init(key)
 
         def step(self, state: jax.Array, action: jax.Array):
@@ -762,7 +752,7 @@ def make_wrapper_complex_state_env() -> Environment:
             )
             return st, info
 
-        def reset(self, key: Key, state: State):
+        def reset(self, state: State, key: Key):
             return self.init(key)
 
         def step(self, state: dict, action: jax.Array):

{jax_envelope-0.2.0 → jax_envelope-0.3.0}/tests/wrappers/test_autoreset_wrapper.py

@@ -544,8 +544,8 @@ def test_auto_reset_passes_state_to_inner_wrapper():
                 received_state_on_reset=False,
             ), info
 
-        def reset(self, key, state):
-            inner_state, info = self.env.reset(key, state.inner_state)
+        def reset(self, state, key):
+            inner_state, info = self.env.reset(state.inner_state, key)
             return self.TrackingState(
                 inner_state=inner_state,
                 received_state_on_reset=True,

{jax_envelope-0.2.0 → jax_envelope-0.3.0}/tests/wrappers/test_clip_action_wrapper.py

@@ -21,8 +21,8 @@ def test_init_reset_delegate_unchanged():
     assert jnp.allclose(state_w, state_e)
     assert jnp.allclose(info_w.obs, info_e.obs)
 
-    state_w, info_w = w.reset(key, state_w)
-    state_e, info_e = env.reset(key, state_e)
+    state_w, info_w = w.reset(state_w, key)
+    state_e, info_e = env.reset(state_e, key)
     assert jnp.allclose(state_w, state_e)
     assert jnp.allclose(info_w.obs, info_e.obs)
     assert w.observation_space.contains(info_w.obs)

{jax_envelope-0.2.0 → jax_envelope-0.3.0}/tests/wrappers/test_continuous_observation_wrapper.py

@@ -17,7 +17,7 @@ def test_init_reset_step_cast_discrete_obs_to_float32():
     key = jax.random.PRNGKey(0)
     state, info = w.init(key)
     assert info.obs.dtype == jnp.float32
-    state, info = w.reset(key, state)
+    state, info = w.reset(state, key)
     assert info.obs.dtype == jnp.float32
     assert w.observation_space.contains(info.obs)
     state, info = w.step(state, jnp.array(0, dtype=jnp.int32))

{jax_envelope-0.2.0 → jax_envelope-0.3.0}/tests/wrappers/test_episode_statistics_wrapper.py

@@ -66,7 +66,7 @@ def test_reset_preserves_stats():
         state, _ = w.step(state, jnp.asarray(0.2))
     reward_before = state.stats.reward
     length_before = state.stats.length
-    state, info = w.reset(key, state)
+    state, info = w.reset(state, key)
     assert jnp.allclose(state.stats.reward, reward_before)
     assert jnp.allclose(state.stats.length, length_before)
     assert jnp.allclose(info.stats.reward, reward_before)
@@ -81,7 +81,7 @@ def test_stats_persist_and_continue_after_reset():
     state, _ = w.init(key)
     for _ in range(3):
         state, _ = w.step(state, jnp.asarray(0.1))
-    state, _ = w.reset(key, state)
+    state, _ = w.reset(state, key)
     for _ in range(2):
         state, _ = w.step(state, jnp.asarray(0.1))
     # Total length = 3 + 2 = 5, reward = 0.1*5 = 0.5

{jax_envelope-0.2.0 → jax_envelope-0.3.0}/tests/wrappers/test_flatten_action_wrapper.py

@@ -60,8 +60,8 @@ def test_init_reset_delegate_unchanged():
     state_e, info_e = env.init(key)
     assert jnp.allclose(state_w, state_e)
     assert jnp.allclose(info_w.obs, info_e.obs)
-    state_w, info_w = w.reset(key, state_w)
-    state_e, info_e = env.reset(key, state_e)
+    state_w, info_w = w.reset(state_w, key)
+    state_e, info_e = env.reset(state_e, key)
     assert jnp.allclose(state_w, state_e)
     assert jnp.allclose(info_w.obs, info_e.obs)
 
@@ -101,7 +101,7 @@ def test_action_space_flattened_discrete():
                 obs=s, reward=0.0, terminated=False, truncated=False
             )
 
-        def reset(self, key, state):
+        def reset(self, state, key):
             return self.init(key)
 
         def step(self, state, action):
@@ -161,7 +161,7 @@ def test_mixed_space_types_raises_value_error():
                 obs=s, reward=0.0, terminated=False, truncated=False
             )
 
-        def reset(self, key, state):
+        def reset(self, state, key):
             return self.init(key)
 
         def step(self, state, action):

{jax_envelope-0.2.0 → jax_envelope-0.3.0}/tests/wrappers/test_flatten_observation_wrapper.py

@@ -29,7 +29,7 @@ def test_reset_step_flatten_pytree_obs():
     key = jax.random.PRNGKey(0)
     state, info = w.init(key)
     assert info.obs.shape == (5,)
-    state, info = w.reset(key, state)
+    state, info = w.reset(state, key)
     assert info.obs.shape == (5,)
     assert w.observation_space.contains(info.obs)
     state, info = w.step(state, jnp.array(0.0))
@@ -126,7 +126,7 @@ def test_mixed_space_types_raises_value_error():
                 obs=obs, reward=0.0, terminated=False, truncated=False
             )
 
-        def reset(self, key, state):
+        def reset(self, state, key):
             return self.init(key)
 
         def step(self, state, action):

{jax_envelope-0.2.0 → jax_envelope-0.3.0}/tests/wrappers/test_pooled_init_vmap_wrapper.py

@@ -108,7 +108,7 @@ def test_reset_vmaps_inner_reset():
     w = PooledInitVmapWrapper(env=env, batch_size=batch_size, pool_size=3)
     key = jax.random.PRNGKey(0)
     state, info = w.init(key)
-    state, info = w.reset(key, state)
+    state, info = w.reset(state, key)
     assert info.obs.shape == (batch_size,)
     assert w.observation_space.contains(info.obs)
 

{jax_envelope-0.2.0 → jax_envelope-0.3.0}/tests/wrappers/test_state_injection_wrapper.py

@@ -66,7 +66,7 @@ class TestStateInjectionCoreFunctionality:
 
         # Reset again, passing the current state (simulates auto-reset)
         key2 = jax.random.PRNGKey(1)
-        state2, info2 = w.reset(key2, state)
+        state2, info2 = w.reset(state, key2)
 
         # Should preserve the injected state
         assert jnp.allclose(state2.reset_state.env_state, jnp.array(42.0))
@@ -132,7 +132,7 @@ class TestStateInjectionCoreFunctionality:
 
         # Reset with this state (no reset_state set) - should delegate to inner env
         key2 = jax.random.PRNGKey(1)
-        state2, info2 = w.reset(key2, state)
+        state2, info2 = w.reset(state, key2)
 
         # Should have done a normal reset - inner_state is fresh from env
         assert jnp.allclose(state2.inner_state.env_state, jnp.array(0.0))
@@ -166,7 +166,7 @@ class TestStateInjectionCoreFunctionality:
         )
 
         with pytest.raises(ValueError, match="must set both"):
-            w.reset(key, state_with_only_reset_state)
+            w.reset(state_with_only_reset_state, key)
 
         # Create state with only reset_obs set (not reset_state)
         state_with_only_reset_obs = w.InjectedState(
@@ -176,7 +176,7 @@ class TestStateInjectionCoreFunctionality:
         )
 
         with pytest.raises(ValueError, match="must set both"):
-            w.reset(key, state_with_only_reset_obs)
+            w.reset(state_with_only_reset_obs, key)
 
 
 # ============================================================================

{jax_envelope-0.2.0 → jax_envelope-0.3.0}/tests/wrappers/test_truncation_wrapper.py

@@ -116,7 +116,7 @@ def test_steps_as_jax_scalar_array_behaves_correctly():
 
 
 def test_reset_with_state_passes_inner_state_down():
-    """reset(key, state) should pass state.inner_state to the inner env's reset."""
+    """reset(state, key) should pass state.inner_state to the inner env's reset."""
     env = StepCounterEnv()
     w = TruncationWrapper(env=env, max_steps=10)
     key = jax.random.PRNGKey(0)
@@ -126,7 +126,7 @@ def test_reset_with_state_passes_inner_state_down():
         state, _ = w.step(state, jnp.asarray(0.1))
     assert state.steps == 5
 
-    new_state, _ = w.reset(jax.random.PRNGKey(1), state)
+    new_state, _ = w.reset(state, jax.random.PRNGKey(1))
 
     # Inner env should be reset
     assert jnp.allclose(new_state.inner_state.env_state, 0.0)

{jax_envelope-0.2.0 → jax_envelope-0.3.0}/uv.lock

@@ -1056,7 +1056,7 @@ wheels = [
 
 [[package]]
 name = "jax-envelope"
-version = "0.2.0"
+version = "0.3.0"
 source = { editable = "." }
 dependencies = [
     { name = "jax" },