jax-envelope 0.1.0__py3-none-any.whl

envelope/wrappers/episode_statistics_wrapper.py

@@ -0,0 +1,47 @@
+from dataclasses import field
+from envelope.wrappers import Wrapper
+from typing import override
+
+from envelope.environment import Environment, Info, State
+from envelope.typing import Key, PyTree, Array
+
+
+class EpisodeStatisticsWrapper(Wrapper):
+    class StatisticsState(WappedState):
+        episode_reward: Array
+        episode_length: Array
+        _pointer: int = field(default=0)
+
+    def reset(
+        self, key: Key, state: State | None = None, **kwargs
+    ) -> tuple[State, Info]:
+        state, info = self.env.reset(key, state=state, **kwargs)
+        info = 
+        return state, info
+
+
+    @override
+    def reset(
+        self, key: Key, state: State | None = None, **kwargs
+    ) -> tuple[State, Info]:
+        state, info = self.env.reset(key, state=state, **kwargs)
+        info = 
+        return state, info
+
+    @override
+    def step(self, state: State, action: PyTree, **kwargs) -> tuple[State, Info]:
+        next_state, info = self.env.step(state, action, **kwargs)
+        info = self._update_episode_statistics(info)
+        return next_state, info
+
+    def _update_episode_statistics(self, info: Info) -> Info:
+        """Update episode statistics in the info dictionary."""
+        if "episode_statistics" not in info:
+            info["episode_statistics"] = {
+                "reward": 0.0,
+                "length": 0,
+            }
+        info["episode_statistics"]["reward"] += info.get("reward", 0.0)
+        info["episode_statistics"]["length"] += 1
+        return info
+

envelope/wrappers/normalization.py

@@ -0,0 +1,56 @@
+from functools import cached_property
+from typing import NamedTuple
+
+import jax
+from jax import numpy as jnp
+
+from envelope.struct import FrozenPyTreeNode
+from envelope.typing import Array, PyTree
+
+
+class MeanVarPair(NamedTuple):
+    mean: Array
+    var: Array
+
+
+class RunningMeanVar(FrozenPyTreeNode):
+    mean: PyTree
+    var: PyTree
+    count: int
+
+    @cached_property
+    def std(self) -> PyTree:
+        return jax.tree.map(jnp.sqrt, self.var)
+
+
+def update_rmv(rmv_state: RunningMeanVar, x: PyTree) -> RunningMeanVar:
+    """
+    Update running mean/variance with a new batch of observations x. We assume x is a
+    PyTree of arrays, each with a leading batch dimension (aligned sizes).
+    """
+    global_count = rmv_state.count
+    batch_count = jax.tree.leaves(x)[0].shape[0]
+    tot_count = global_count + batch_count
+
+    def _update_arr(mean: Array, var: Array, x_arr: Array) -> MeanVarPair:
+        batch_mean = x_arr.mean(axis=0)
+        batch_var = x_arr.var(axis=0)
+
+        # Combine variances using parallel algorithm
+        m_a = var * global_count
+        m_b = batch_var * batch_count
+        delta = batch_mean - mean
+        m2 = m_a + m_b + (delta**2) * (global_count * batch_count) / tot_count
+
+        new_mean = mean + delta * (batch_count / tot_count)
+        new_var = m2 / tot_count
+        return MeanVarPair(mean=new_mean, var=new_var)
+
+    def is_pair(z):
+        return isinstance(z, MeanVarPair)
+
+    # jax.tree.map returns a PyTree whose leaves are MeanVarPairs
+    mean_var_pairs = jax.tree.map(_update_arr, rmv_state.mean, rmv_state.var, x)
+    new_mean = jax.tree.map(lambda mv: mv.mean, mean_var_pairs, is_leaf=is_pair)
+    new_var = jax.tree.map(lambda mv: mv.var, mean_var_pairs, is_leaf=is_pair)
+    return RunningMeanVar(mean=new_mean, var=new_var, count=tot_count)

envelope/wrappers/observation_normalization_wrapper.py

