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

envelope/__init__.py

@@ -0,0 +1,42 @@
+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.wrappers import (
+    Wrapper,
+    WrappedState,
+    AutoResetWrapper,
+    ObservationNormalizationWrapper,
+    StateInjectionWrapper,
+    TruncationWrapper,
+    VmapWrapper,
+    VmapEnvsWrapper,
+)
+
+__all__ = [
+    # Basic functionality
+    "create",
+    "Environment",
+    "Info",
+    "InfoContainer",
+    # Spaces
+    "Space",
+    "BatchedSpace",
+    "Continuous",
+    "Discrete",
+    "PyTreeSpace",
+    # Struct
+    "field",
+    "static_field",
+    "FrozenPyTreeNode",
+    "Container",
+    # Wrappers
+    "Wrapper",
+    "WrappedState",
+    "AutoResetWrapper",
+    "ObservationNormalizationWrapper",
+    "StateInjectionWrapper",
+    "TruncationWrapper",
+    "VmapWrapper",
+    "VmapEnvsWrapper",
+]

envelope/typing.py

@@ -1,4 +1,3 @@
-from enum import Enum
 from typing import Any, TypeAlias
 
 import jax
@@ -6,18 +5,3 @@ import jax
 PyTree: TypeAlias = Any
 Key: TypeAlias = jax.Array
 Array: TypeAlias = jax.Array
-
-
-class BatchKind(Enum):
-    """
-    Batch semantics for environments.
-    - VMAP: environment represents instances compatible with `jax.vmap`, for example by
-      wrapping it in a `VmapWrapper`.
-    - NATIVE_POOL: environment represents a batch of instances via a native pool. This
-      is the case when it is wrapping a non-jax-based environment that supports native
-      batching, for example those provided by envpool. Environments in this mode cannot
-      be vmapped, as it would break the native batching semantics.
-    """
-
-    VMAP = "vmap"
-    NATIVE_POOL = "native_pool"

envelope/wrappers/__init__.py

@@ -0,0 +1,20 @@
+from envelope.wrappers.autoreset_wrapper import AutoResetWrapper
+from envelope.wrappers.observation_normalization_wrapper import (
+    ObservationNormalizationWrapper,
+)
+from envelope.wrappers.state_injection_wrapper import StateInjectionWrapper
+from envelope.wrappers.truncation_wrapper import TruncationWrapper
+from envelope.wrappers.vmap_wrapper import VmapWrapper
+from envelope.wrappers.vmap_envs_wrapper import VmapEnvsWrapper
+from envelope.wrappers.wrapper import Wrapper, WrappedState
+
+__all__ = [
+    "Wrapper",
+    "WrappedState",
+    "AutoResetWrapper",
+    "ObservationNormalizationWrapper",
+    "StateInjectionWrapper",
+    "TruncationWrapper",
+    "VmapWrapper",
+    "VmapEnvsWrapper",
+]

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: jax-envelope
-Version: 0.1.0
+Version: 0.1.1
 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
@@ -81,7 +81,7 @@ let's you create environments from any of the above!
 pip install jax-envelope
 ```
 
-## 💞 Related projects
+## 💞 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!

{jax_envelope-0.1.0.dist-info → jax_envelope-0.1.1.dist-info}/RECORD

@@ -1,8 +1,8 @@
-envelope/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,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=dcftDRNM0luCqHDEJ_qVmZnjUOEwn-HiFvwuV1z1Bn0,734
+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
@@ -11,17 +11,17 @@ envelope/compat/jumanji_envelope.py,sha256=_Tox7DfiyE-Z4drszPUe-arPurz6hMeFtRNfy
 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/timestep_wrapper.py,sha256=6jS-80AwnIIct3Ool9zq_iGOyfodNbs89NfuyYMfYIE,874
 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.0.dist-info/METADATA,sha256=-ASxsWLgF9CwcQ-mp1NJAmdFDS4qzs8bekNJH291R00,4653
-jax_envelope-0.1.0.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-jax_envelope-0.1.0.dist-info/licenses/LICENSE,sha256=VyF-MK-gY2_fZlhf8uEnE2y8ziIXK-w55GM12eOgXrQ,1069
-jax_envelope-0.1.0.dist-info/RECORD,,
+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,,

envelope/wrappers/timestep_wrapper.py

@@ -1,22 +0,0 @@
-from envelope.environment import Info
-from envelope.struct import field
-from envelope.typing import Key, PyTree
-from envelope.wrappers.wrapper import WrappedState, Wrapper
-
-
-class TimeStepWrapper(Wrapper):
-    class TimeStepState(WrappedState):
-        steps: PyTree = field(default=0)
-
-    def reset(
-        self, key: Key, state: PyTree | None = None, **kwargs
-    ) -> tuple[WrappedState, Info]:
-        inner_state, info = self.env.reset(key, state, **kwargs)
-        return self.TimeStepState(inner_state=inner_state, steps=0), info
-
-    def step(
-        self, state: WrappedState, action: PyTree, **kwargs
-    ) -> tuple[WrappedState, Info]:
-        next_inner_state, info = self.env.step(state.inner_state, action, **kwargs)
-        next_steps = getattr(state, "steps", 0) + 1
-        return self.TimeStepState(inner_state=next_inner_state, steps=next_steps), info