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

envelope/wrappers/vmap_wrapper.py

@@ -2,50 +2,65 @@ from functools import cached_property
 from typing import override
 
 import jax
+import jax.numpy as jnp
 
 from envelope import spaces
 from envelope.environment import Info
-from envelope.struct import field
+from envelope.struct import static_field
 from envelope.typing import Key, PyTree
 from envelope.wrappers.wrapper import WrappedState, Wrapper
 
 
+def is_single_key(key):
+    # New-style typed keys have dtype like key<fry>
+    if jnp.issubdtype(key.dtype, jax.dtypes.prng_key):
+        return key.ndim == 0
+    return key.shape == (2,)
+
+
+def _split_or_keep_key(key: Key, batch_size: int) -> Key:
+    if is_single_key(key):
+        return jax.random.split(key, batch_size)
+    elif key.shape[0] == batch_size:
+        return key
+    raise ValueError(
+        f"reset key's leading dimension ({key.shape[0]}) must match "
+        f"batch_size ({batch_size})."
+    )
+
+
 class VmapWrapper(Wrapper):
-    """Does not forward kwargs to the underlying env. Does not wrap the state."""
+    """Does not wrap the state."""
 
-    batch_size: int = field(kw_only=True)
+    batch_size: int = static_field(kw_only=True)
+
+    @override
+    def init(self, key: Key) -> tuple[WrappedState, Info]:
+        keys = _split_or_keep_key(key, self.batch_size)
+        state, info = jax.vmap(self.env.init)(keys)
+        return state, info
 
     @override
-    def reset(
-        self, key: Key, state: PyTree | None = None, **kwargs
-    ) -> tuple[WrappedState, Info]:
-        # Accept single key or batched keys
-        if key.shape == (2,):
-            keys = jax.random.split(key, self.batch_size)
-        else:
-            if key.shape[0] != self.batch_size:
-                raise ValueError(
-                    f"reset key's leading dimension ({key.shape[0]}) must match "
-                    f"batch_size ({self.batch_size})."
-                )
-            keys = key
-
-        state, info = jax.vmap(self.env.reset)(keys, state)
+    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)(state, keys)
         return state, info
 
     @override
-    def step(
-        self, state: WrappedState, action: PyTree, **kwargs
-    ) -> tuple[WrappedState, Info]:
+    def step(self, state: WrappedState, action: PyTree) -> tuple[WrappedState, Info]:
         state, info = jax.vmap(self.env.step)(state, action)
         return state, info
 
     @override
     @cached_property
     def observation_space(self) -> spaces.Space:
-        return spaces.batch_space(self.env.observation_space, self.batch_size)
+        return spaces.BatchedSpace(
+            space=self.env.observation_space, batch_size=self.batch_size
+        )
 
     @override
     @cached_property
     def action_space(self) -> spaces.Space:
-        return spaces.batch_space(self.env.action_space, self.batch_size)
+        return spaces.BatchedSpace(
+            space=self.env.action_space, batch_size=self.batch_size
+        )

envelope/wrappers/wrapper.py

@@ -25,16 +25,16 @@ class Wrapper(Environment):
     env: Environment = field(kw_only=True)
 
     @override
-    def reset(
-        self, key: Key, state: State | None = None, **kwargs
-    ) -> tuple[State, Info]:
-        return self.env.reset(key, state=state, **kwargs)
+    def init(self, key: Key) -> tuple[State, Info]:
+        return self.env.init(key)
 
     @override
-    def step(
-        self, state: WrappedState, action: PyTree, **kwargs
-    ) -> tuple[WrappedState, Info]:
-        return self.env.step(state, action, **kwargs)
+    def reset(self, state: State, key: Key) -> tuple[State, Info]:
+        return self.env.reset(state, key)
+
+    @override
+    def step(self, state: State, action: PyTree) -> tuple[State, Info]:
+        return self.env.step(state, action)
 
     @override
     @cached_property

{jax_envelope-0.1.1.dist-info → jax_envelope-0.3.0.dist-info}/METADATA

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