@@ -0,0 +1,114 @@
+from typing import override
+
+import jax
+from jax import numpy as jnp
+
+from envelope.environment import Info
+from envelope.spaces import BatchedSpace, PyTreeSpace, Space
+from envelope.struct import field, static_field
+from envelope.typing import Key, PyTree
+from envelope.wrappers.normalization import RunningMeanVar, update_rmv
+from envelope.wrappers.wrapper import WrappedState, Wrapper
+
+
+class ObservationNormalizationWrapper(Wrapper):
+    class ObservationNormalizationState(WrappedState):
+        rmv_state: RunningMeanVar = field()
+
+    stats_spec: PyTree | None = static_field(default=None)
+    """Per-leaf normalization statistics spec as a pytree of jax.ShapeDtypeStruct.
+    Shapes must be broadcastable to the observation leaves. If None, inferred from
+    the observation_space with BatchedSpace ignored; each leaf must have a floating
+    dtype."""
+
+    def __post_init__(self):
+        if self.stats_spec is None:
+            stats_spec = _infer_stats_spec(self.env.observation_space)
+            object.__setattr__(self, "stats_spec", stats_spec)
+
+    def _init_rmv_state(self) -> RunningMeanVar:
+        def zeros(sd: jax.ShapeDtypeStruct) -> jax.Array:
+            return jnp.zeros(sd.shape, dtype=sd.dtype)
+
+        def ones(sd: jax.ShapeDtypeStruct) -> jax.Array:
+            return jnp.ones(sd.shape, dtype=sd.dtype)
+
+        mean = jax.tree.map(zeros, self.stats_spec)
+        var = jax.tree.map(ones, self.stats_spec)
+
+        return RunningMeanVar(mean=mean, var=var, count=0)
+
+    def _normalize_obs(self, obs: PyTree, rmv: RunningMeanVar) -> PyTree:
+        def norm_leaf(x, mean, std, spec):
+            mean = jnp.broadcast_to(mean, x.shape)
+            std = jnp.broadcast_to(std, x.shape)
+            obs = (x - mean) / (std + 1e-8)
+            return obs.astype(spec.dtype)
+
+        return jax.tree.map(norm_leaf, obs, rmv.mean, rmv.std, self.stats_spec)
+
+    def _normalize_and_update(
+        self, state: WrappedState, info: Info
+    ) -> tuple[WrappedState, Info]:
+        # Ensure each observation leaf is shaped as (-1, *spec.shape)
+        reshaped_obs = jax.tree.map(
+            lambda x, spec: x.reshape((-1,) + tuple(spec.shape)),
+            info.obs,
+            self.stats_spec,
+        )
+        rmv_state = update_rmv(state.rmv_state, reshaped_obs)
+        norm_obs = self._normalize_obs(info.obs, rmv_state)
+
+        state = self.ObservationNormalizationState(
+            inner_state=state.inner_state, rmv_state=rmv_state
+        )
+        info = info.update(obs=norm_obs, unnormalized_obs=info.obs)
+        return state, info
+
+    @override
+    def reset(
+        self, key: Key, state: PyTree | None = None, **kwargs
+    ) -> tuple[WrappedState, Info]:
+        inner_state = None
+        rmv_state = self._init_rmv_state()
+        if state:
+            inner_state = state.inner_state
+            rmv_state = state.rmv_state
+
+        inner_state, info = self.env.reset(key, inner_state, **kwargs)
+        next_state = self.ObservationNormalizationState(
+            inner_state=inner_state, rmv_state=rmv_state
+        )
+        return self._normalize_and_update(next_state, info)
+
+    @override
+    def step(
+        self, state: WrappedState, action: PyTree, **kwargs
+    ) -> tuple[WrappedState, Info]:
+        inner_state, info = self.env.step(state.inner_state, action, **kwargs)
+        state = state.replace(inner_state=inner_state)
+        return self._normalize_and_update(state, info)
+
+
+def _infer_stats_spec(space: Space) -> PyTree:
+    """
+    Build a PyTree of jax.ShapeDtypeStruct for stats. Strip BatchedSpace layers,
+    and for leaf spaces return (shape=space.shape, dtype=inferred).
+    """
+
+    def descend(sp: Space):
+        if isinstance(sp, BatchedSpace):
+            return descend(sp.space)
+        if isinstance(sp, PyTreeSpace):
+            return jax.tree.map(
+                lambda s: descend(s),
+                sp.tree,
+                is_leaf=lambda n: isinstance(n, Space),
+            )
+        if not jnp.issubdtype(sp.dtype, jnp.floating):
+            raise ValueError(
+                f"Space {sp} has dtype {sp.dtype} which is not a floating point dtype"
+            )
+        return jax.ShapeDtypeStruct(tuple(sp.shape), sp.dtype)
+
+    return descend(space)

envelope/wrappers/state_injection_wrapper.py

