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

envelope/__init__.py

@@ -1,4 +1,4 @@
-from envelope.compat import create
+from envelope.adapters import create
 from envelope.environment import Environment, Info, InfoContainer
 from envelope.spaces import BatchedSpace, Continuous, Discrete, PyTreeSpace, Space
 from envelope.struct import Container, FrozenPyTreeNode, field, static_field

envelope/{compat → adapters}/__init__.py

@@ -4,14 +4,14 @@ from typing import Any, Protocol, Self
 
 # Lazy imports to avoid requiring all dependencies at once
 _env_module_map = {
-    "gymnax": ("envelope.compat.gymnax_envelope", "GymnaxEnvelope"),
-    "brax": ("envelope.compat.brax_envelope", "BraxEnvelope"),
-    "navix": ("envelope.compat.navix_envelope", "NavixEnvelope"),
-    "jumanji": ("envelope.compat.jumanji_envelope", "JumanjiEnvelope"),
-    "kinetix": ("envelope.compat.kinetix_envelope", "KinetixEnvelope"),
-    "craftax": ("envelope.compat.craftax_envelope", "CraftaxEnvelope"),
+    "gymnax": ("envelope.adapters.gymnax_envelope", "GymnaxEnvelope"),
+    "brax": ("envelope.adapters.brax_envelope", "BraxEnvelope"),
+    "navix": ("envelope.adapters.navix_envelope", "NavixEnvelope"),
+    "jumanji": ("envelope.adapters.jumanji_envelope", "JumanjiEnvelope"),
+    "kinetix": ("envelope.adapters.kinetix_envelope", "KinetixEnvelope"),
+    "craftax": ("envelope.adapters.craftax_envelope", "CraftaxEnvelope"),
     "mujoco_playground": (
-        "envelope.compat.mujoco_playground_envelope",
+        "envelope.adapters.mujoco_playground_envelope",
         "MujocoPlaygroundEnvelope",
     ),
 }

envelope/{compat → adapters}/craftax_envelope.py

@@ -10,7 +10,7 @@ from craftax.craftax_classic.envs.craftax_state import (
 from craftax.craftax_env import make_craftax_env_from_name
 
 from envelope import spaces as envelope_spaces
-from envelope.compat.gymnax_envelope import _convert_space as _convert_gymnax_space
+from envelope.adapters.gymnax_envelope import _convert_space as _convert_gymnax_space
 from envelope.environment import Environment, Info, InfoContainer, State
 from envelope.struct import Container, static_field
 from envelope.typing import Key, PyTree, TypeAlias

envelope/{compat → adapters}/kinetix_envelope.py

@@ -1,7 +1,7 @@
 """Kinetix compatibility wrapper.
 
 This module exposes Kinetix environments through the `envelope.environment.Environment`
-API. It mirrors envelope's compat philosophy:
+API. It mirrors envelope's adapters philosophy:
 - prefer *no* environment-side auto-reset (use `AutoResetWrapper` in envelope)
 - prefer *no* fixed episode time-limits (use `TruncationWrapper` in envelope)
 
@@ -30,7 +30,7 @@ from kinetix.util.saving import load_from_json_file
 
 from envelope import field
 from envelope import spaces as envelope_spaces
-from envelope.compat.gymnax_envelope import _convert_space as _convert_gymnax_space
+from envelope.adapters.gymnax_envelope import _convert_space as _convert_gymnax_space
 from envelope.environment import Environment, Info, InfoContainer, State
 from envelope.struct import Container, static_field
 from envelope.typing import Key, PyTree

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: jax-envelope
-Version: 0.3.0
+Version: 0.4.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
@@ -42,6 +42,7 @@ 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 `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`.
@@ -53,8 +54,6 @@ env = envelope.wrappers.ObservationNormalizationWrapper(env)
 - **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
 
 
@@ -78,11 +77,11 @@ 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`
-  - `uv run pytest -m compat`
-  - If any compat dependency is missing/broken, the run will fail fast with an error telling you what to install.
+- **Default (no optional adapters deps required)**: `uv run pytest -m "not adapters"`
+- **Adapters suite (requires full adapters dependency group)**:
+  - `uv sync --group adapters`
+  - `uv run pytest -m adapters`
+  - If any adapter dependency is missing/broken, the run will fail fast with an error telling you what to install.
 
 ## 🏗️ Installation
 

{jax_envelope-0.3.0.dist-info → jax_envelope-0.4.0.dist-info}/RECORD

@@ -1,16 +1,16 @@
-envelope/__init__.py,sha256=AAk0b9UHGUTC4EzERXYPOrHuECgO0fx_UN4_gpKvgT0,1329
+envelope/__init__.py,sha256=3IaAPQ12MPozGX1a73c6_EGaI26-UdXYtu5rqJfPBYs,1331
 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/adapters/__init__.py,sha256=Xn2jNxQK1tdS1P1FEzCNTO8ulLuKfDaiozWkOqNRTPQ,3472
+envelope/adapters/brax_envelope.py,sha256=xToV4fJYU9gPfF5SDs1Y3BaP96daz9iW8gKs2QYQDNA,3476
+envelope/adapters/craftax_envelope.py,sha256=G9jLBJjV_YQh_dnnQlk0hm_hW5B2lyjz7GiRZnIo-FM,3961
+envelope/adapters/gymnax_envelope.py,sha256=X0pfFKaFKzvBl2Zih4utiXSEI-lc5f9miJGMxeT-Xfw,4690
+envelope/adapters/jumanji_envelope.py,sha256=zl6yE1yqdCxZmjVCexu5vD10f5a-M82EpjzyVlj13wM,4909
+envelope/adapters/kinetix_envelope.py,sha256=Nh98HXzuXB_gCWKOdYWBiaxa2v1zyhmYzrUw2xmfwWo,6857
+envelope/adapters/mujoco_playground_envelope.py,sha256=7iPaRuwXUHbP8hHNgSFipW8rJVruootMQPYBWNpPiHA,3705
+envelope/adapters/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
@@ -26,7 +26,7 @@ envelope/wrappers/truncation_wrapper.py,sha256=k3NZBMmnYiCN35jA1Yyi3rGdLcqM52Czu
 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.4.0.dist-info/METADATA,sha256=gxqlnwoNe2wcVPfYMZQE9HRROD8q7CDkni1ND1akrKg,5056
+jax_envelope-0.4.0.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+jax_envelope-0.4.0.dist-info/licenses/LICENSE,sha256=VyF-MK-gY2_fZlhf8uEnE2y8ziIXK-w55GM12eOgXrQ,1069
+jax_envelope-0.4.0.dist-info/RECORD,,