@@ -1,27 +0,0 @@
-envelope/__init__.py,sha256=x4DtJ3WWPsPde1wSrXAXCdJVuXV3pioAojrrOvpJNqw,975
-envelope/environment.py,sha256=4T0IosoRX-6gVvZCKBUEitlM_5cAT43hrKQ0Hr53XGk,1735
-envelope/spaces.py,sha256=fvJ3-ZI3iyFlpfEN1barTd9cbWvQ11nl-7Y45KgFizs,6433
-envelope/struct.py,sha256=Sb8GsxZ7rFF5A5128oZWEQVnFK79xUFe6EaU2bghBOw,5158
-envelope/typing.py,sha256=2rWKiZnKXaUU6JIl51Wcj_txRwe8wQNRSLSiRuT3OP4,127
-envelope/compat/__init__.py,sha256=6q2B2bfTIu7MlnIXc702ysqEnyURz-MEjTBMBCo8rdQ,3458
-envelope/compat/brax_envelope.py,sha256=T0Xefrpgft16SdRJKMIjzdJq5tbocQ4jeI763PDzZ_4,3445
-envelope/compat/craftax_envelope.py,sha256=t151fDCrgPWfynERXHYXiWBwCGAXbFcreuIyH7GCApg,3162
-envelope/compat/gymnax_envelope.py,sha256=iau4CKFLVIJkRA9WILgOqkp3mXIAzkilKcrKF8vkld4,3790
-envelope/compat/jumanji_envelope.py,sha256=_Tox7DfiyE-Z4drszPUe-arPurz6hMeFtRNfy8o9Ipc,4869
-envelope/compat/kinetix_envelope.py,sha256=kPLWwSoQ49kNXuvfbStOtoGby9EkV7bKAetxUa1_YPA,6817
-envelope/compat/mujoco_playground_envelope.py,sha256=UcYZwOPD27473wmDDb5n2eK46uPwOeh0LZOgyov4rww,3706
-envelope/compat/navix_envelope.py,sha256=K7AGae9_kDETFAUuWH9sCCLDo5fuEFbc0eKTHcvSxek,3027
-envelope/wrappers/__init__.py,sha256=r0Qz7gNkgJQF848Iv01uMGbeXTKHUdK2tjyom4nPCGw,701
-envelope/wrappers/autoreset_wrapper.py,sha256=3OUUNb4L6P4Ncn57Uy2ldoLFr1IKnPSWuzyjf20N0Y8,1299
-envelope/wrappers/episode_statistics_wrapper.py,sha256=Mj5Ua7cLtBQYtFn3oQB9eJ7TclKwAFUvsQHIRkcQJvs,1509
-envelope/wrappers/normalization.py,sha256=xHezXsb1J09D3IZACC1xxMy9NuUSBWqCyYH2wFAItPs,1811
-envelope/wrappers/observation_normalization_wrapper.py,sha256=NNHz_THFe2eQUY0Pued9_hJztIGU6SU7kae_AKFbe4k,4248
-envelope/wrappers/state_injection_wrapper.py,sha256=yk7he1zPaEdzjks_NDnYDC91e5bXTlCFPKFMwm-A7jk,3687
-envelope/wrappers/truncation_wrapper.py,sha256=TzJqxjAjipk7pTk-drXRvQ41uR-jWsHlAgYJpbHI32M,1132
-envelope/wrappers/vmap_envs_wrapper.py,sha256=TezSCeOf3oJKj8w7i_UaAA2LsoWkhhp835W545Q-waI,2561
-envelope/wrappers/vmap_wrapper.py,sha256=zYMYVvuVQfRM7d2ygxxcavRNtkdzkqhArakE507tcU4,1587
-envelope/wrappers/wrapper.py,sha256=ydolcZR9ogW3z4ui07845mbMue2XtEoEJ64tHCzAAtw,1504
-jax_envelope-0.1.1.dist-info/METADATA,sha256=JdPA7sTzarKPGRTiCckPxwPpy2MU_Wm-5vaKW6dIelg,4652
-jax_envelope-0.1.1.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-jax_envelope-0.1.1.dist-info/licenses/LICENSE,sha256=VyF-MK-gY2_fZlhf8uEnE2y8ziIXK-w55GM12eOgXrQ,1069
-jax_envelope-0.1.1.dist-info/RECORD,,