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

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

@@ -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 → jax_envelope-0.1.1}/README.md

@@ -55,7 +55,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 → jax_envelope-0.1.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "jax-envelope"
-version = "0.1.0"
+version = "0.1.1"
 description = "A JAX-native environment interface with powerful wrappers and adapters for popular RL environment suites"
 readme = "README.md"
 requires-python = ">=3.12"

jax_envelope-0.1.1/src/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",
+]

jax_envelope-0.1.1/src/envelope/typing.py

@@ -0,0 +1,7 @@
+from typing import Any, TypeAlias
+
+import jax
+
+PyTree: TypeAlias = Any
+Key: TypeAlias = jax.Array
+Array: TypeAlias = jax.Array

jax_envelope-0.1.1/src/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 → jax_envelope-0.1.1}/uv.lock

@@ -1049,7 +1049,7 @@ wheels = [
 
 [[package]]
 name = "jax-envelope"
-version = "0.1.0"
+version = "0.1.1"
 source = { editable = "." }
 dependencies = [
     { name = "jax" },

jax_envelope-0.1.0/src/envelope/typing.py

@@ -1,23 +0,0 @@
-from enum import Enum
-from typing import Any, TypeAlias
-
-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"

jax_envelope-0.1.0/src/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