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

envelope/wrappers/vmap_envs_wrapper.py

@@ -17,34 +17,39 @@ class VmapEnvsWrapper(Wrapper):
     Usage:
         envs = jax.vmap(make_env)(params_batch)     # env pytree batched on leading axis
         wrapped = VmapEnvsWrapper(env=envs, batch_size=B)
-        state, info = wrapped.reset(keys)           # keys shape (B, 2) or single key
+        state, info = wrapped.init(keys)             # keys shape (B, 2) or single key
         next_state, info = wrapped.step(state, action)
     """
 
     batch_size: int = field(kw_only=True)
 
-    @override
-    def reset(
-        self, key: Key, state: PyTree | None = None, **kwargs
-    ) -> tuple[WrappedState, Info]:
+    def _split_keys(self, key: Key) -> Key:
         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
-        # vmap over env 'self' and keys
-        state, info = jax.vmap(lambda e, k: e.reset(k, state, **kwargs))(self.env, keys)
+            return jax.random.split(key, self.batch_size)
+        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})."
+            )
+        return key
+
+    @override
+    def init(self, key: Key) -> tuple[WrappedState, Info]:
+        keys = self._split_keys(key)
+        state, info = jax.vmap(lambda e, k: e.init(k))(self.env, keys)
+        return state, info
+
+    @override
+    def reset(self, key: Key, state: PyTree) -> tuple[WrappedState, Info]:
+        keys = self._split_keys(key)
+        state, info = jax.vmap(lambda e, k, s: e.reset(k, s))(
+            self.env, keys, state
+        )
         return state, info
 
     @override
-    def step(
-        self, state: WrappedState, action: PyTree, **kwargs
-    ) -> tuple[WrappedState, Info]:
-        next_state, info = jax.vmap(lambda e, s, a: e.step(s, a, **kwargs))(
+    def step(self, state: WrappedState, action: PyTree) -> tuple[WrappedState, Info]:
+        next_state, info = jax.vmap(lambda e, s, a: e.step(s, a))(
             self.env, state, action
         )
         return next_state, info
@@ -53,13 +58,13 @@ class VmapEnvsWrapper(Wrapper):
     @property
     def observation_space(self) -> spaces.Space:
         env0 = _index_env(self.env, 0, self.batch_size)
-        return spaces.batch_space(env0.observation_space, self.batch_size)
+        return spaces.BatchedSpace(space=env0.observation_space, batch_size=self.batch_size)
 
     @override
     @cached_property
     def action_space(self) -> spaces.Space:
         env0 = _index_env(self.env, 0, self.batch_size)
-        return spaces.batch_space(env0.action_space, self.batch_size)
+        return spaces.BatchedSpace(space=env0.action_space, batch_size=self.batch_size)
 
     @override
     @property

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 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
+    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) -> tuple[WrappedState, Info]:
+        keys = _split_or_keep_key(key, self.batch_size)
         state, info = jax.vmap(self.env.reset)(keys, state)
         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, key: Key, state: State) -> tuple[State, Info]:
+        return self.env.reset(key, state)
+
+    @override
+    def step(self, state: WrappedState, action: PyTree) -> tuple[WrappedState, Info]:
+        return self.env.step(state, action)
 
     @override
     @cached_property

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: jax-envelope
-Version: 0.1.1
+Version: 0.2.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
@@ -82,6 +82,6 @@ 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.
+* [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.dist-info/RECORD

@@ -0,0 +1,32 @@
+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,,

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,,