@@ -0,0 +1,91 @@
+from envelope.environment import Info, InfoContainer
+from envelope.struct import field
+from envelope.typing import Key, PyTree
+from envelope.wrappers.wrapper import WrappedState, Wrapper
+
+
+class StateInjectionWrapper(Wrapper):
+    """Stores a state that all resets return to.
+
+    For UED: use set_reset_state() to update the injected state, then all resets
+    (including auto-reset) return to that state until it's changed again.
+
+    Usage:
+        env = AutoResetWrapper(StateInjectionWrapper(env=base_env))
+        state, info = env.reset(key)
+
+        for outer_iter in range(num_outer_iters):
+            # Sample a new task and set it as the reset state
+            task_state, task_obs = sample_task(key)
+            state = env.set_reset_state(state, task_state, task_obs)
+
+            # Run episode - auto-resets return to task_state
+            for inner_step in range(num_inner_steps):
+                state, info = env.step(state, policy(info.obs))
+    """
+
+    class InjectedState(WrappedState):
+        reset_state: PyTree | None = field(default=None)
+        reset_obs: PyTree | None = field(default=None)
+
+    def set_reset_state(
+        self, state: WrappedState, reset_state: PyTree, reset_obs: PyTree
+    ) -> WrappedState:
+        """Update the state that resets will return to.
+
+        This method traverses the wrapped state hierarchy to find and update
+        the InjectedState, then reconstructs the full state tree.
+
+        Args:
+            state: Current state (can be from any outer wrapper)
+            reset_state: The state to reset to (inner environment state)
+            reset_obs: The observation to return on reset
+
+        Returns:
+            New state with updated reset fields at the appropriate level
+        """
+
+        def update_injected(s: WrappedState) -> WrappedState:
+            # If this is our InjectedState, update it
+            if isinstance(s, self.InjectedState):
+                return self.InjectedState(
+                    inner_state=reset_state,
+                    reset_state=reset_state,
+                    reset_obs=reset_obs,
+                )
+            # Otherwise, recurse into inner_state and rebuild
+            if hasattr(s, "inner_state"):
+                return s.replace(inner_state=update_injected(s.inner_state))
+            raise ValueError("Could not find InjectedState in given state")
+
+        return update_injected(state)
+
+    def reset(
+        self, key: Key, state: PyTree | None = None, **kwargs
+    ) -> tuple[WrappedState, Info]:
+        # Default state has no inner state to reset to
+        if state is None:
+            state = self.InjectedState(inner_state=None)
+
+        # If no reset state is set, reset wrapped environment
+        if state.reset_state is None and state.reset_obs is None:
+            inner_state, info = self.env.reset(key, state=state.inner_state, **kwargs)
+
+        # If reset state is set, use it
+        elif state.reset_state is not None and state.reset_obs is not None:
+            inner_state = state.reset_state
+            info = InfoContainer(obs=state.reset_obs, reward=0.0, terminated=False)
+
+        # If only one of reset_state or reset_obs is set, raise error
+        else:
+            raise ValueError("State must set both reset_state and reset_obs or neither")
+
+        # Return new state with updated inner state
+        state = state.replace(inner_state=inner_state)
+        return state, info
+
+    def step(
+        self, state: WrappedState, action: PyTree, **kwargs
+    ) -> tuple[WrappedState, Info]:
+        inner_state, info = self.env.step(state.inner_state, action, **kwargs)
+        return state.replace(inner_state=inner_state), info

envelope/wrappers/timestep_wrapper.py

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

envelope/wrappers/truncation_wrapper.py

@@ -0,0 +1,31 @@
+import jax.numpy as jnp
+
+from envelope.environment import Info
+from envelope.struct import field
+from envelope.typing import Key, PyTree
+from envelope.wrappers.wrapper import WrappedState, Wrapper
+
+
+class TruncationWrapper(Wrapper):
+    max_steps: int = field(kw_only=True)
+
+    class TruncationState(WrappedState):
+        steps: jnp.ndarray | int = field(default=0)
+
+    def reset(
+        self, key: Key, state: PyTree | None = None, **kwargs
+    ) -> tuple[WrappedState, Info]:
+        inner_state, info = self.env.reset(key)
+        state = self.TruncationState(inner_state=inner_state, steps=0)
+        return state, info.update(truncated=self.max_steps <= 0)
+
+    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 = state.steps + 1
+        next_state = self.TruncationState(
+            inner_state=next_inner_state, steps=next_steps
+        )
+        truncated = jnp.asarray(next_steps) >= self.max_steps
+        return next_state, info.update(truncated=truncated)

envelope/wrappers/vmap_envs_wrapper.py

