jax-envelope 0.2.0__py3-none-any.whl → 0.3.0__py3-none-any.whl

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

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)

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)
 

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

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

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
 

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)
 

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]:

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

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)
 

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")
 

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)
 

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

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

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.dist-info → jax_envelope-0.3.0.dist-info}/METADATA

@@ -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.3.0.dist-info/RECORD

@@ -0,0 +1,32 @@
+envelope/__init__.py,sha256=AAk0b9UHGUTC4EzERXYPOrHuECgO0fx_UN4_gpKvgT0,1329
+envelope/environment.py,sha256=ZPTznHjMdpUkadDB0rpiGuG_E5ZhfgHt8SxfUhDgySU,2015
+envelope/spaces.py,sha256=E05OsM999nJw03lyg-ALXbvuw8xQT4AJhLFyahDKFaA,6938
+envelope/struct.py,sha256=tO1LLxk0bYc4kBCXraDv93T5d-xFRzpv4QCjASB0Lv4,5522
+envelope/typing.py,sha256=2rWKiZnKXaUU6JIl51Wcj_txRwe8wQNRSLSiRuT3OP4,127
+envelope/compat/__init__.py,sha256=6q2B2bfTIu7MlnIXc702ysqEnyURz-MEjTBMBCo8rdQ,3458
+envelope/compat/brax_envelope.py,sha256=xToV4fJYU9gPfF5SDs1Y3BaP96daz9iW8gKs2QYQDNA,3476
+envelope/compat/craftax_envelope.py,sha256=ey_mhJh5KCDRvjpCAJbFexpWiFjzQr7BZuRYLM9g-yA,3959
+envelope/compat/gymnax_envelope.py,sha256=X0pfFKaFKzvBl2Zih4utiXSEI-lc5f9miJGMxeT-Xfw,4690
+envelope/compat/jumanji_envelope.py,sha256=zl6yE1yqdCxZmjVCexu5vD10f5a-M82EpjzyVlj13wM,4909
+envelope/compat/kinetix_envelope.py,sha256=99RNN6nBN9CQEEe2fX2T1n8lL9ki0uKaz5BK052IU1o,6853
+envelope/compat/mujoco_playground_envelope.py,sha256=7iPaRuwXUHbP8hHNgSFipW8rJVruootMQPYBWNpPiHA,3705
+envelope/compat/navix_envelope.py,sha256=2n8vgr6gVeAYo12D1kDPxUrA9Hm-If3E7rV7TF_ksSA,3026
+envelope/wrappers/__init__.py,sha256=GFs54kY1aWbf9xymzm22x6M2AhZCwFLx1f5LF23jX94,1409
+envelope/wrappers/autoreset_wrapper.py,sha256=MIP9Svl1KyX0B7tfsSpcDJeEzPrAJQNWwpixrMZkF_s,3382
+envelope/wrappers/clip_action_wrapper.py,sha256=YKJbKBUbf-LADxIfCE6mMOKo_NEa2xPfDncNKg9iEHM,955
+envelope/wrappers/continuous_observation_wrapper.py,sha256=ud6MgMe7hY91q_26-A9t24RJSPrHI8wXue1LTdrm8O8,2041
+envelope/wrappers/episode_statistics_wrapper.py,sha256=F5NM6U0Z5kImSr7uflPb6nnZbtDtyqewAlryzqiWnwU,1433
+envelope/wrappers/flatten_action_wrapper.py,sha256=znR7WSebLy2wZnEWHi9SshxgV9b29khg-6ogQlmXq9U,2609
+envelope/wrappers/flatten_observation_wrapper.py,sha256=-i7mH09MoTwlA9xT7768n2Pb0TrVYSTLiUdoPKR76Gc,2829
+envelope/wrappers/normalization.py,sha256=YkNEZqvFSaPurBA-WfXDynV51552ucNkY0JqNSAKNUE,1819
+envelope/wrappers/observation_normalization_wrapper.py,sha256=crCm4FREqUumFnNdr30-cORLp4rIOTiyKywMzvQz0Xk,4921
+envelope/wrappers/pooled_init_vmap_wrapper.py,sha256=foqZBx-7u-n8VMidF_SoGFY0-Y-RFKB9v-Nj7ZTtmqk,4985
+envelope/wrappers/state_injection_wrapper.py,sha256=2Lr_LG2-0EF1q2KP7yQ954xfrWTU9moSC4bB1VLK3Vs,3553
+envelope/wrappers/truncation_wrapper.py,sha256=k3NZBMmnYiCN35jA1Yyi3rGdLcqM52CzuPabtvwGxe8,1345
+envelope/wrappers/vmap_envs_wrapper.py,sha256=b-xtzhVcT7P0IV4dvVvQk-rX5zv9H4Msv3Ftf5KnZp4,2757
+envelope/wrappers/vmap_wrapper.py,sha256=8Qxml26HZqjdKV7ey43wLszjMUxURQUFBHwJ1z5F9y8,1973
+envelope/wrappers/wrapper.py,sha256=SSnKFs_9Rhf0zorMUTTSlzj81OwR7KWo4PgCgycROwo,1503
+jax_envelope-0.3.0.dist-info/METADATA,sha256=I33XYdMzK_nJwPegfLUUU3TBusSw4oXy11_h3F3B_Is,5044
+jax_envelope-0.3.0.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+jax_envelope-0.3.0.dist-info/licenses/LICENSE,sha256=VyF-MK-gY2_fZlhf8uEnE2y8ziIXK-w55GM12eOgXrQ,1069
+jax_envelope-0.3.0.dist-info/RECORD,,

