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

{jax_envelope-0.1.1 → jax_envelope-0.3.0}/.gitignore

@@ -166,7 +166,8 @@ wandb/
 # ruff
 .ruff_cache/
 
-# Cursor
+# Vibecoding
 .cursor
 .cursorignore
 AGENTS.md
+CLAUDE.md

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: jax-envelope
-Version: 0.1.1
+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
-* [stoax](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.1.1 → 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
-* [stoax](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.1.1 → jax_envelope-0.3.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "jax-envelope"
-version = "0.1.1"
+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"
@@ -51,18 +51,15 @@ allow-direct-references = true
 [tool.hatch.build.targets.wheel]
 packages = ["src/envelope"]
 
-[tool.uv.sources]
-gymnax = { git = "https://github.com/RobertTLange/gymnax" }
-
 [dependency-groups]
 compat = [
-    "gymnax @ git+https://github.com/RobertTLange/gymnax@main",
     "brax>=0.13.0",
     "craftax>=1.4.3",
     "navix>=0.7.0",
     "jumanji>=1.0.1",
-    "kinetix-env>=2.0.0",
     "playground>=0.1.0",
+    "gymnax",
+    "kinetix-env",
 ]
 dev = [
     "hypothesis>=6.148.1",
@@ -71,7 +68,17 @@ dev = [
     "ruff>=0.14.2",
 ]
 
+[tool.uv]
+override-dependencies = [
+    "tensorflow-probability>=0.26.0.dev20260116",
+]
+
+[tool.uv.sources]
+gymnax = { git = "https://github.com/RobertTLange/gymnax.git" }
+kinetix-env = { git = "https://github.com/FLAIROx/Kinetix.git" }
+
 [tool.pytest.ini_options]
+testpaths = ["tests"]
 markers = [
   "compat: tests requiring optional compat dependencies",
 ]

{jax_envelope-0.1.1 → jax_envelope-0.3.0}/src/envelope/__init__.py

@@ -1,16 +1,22 @@
 from envelope.compat import create
 from envelope.environment import Environment, Info, InfoContainer
 from envelope.spaces import BatchedSpace, Continuous, Discrete, PyTreeSpace, Space
-from envelope.struct import field, static_field, FrozenPyTreeNode, Container
+from envelope.struct import Container, FrozenPyTreeNode, field, static_field
 from envelope.wrappers import (
-    Wrapper,
-    WrappedState,
     AutoResetWrapper,
+    ClipActionWrapper,
+    ContinuousObservationWrapper,
+    EpisodeStatisticsWrapper,
+    FlattenActionWrapper,
+    FlattenObservationWrapper,
     ObservationNormalizationWrapper,
+    PooledInitVmapWrapper,
     StateInjectionWrapper,
     TruncationWrapper,
-    VmapWrapper,
     VmapEnvsWrapper,
+    VmapWrapper,
+    WrappedState,
+    Wrapper,
 )
 
 __all__ = [
@@ -34,7 +40,13 @@ __all__ = [
     "Wrapper",
     "WrappedState",
     "AutoResetWrapper",
+    "ClipActionWrapper",
+    "ContinuousObservationWrapper",
+    "EpisodeStatisticsWrapper",
+    "FlattenActionWrapper",
+    "FlattenObservationWrapper",
     "ObservationNormalizationWrapper",
+    "PooledInitVmapWrapper",
     "StateInjectionWrapper",
     "TruncationWrapper",
     "VmapWrapper",

{jax_envelope-0.1.1 → jax_envelope-0.3.0}/src/envelope/compat/brax_envelope.py

@@ -48,7 +48,7 @@ class BraxEnvelope(Environment):
     def default_max_steps(self) -> int:
         return _BRAX_DEFAULT_EPISODE_LENGTH
 
-    def __post_init__(self) -> "BraxEnvelope":
+    def __post_init__(self):
         if isinstance(self.brax_env, BraxWrapper):
             warnings.warn(
                 "Environment wrapping should be handled by envelope. "
@@ -57,7 +57,7 @@ class BraxEnvelope(Environment):
             object.__setattr__(self, "brax_env", self.brax_env.unwrapped)
 
     @override
-    def reset(self, key: Key) -> tuple[State, Info]:
+    def init(self, key: Key) -> tuple[State, Info]:
         brax_state = self.brax_env.reset(key)
         info = InfoContainer(obs=brax_state.obs, reward=0.0, terminated=False)
         info = info.update(**dataclasses.asdict(brax_state))
@@ -67,7 +67,9 @@ class BraxEnvelope(Environment):
     def step(self, state: State, action: PyTree) -> tuple[State, Info]:
         brax_state = self.brax_env.step(state, action)
         info = InfoContainer(
-            obs=brax_state.obs, reward=brax_state.reward, terminated=brax_state.done
+            obs=brax_state.obs,
+            reward=brax_state.reward,
+            terminated=jnp.asarray(brax_state.done, dtype=bool),
         )
         info = info.update(**dataclasses.asdict(brax_state))
         return brax_state, info

{jax_envelope-0.1.1 → jax_envelope-0.3.0}/src/envelope/compat/craftax_envelope.py

@@ -22,7 +22,7 @@ class CraftaxEnvelope(Environment):
     """Wrapper to convert a Craftax environment to a envelope environment."""
 
     craftax_env: Any = static_field()
-    env_params: PyTree
+    env_params: PyTree = static_field()  # TODO: remove static marker as soon as craftax merges https://github.com/MichaelTMatthews/Craftax/pull/48
 
     @classmethod
     def from_name(
@@ -54,12 +54,27 @@ class CraftaxEnvelope(Environment):
     def default_max_steps(self) -> int:
         return int(self.craftax_env.default_params.max_timesteps)
 
+    @cached_property
+    def _craftax_info_placeholder(self) -> PyTree:
+        key = jax.random.PRNGKey(0)
+        _, state = self.craftax_env.reset(key, self.env_params)
+        _, _, _, _, info = self.craftax_env.step(
+            key,
+            state,
+            self.craftax_env.action_space(self.env_params).sample(key),
+            self.env_params,
+        )
+        return jax.tree.map(lambda x: jnp.full_like(x, jnp.nan), info)
+
     @override
-    def reset(self, key: Key) -> tuple[State, Info]:
+    def init(self, key: Key) -> tuple[State, Info]:
+        # TODO: this function does not add env_info (or comparable) to the info
+        # container. We should add tests for this (and all other envelopes) and fix it.
         key, subkey = jax.random.split(key)
         obs, env_state = self.craftax_env.reset(subkey, self.env_params)
         state = Container().update(key=key, env_state=env_state)
         info = InfoContainer(obs=obs, reward=0.0, terminated=False)
+        info = info.update(info=self._craftax_info_placeholder)
         return state, info
 
     @override

{jax_envelope-0.1.1 → jax_envelope-0.3.0}/src/envelope/compat/gymnax_envelope.py

@@ -1,5 +1,5 @@
 from functools import cached_property
-from typing import Any, override
+from typing import Any, Callable, cast, override
 
 import jax
 import jax.numpy as jnp
@@ -10,15 +10,24 @@ from gymnax.environments.environment import EnvParams as GymnaxEnvParams
 
 from envelope import spaces as envelope_spaces
 from envelope.environment import Environment, Info, InfoContainer, State
-from envelope.struct import Container, static_field
+from envelope.struct import Container, field, static_field
 from envelope.typing import Key, PyTree
 
+_GymnaxReset = Callable[
+    [Key, GymnaxEnvParams],
+    tuple[PyTree, Any],
+]
+_GymnaxStep = Callable[
+    [Key, Any, PyTree, GymnaxEnvParams],
+    tuple[PyTree, Any, jnp.ndarray, jnp.ndarray, PyTree],
+]
+
 
 class GymnaxEnvelope(Environment):
     """Wrapper to convert a Gymnax environment to a envelope environment."""
 
     gymnax_env: GymnaxEnv = static_field()
-    env_params: PyTree
+    env_params: PyTree = field()
 
     @classmethod
     def from_name(
@@ -43,19 +52,37 @@ class GymnaxEnvelope(Environment):
     def default_max_steps(self) -> int:
         return int(self.gymnax_env.default_params.max_steps_in_episode)
 
+    @cached_property
+    def _gymnax_info_placeholder(self) -> PyTree:
+        reset_fn = cast(_GymnaxReset, self.gymnax_env.reset)
+        step_fn = cast(_GymnaxStep, self.gymnax_env.step)
+
+        key = jax.random.PRNGKey(0)
+        _, state = reset_fn(key, self.env_params)
+        _, _, _, _, info = step_fn(
+            key,
+            state,
+            self.gymnax_env.action_space(self.env_params).sample(key),
+            self.env_params,
+        )
+        return jax.tree.map(lambda x: jnp.full_like(x, jnp.nan, dtype=float), info)
+
     @override
-    def reset(self, key: Key) -> tuple[State, Info]:
+    def init(self, key: Key) -> tuple[State, Info]:
+        reset_fn = cast(_GymnaxReset, self.gymnax_env.reset)
+
         key, subkey = jax.random.split(key)
-        obs, env_state = self.gymnax_env.reset(subkey, self.env_params)
+        obs, env_state = reset_fn(subkey, self.env_params)
         state = Container().update(key=key, env_state=env_state)
         info = InfoContainer(obs=obs, reward=0.0, terminated=False)
-        info = info.update(info=None)
+        info = info.update(info=self._gymnax_info_placeholder)
         return state, info
 
     @override
     def step(self, state: State, action: PyTree) -> tuple[State, Info]:
         key, subkey = jax.random.split(state.key)
-        obs, env_state, reward, done, env_info = self.gymnax_env.step(
+        step_fn = cast(_GymnaxStep, self.gymnax_env.step)
+        obs, env_state, reward, done, env_info = step_fn(
             subkey, state.env_state, action, self.env_params
         )
         state = state.update(key=key, env_state=env_state)

{jax_envelope-0.1.1 → jax_envelope-0.3.0}/src/envelope/compat/jumanji_envelope.py

@@ -48,7 +48,7 @@ class JumanjiEnvelope(Environment):
         return self._default_time_limit
 
     @override
-    def reset(self, key: Key) -> tuple[State, Info]:
+    def init(self, key: Key) -> tuple[State, Info]:
         env_state, timestep = self.jumanji_env.reset(key)
         info = convert_jumanji_to_envelope_info(timestep)
         return env_state, info
@@ -81,8 +81,9 @@ class JumanjiEnvelope(Environment):
 
 
 def convert_jumanji_to_envelope_info(timestep: JumanjiTimeStep) -> InfoContainer:
+    term = jnp.asarray(timestep.last(), dtype=bool)
     info = InfoContainer(
-        obs=timestep.observation, reward=timestep.reward, terminated=timestep.last()
+        obs=timestep.observation, reward=timestep.reward, terminated=term
     ).update(**timestep.extras)
     return info
 

{jax_envelope-0.1.1 → jax_envelope-0.3.0}/src/envelope/compat/kinetix_envelope.py

@@ -28,6 +28,7 @@ from kinetix.environment import (
 from kinetix.environment.ued.ued import make_reset_fn_sample_kinetix_level
 from kinetix.util.saving import load_from_json_file
 
+from envelope import field
 from envelope import spaces as envelope_spaces
 from envelope.compat.gymnax_envelope import _convert_space as _convert_gymnax_space
 from envelope.environment import Environment, Info, InfoContainer, State
@@ -67,7 +68,7 @@ class KinetixEnvelope(Environment):
     """Wrapper to convert a Kinetix environment to a envelope environment."""
 
     kinetix_env: Any = static_field()
-    env_params: Any
+    env_params: Any = field()
 
     @property
     def default_max_steps(self) -> int:
@@ -162,7 +163,7 @@ class KinetixEnvelope(Environment):
         return cls(kinetix_env=kinetix_env, env_params=env_params)
 
     @override
-    def reset(self, key: Key) -> tuple[State, Info]:
+    def init(self, key: Key) -> tuple[State, Info]:
         key, subkey = jax.random.split(key)
         obs, env_state = self.kinetix_env.reset(subkey, self.env_params)
         state_out = Container().update(key=key, env_state=env_state)

{jax_envelope-0.1.1 → jax_envelope-0.3.0}/src/envelope/compat/mujoco_playground_envelope.py

@@ -56,7 +56,7 @@ class MujocoPlaygroundEnvelope(Environment):
         return self._default_max_steps
 
     @override
-    def reset(self, key: Key) -> tuple[State, Info]:
+    def init(self, key: Key) -> tuple[State, Info]:
         env_state = self.mujoco_playground_env.reset(key)
         info = InfoContainer(obs=env_state.obs, reward=0.0, terminated=False)
         info = info.update(**dataclasses.asdict(env_state))

{jax_envelope-0.1.1 → jax_envelope-0.3.0}/src/envelope/compat/navix_envelope.py

@@ -38,7 +38,7 @@ class NavixEnvelope(Environment):
         return _NAVIX_DEFAULT_MAX_STEPS
 
     @override
-    def reset(self, key: Key) -> tuple[State, Info]:
+    def init(self, key: Key) -> tuple[State, Info]:
         timestep = self.navix_env.reset(key)
         return timestep, convert_navix_to_envelope_info(timestep)
 

{jax_envelope-0.1.1 → jax_envelope-0.3.0}/src/envelope/environment.py

@@ -5,7 +5,7 @@ from typing import Protocol, runtime_checkable
 
 from envelope import spaces
 from envelope.struct import Container, FrozenPyTreeNode
-from envelope.typing import Key, PyTree
+from envelope.typing import Array, Key, PyTree
 
 __all__ = ["Environment", "State", "Info", "InfoContainer"]
 
@@ -23,7 +23,7 @@ class Info(Protocol):
 
 class InfoContainer(Container):
     obs: PyTree
-    reward: float
+    reward: float | Array
     terminated: bool
     truncated: bool = field(default=False)
 
@@ -38,18 +38,25 @@ class Environment(ABC, FrozenPyTreeNode):
 
     State is an opaque PyTree owned by each environment; wrappers that stack
     environments should expose their wrapped env state as `inner_state` while
-    adding any wrapper-specific fields. `reset` may optionally receive a prior
-    state (for cross-episode persistence) and arbitrary **kwargs that wrappers
-    or environments can use.
+    adding any wrapper-specific fields.
+
+    Two distinct lifecycle methods:
+        init(key) - Initialize environment and all state from scratch.
+        reset(state, key) - Reset the inner environment while preserving
+            episode-persistent state.
     """
 
     @abstractmethod
-    def reset(
-        self, key: Key, state: State | None = None, **kwargs
-    ) -> tuple[State, Info]: ...
+    def init(self, key: Key) -> tuple[State, Info]:
+        """Initialize environment and all state from scratch."""
+        ...
+
+    def reset(self, state: State, key: Key) -> tuple[State, Info]:
+        """Reset the inner environment while preserving episode-persistent state."""
+        return self.init(key)
 
     @abstractmethod
-    def step(self, state: State, action: PyTree, **kwargs) -> tuple[State, Info]: ...
+    def step(self, state: State, action: PyTree) -> tuple[State, Info]: ...
 
     @abstractmethod
     @cached_property