@@ -0,0 +1,77 @@
+from functools import cached_property
+from typing import override
+
+import jax
+
+from envelope import spaces
+from envelope.environment import Environment, Info
+from envelope.struct import field
+from envelope.typing import Key, PyTree
+from envelope.wrappers.wrapper import WrappedState, Wrapper
+
+
+class VmapEnvsWrapper(Wrapper):
+    """
+    Vectorizes over a batched collection of environment instances (vmapping over 'self').
+
+    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
+        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]:
+        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 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))(
+            self.env, state, action
+        )
+        return next_state, info
+
+    @override
+    @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)
+
+    @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)
+
+    @override
+    @property
+    def unwrapped(self) -> Environment:
+        return self.env.unwrapped
+
+
+def _index_env(env: Environment, idx: int, batch_size: int) -> Environment:
+    def idx_or_keep(x):
+        if hasattr(x, "shape") and isinstance(getattr(x, "shape"), tuple):
+            if len(x.shape) > 0 and x.shape[0] == batch_size:
+                return x[idx]
+        return x
+
+    return jax.tree.map(lambda x: idx_or_keep(x), env)

envelope/wrappers/vmap_wrapper.py

@@ -0,0 +1,51 @@
+from functools import cached_property
+from typing import override
+
+import jax
+
+from envelope import spaces
+from envelope.environment import Info
+from envelope.struct import field
+from envelope.typing import Key, PyTree
+from envelope.wrappers.wrapper import WrappedState, Wrapper
+
+
+class VmapWrapper(Wrapper):
+    """Does not forward kwargs to the underlying env. Does not wrap the state."""
+
+    batch_size: int = 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
+
+        state, info = jax.vmap(self.env.reset)(keys, state)
+        return state, info
+
+    @override
+    def step(
+        self, state: WrappedState, action: PyTree, **kwargs
+    ) -> 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)
+
+    @override
+    @cached_property
+    def action_space(self) -> spaces.Space:
+        return spaces.batch_space(self.env.action_space, self.batch_size)

envelope/wrappers/wrapper.py

@@ -0,0 +1,57 @@
+from dataclasses import KW_ONLY
+from functools import cached_property
+from typing import override
+
+from envelope import spaces
+from envelope.environment import Environment, Info, State
+from envelope.struct import FrozenPyTreeNode, field
+from envelope.typing import Key, PyTree
+
+
+class WrappedState(FrozenPyTreeNode):
+    inner_state: State = field()
+    _: KW_ONLY
+
+    @property
+    def unwrapped(self) -> State:
+        if hasattr(self.inner_state, "unwrapped"):
+            return self.inner_state.unwrapped
+        return self.inner_state
+
+
+class Wrapper(Environment):
+    """Wrapper for environments."""
+
+    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)
+
+    @override
+    def step(
+        self, state: WrappedState, action: PyTree, **kwargs
+    ) -> tuple[WrappedState, Info]:
+        return self.env.step(state, action, **kwargs)
+
+    @override
+    @cached_property
+    def observation_space(self) -> spaces.Space:
+        return self.env.observation_space
+
+    @override
+    @cached_property
+    def action_space(self) -> spaces.Space:
+        return self.env.action_space
+
+    @override
+    @property
+    def unwrapped(self) -> Environment:
+        return self.env.unwrapped
+
+    def __getattr__(self, name):
+        if name == "__setstate__":
+            raise AttributeError(name)
+        return getattr(self.env, name)

jax_envelope-0.1.0.dist-info/METADATA

@@ -0,0 +1,87 @@
+Metadata-Version: 2.4
+Name: jax-envelope
+Version: 0.1.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
+Project-URL: Documentation, https://github.com/keraJLi/envelope#readme
+Project-URL: Issues, https://github.com/keraJLi/envelope/issues
+Project-URL: Changelog, https://github.com/keraJLi/envelope/releases
+Author-email: Jarek Liesen <jarek.liesen@reuben.ox.ac.uk>
+License: MIT
+License-File: LICENSE
+Keywords: deep-learning,environments,gymnasium,hardware-acceleration,jax,machine-learning,reinforcement-learning,vectorization
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Typing :: Typed
+Requires-Python: >=3.12
+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)
+states, infos = jax.lax.scan(env.step, state, actions)
+plt.plot(infos.reward.cumsum())
+
+# ... and enjoy a powerful ecosystem of wrappers.
+env = envelope.wrappers.AutoResetWrapper(env)
+env = envelope.wrappers.VmapWrapper(env)
+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!
+
+## 💪 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. -->
+
+## 🔌 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 |
+
+```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`
+  - `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.
+
+## 🏗️ 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!

jax_envelope-0.1.0.dist-info/RECORD

@@ -0,0 +1,27 @@
+envelope/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+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/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/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.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.28.0
+Root-Is-Purelib: true
+Tag: py3-none-any