jax_envelope-0.2.0.dist-info/RECORD

@@ -1,32 +0,0 @@
-envelope/__init__.py,sha256=AAk0b9UHGUTC4EzERXYPOrHuECgO0fx_UN4_gpKvgT0,1329
-envelope/environment.py,sha256=OC15ooivLvD8EpmgRQGL-UjRpF0UXO0fOGu61I8HmRs,2015
-envelope/spaces.py,sha256=Nb8T11ilCtqKPSLZ98oqhnuw_8JNLOGCZdsy2X_pSwU,6944
-envelope/struct.py,sha256=tO1LLxk0bYc4kBCXraDv93T5d-xFRzpv4QCjASB0Lv4,5522
-envelope/typing.py,sha256=2rWKiZnKXaUU6JIl51Wcj_txRwe8wQNRSLSiRuT3OP4,127
-envelope/compat/__init__.py,sha256=6q2B2bfTIu7MlnIXc702ysqEnyURz-MEjTBMBCo8rdQ,3458
-envelope/compat/brax_envelope.py,sha256=XV07CYApXmf5g-Bin4pKWqExTZ7gd6AO2BCx3bk1Z1c,3482
-envelope/compat/craftax_envelope.py,sha256=ey_mhJh5KCDRvjpCAJbFexpWiFjzQr7BZuRYLM9g-yA,3959
-envelope/compat/gymnax_envelope.py,sha256=X0pfFKaFKzvBl2Zih4utiXSEI-lc5f9miJGMxeT-Xfw,4690
-envelope/compat/jumanji_envelope.py,sha256=WhU4WJLW2Fn4BG1_69EJO9ZiBDWEaY387dhkRZjqcGc,4916
-envelope/compat/kinetix_envelope.py,sha256=99RNN6nBN9CQEEe2fX2T1n8lL9ki0uKaz5BK052IU1o,6853
-envelope/compat/mujoco_playground_envelope.py,sha256=7iPaRuwXUHbP8hHNgSFipW8rJVruootMQPYBWNpPiHA,3705
-envelope/compat/navix_envelope.py,sha256=2n8vgr6gVeAYo12D1kDPxUrA9Hm-If3E7rV7TF_ksSA,3026
-envelope/wrappers/__init__.py,sha256=GFs54kY1aWbf9xymzm22x6M2AhZCwFLx1f5LF23jX94,1409
-envelope/wrappers/autoreset_wrapper.py,sha256=MJlbOj94_6I1KQ4oBNsKEB416YWp9ZMiqn85b1PdhKQ,3382
-envelope/wrappers/clip_action_wrapper.py,sha256=YKJbKBUbf-LADxIfCE6mMOKo_NEa2xPfDncNKg9iEHM,955
-envelope/wrappers/continuous_observation_wrapper.py,sha256=geDOeQDuASiRDEv5U3fF7GxblfB0Ku99GvTUGZSteJs,2041
-envelope/wrappers/episode_statistics_wrapper.py,sha256=-LzypX4aYS9mCkl33M2mf3aT3OhXWtDc_kVU6VlSJ0o,1433
-envelope/wrappers/flatten_action_wrapper.py,sha256=znR7WSebLy2wZnEWHi9SshxgV9b29khg-6ogQlmXq9U,2609
-envelope/wrappers/flatten_observation_wrapper.py,sha256=xgoToXdJ_6M7cwfL7pu91ZnnFh3LHPR_DSe7XJwkh2Q,2839
-envelope/wrappers/normalization.py,sha256=YkNEZqvFSaPurBA-WfXDynV51552ucNkY0JqNSAKNUE,1819
-envelope/wrappers/observation_normalization_wrapper.py,sha256=pAwyd01lAxvEu1Z4KZfhSRoVNeLpeohtw1t1wsE53w4,4921
-envelope/wrappers/pooled_init_vmap_wrapper.py,sha256=LAgrEjLXlq-gJiCGzI357a41tx1mVgb1ZQwtgN8mTuw,4985
-envelope/wrappers/state_injection_wrapper.py,sha256=L8f9yrlSigk0jNx6D6wSu0cR0bgssq2uJNd9OYPAY0E,3553
-envelope/wrappers/truncation_wrapper.py,sha256=-CSoXs-mDy1AiJVKLUc2zBHQxalhYjjyAAu1NmiIUsE,1345
-envelope/wrappers/vmap_envs_wrapper.py,sha256=gE_TAmfsu2PJlDGWTz5xENJBJ-6nSJpTUlKhOnmZPi0,2757
-envelope/wrappers/vmap_wrapper.py,sha256=B43lFNdAN_d8LgN6Rz6Ty85r9zcXLB3Wrn8R-Z2qHps,1973
-envelope/wrappers/wrapper.py,sha256=tWluyIg4eGl60R2v3qaahU5dlmKKkBsJiNj7GnWR06E,1517
-jax_envelope-0.2.0.dist-info/METADATA,sha256=jCWoAuoFsrLJakIhecRdz1zOIGCdrD4dBS0JSA2LVvw,4651
-jax_envelope-0.2.0.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-jax_envelope-0.2.0.dist-info/licenses/LICENSE,sha256=VyF-MK-gY2_fZlhf8uEnE2y8ziIXK-w55GM12eOgXrQ,1069
-jax_envelope-0.2.0.dist-info/RECORD,,