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

envelope/wrappers/continuous_observation_wrapper.py

@@ -0,0 +1,61 @@
+from functools import cached_property
+from typing import override
+
+import jax
+import jax.numpy as jnp
+
+from envelope.environment import Info, State
+from envelope.spaces import BatchedSpace, Continuous, Discrete, Space, peel_batched
+from envelope.typing import Key, PyTree
+from envelope.wrappers.wrapper import Wrapper
+
+
+def to_float(obs: PyTree) -> PyTree:
+    return jax.tree.map(lambda x: x.astype(jnp.float32), obs)
+
+
+def to_continuous(space: Discrete | Continuous) -> Continuous:
+    if isinstance(space, Continuous):
+        low = jnp.asarray(space.low, dtype=jnp.float32)
+        high = jnp.asarray(space.high, dtype=jnp.float32)
+        return Continuous(low=low, high=high)
+    elif isinstance(space, Discrete):
+        n = jnp.asarray(space.n)
+        low = jnp.zeros_like(n, dtype=jnp.float32)
+        high = jnp.asarray(n - 1, dtype=jnp.float32)
+        return Continuous(low=low, high=high)
+    raise TypeError(f"Expected Discrete or Continuous, got {type(space)}")
+
+
+class ContinuousObservationWrapper(Wrapper):
+    @override
+    def init(self, key: Key) -> tuple[State, Info]:
+        state, info = self.env.init(key)
+        info = info.update(obs=to_float(info.obs))
+        return state, info
+
+    @override
+    def reset(self, key: Key, state: State) -> tuple[State, Info]:
+        state, info = self.env.reset(key, state)
+        info = info.update(obs=to_float(info.obs))
+        return state, info
+
+    @override
+    def step(self, state: State, action: PyTree) -> tuple[State, Info]:
+        state, info = self.env.step(state, action)
+        info = info.update(obs=to_float(info.obs))
+        return state, info
+
+    @override
+    @cached_property
+    def observation_space(self) -> Space:
+        batch_dims, base = peel_batched(self.env.observation_space)
+
+        def is_leaf(x):
+            return isinstance(x, (Discrete, Continuous))
+
+        space = jax.tree.map(to_continuous, base, is_leaf=is_leaf)
+
+        for batch_dim in batch_dims:
+            space = BatchedSpace(space, batch_dim)
+        return space

envelope/wrappers/episode_statistics_wrapper.py

@@ -1,47 +1,40 @@
-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
+import jax
 
+from envelope.environment import Info, State
+from envelope.struct import FrozenPyTreeNode, field
+from envelope.typing import Key, PyTree
+from envelope.wrappers.wrapper import WrappedState, Wrapper
 
-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
+class EpisodeStatistics(FrozenPyTreeNode):
+    reward: jax.Array = field(default=0)
+    length: jax.Array = field(default=0)
+
 
+class EpisodeStatisticsWrapper(Wrapper):
+    class EpisodeStatisticsState(WrappedState):
+        stats: EpisodeStatistics = field(default=EpisodeStatistics())
 
     @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
+    def init(self, key: Key) -> tuple[State, Info]:
+        inner_state, info = self.env.init(key)
+        state = self.EpisodeStatisticsState(inner_state=inner_state)
+        return state, info.update(stats=state.stats)
 
     @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
+    def reset(self, key: Key, state: State) -> tuple[State, Info]:
+        inner_state, info = self.env.reset(key, state.inner_state)
+        state = state.replace(inner_state=inner_state)
+        return state, info.update(stats=state.stats)
 
+    @override
+    def step(self, state: State, action: PyTree) -> tuple[State, Info]:
+        inner_state, info = self.env.step(state.inner_state, action)
+        stats = state.stats.replace(
+            reward=state.stats.reward + info.reward,
+            length=state.stats.length + 1,
+        )
+        state = state.replace(inner_state=inner_state, stats=stats)
+        return state, info.update(stats=stats)

envelope/wrappers/flatten_action_wrapper.py

@@ -0,0 +1,75 @@
+from functools import cached_property
+from typing import override
+
+import jax
+import jax.numpy as jnp
+
+from envelope.environment import Info, State
+from envelope.spaces import (
+    BatchedSpace,
+    Continuous,
+    Discrete,
+    PyTreeSpace,
+    Space,
+    peel_batched,
+)
+from envelope.typing import PyTree
+from envelope.wrappers.wrapper import Wrapper
+
+
+def flatten_space(space: PyTreeSpace | Continuous | Discrete):
+    def is_leaf(x):
+        # Tuples containing only integers are shape tuples (leaves)
+        # PyTreeSpace can only have tuples that contain at least a Space, so
+        # tuples with only integers must be shape tuples from leaf spaces
+        return isinstance(x, tuple) and all(isinstance(i, int) for i in x)
+
+    shapes, treedef = jax.tree.flatten(space.shape, is_leaf=is_leaf)
+    dims = [jnp.prod(jnp.asarray(shape)) for shape in shapes]
+    return treedef, shapes, dims
+
+
+def unflatten_x(x: jax.Array, treedef, shapes, dims):
+    indices = jnp.cumsum(jnp.array(dims))[:-1]  # last split is the remainder
+    xs = jnp.split(x, indices)
+    xs = jax.tree.map(lambda x, shape: x.reshape(shape), xs, shapes)
+    return jax.tree.unflatten(treedef, xs)
+
+
+class FlattenActionWrapper(Wrapper):
+    @override
+    def step(self, state: State, action: PyTree) -> tuple[State, Info]:
+        treedef, shapes, dims = flatten_space(self.env.action_space)
+        action = unflatten_x(action, treedef, shapes, dims)
+        return self.env.step(state, action)
+
+    @override
+    @cached_property
+    def action_space(self) -> Space:
+        batch_dims, base = peel_batched(self.env.action_space)
+
+        def is_leaf(x):
+            return isinstance(x, (Continuous, Discrete))
+
+        spaces = jax.tree.leaves(base, is_leaf=is_leaf)
+        act_cls = type(spaces[0])
+
+        if not all(isinstance(space, act_cls) for space in spaces):
+            raise ValueError("All spaces must be of the same type")
+
+        if act_cls == Continuous:
+            lows = [jnp.asarray(s.low).reshape(-1) for s in spaces]
+            highs = [jnp.asarray(s.high).reshape(-1) for s in spaces]
+            low = jnp.concatenate(lows, axis=0)
+            high = jnp.concatenate(highs, axis=0)
+            space = Continuous(low=low, high=high)
+        elif act_cls == Discrete:
+            ns = [jnp.asarray(s.n).reshape(-1) for s in spaces]
+            n = jnp.concatenate(ns, axis=0)
+            space = Discrete(n=n)
+        else:
+            raise ValueError(f"Unsupported space type: {act_cls}")
+
+        for batch_dim in batch_dims:
+            space = BatchedSpace(space, batch_dim)
+        return space

envelope/wrappers/flatten_observation_wrapper.py

@@ -0,0 +1,81 @@
+from functools import cached_property
+from typing import override
+
+import jax
+import jax.numpy as jnp
+
+from envelope.environment import Info, State
+from envelope.spaces import BatchedSpace, Continuous, Discrete, Space, peel_batched
+from envelope.typing import Key, PyTree
+from envelope.wrappers.wrapper import Wrapper
+
+
+def flatten_space(space: Space):
+    def is_leaf(x):
+        # Tuples containing only integers are shape tuples (leaves)
+        # PyTreeSpace can only have tuples that contain at least a Space, so
+        # tuples with only integers must be shape tuples from leaf spaces
+        return isinstance(x, tuple) and all(isinstance(i, int) for i in x)
+
+    shapes, treedef = jax.tree.flatten(space.shape, is_leaf=is_leaf)
+    dims = [jnp.prod(jnp.asarray(shape)) for shape in shapes]
+    return treedef, shapes, dims
+
+
+def flatten_x(x: PyTree):
+    leaves = jax.tree.leaves(x)
+    xs = jax.tree.map(lambda x: jnp.asarray(x).reshape(-1), leaves)
+    x = jnp.concatenate(xs, axis=0)
+    return x
+
+
+class FlattenObservationWrapper(Wrapper):
+    @override
+    def init(self, key: Key) -> tuple[State, Info]:
+        state, info = self.env.init(key)
+        info = info.update(obs=flatten_x(info.obs))
+        return state, info
+
+    @override
+    def reset(self, key: Key, state: State) -> tuple[State, Info]:
+        state, info = self.env.reset(key, state)
+        info = info.update(obs=flatten_x(info.obs))
+        return state, info
+
+    @override
+    def step(self, state: State, action: PyTree) -> tuple[State, Info]:
+        state, info = self.env.step(state, action)
+        info = info.update(obs=flatten_x(info.obs))
+        return state, info
+
+    @override
+    @cached_property
+    def observation_space(self) -> Space:
+        batch_dims, base = peel_batched(self.env.observation_space)
+
+        def is_leaf(x):
+            spaces = (Continuous, Discrete)
+            return isinstance(x, spaces)
+
+        spaces = jax.tree.leaves(base, is_leaf=is_leaf)
+        obs_cls = type(spaces[0])
+
+        if not all(isinstance(space, obs_cls) for space in spaces):
+            raise ValueError("All spaces must be of the same type")
+
+        if obs_cls == Continuous:
+            lows = [jnp.asarray(s.low).reshape(-1) for s in spaces]
+            highs = [jnp.asarray(s.high).reshape(-1) for s in spaces]
+            low = jnp.concatenate(lows, axis=0)
+            high = jnp.concatenate(highs, axis=0)
+            space = Continuous(low=low, high=high)
+        elif obs_cls == Discrete:
+            ns = [jnp.asarray(s.n).reshape(-1) for s in spaces]
+            n = jnp.concatenate(ns, axis=0)
+            space = Discrete(n=n)
+        else:
+            raise ValueError(f"Unsupported space type: {obs_cls}")
+
+        for batch_dim in batch_dims:
+            space = BatchedSpace(space, batch_dim)
+        return space

envelope/wrappers/normalization.py

@@ -16,7 +16,7 @@ class MeanVarPair(NamedTuple):
 class RunningMeanVar(FrozenPyTreeNode):
     mean: PyTree
     var: PyTree
-    count: int
+    count: int | Array
 
     @cached_property
     def std(self) -> PyTree:

envelope/wrappers/observation_normalization_wrapper.py

@@ -1,10 +1,11 @@
-from typing import override
+from functools import cached_property
+from typing import cast, override
 
 import jax
 from jax import numpy as jnp
 
 from envelope.environment import Info
-from envelope.spaces import BatchedSpace, PyTreeSpace, Space
+from envelope.spaces import BatchedSpace, Continuous, Discrete, PyTreeSpace, Space
 from envelope.struct import field, static_field
 from envelope.typing import Key, PyTree
 from envelope.wrappers.normalization import RunningMeanVar, update_rmv
@@ -36,7 +37,7 @@ class ObservationNormalizationWrapper(Wrapper):
         mean = jax.tree.map(zeros, self.stats_spec)
         var = jax.tree.map(ones, self.stats_spec)
 
-        return RunningMeanVar(mean=mean, var=var, count=0)
+        return RunningMeanVar(mean=mean, var=var, count=jnp.asarray(0))
 
     def _normalize_obs(self, obs: PyTree, rmv: RunningMeanVar) -> PyTree:
         def norm_leaf(x, mean, std, spec):
@@ -66,29 +67,40 @@ class ObservationNormalizationWrapper(Wrapper):
         return state, info
 
     @override
-    def reset(
-        self, key: Key, state: PyTree | None = None, **kwargs
-    ) -> tuple[WrappedState, Info]:
-        inner_state = None
+    def init(self, key: Key) -> tuple[WrappedState, Info]:
+        inner_state, info = self.env.init(key)
         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)
+    def reset(self, key: Key, state: WrappedState) -> tuple[WrappedState, Info]:
+        inner_state, info = self.env.reset(key, state.inner_state)
+        # Preserve running statistics across resets
+        next_state = self.ObservationNormalizationState(
+            inner_state=inner_state, rmv_state=state.rmv_state
+        )
+        return self._normalize_and_update(next_state, info)
+
+    @override
+    def step(self, state: WrappedState, action: PyTree) -> tuple[WrappedState, Info]:
+        inner_state, info = self.env.step(state.inner_state, action)
         state = state.replace(inner_state=inner_state)
         return self._normalize_and_update(state, info)
 
+    @override
+    @cached_property
+    def observation_space(self) -> Space:
+        def to_continuous(space: Continuous | Discrete) -> Continuous:
+            return Continuous.from_shape(low=-jnp.inf, high=jnp.inf, shape=space.shape)
+
+        def is_leaf(space: Space) -> bool:
+            return isinstance(space, (Discrete, Continuous))
+
+        return jax.tree.map(to_continuous, self.env.observation_space, is_leaf=is_leaf)
+
 
 def _infer_stats_spec(space: Space) -> PyTree:
     """

envelope/wrappers/pooled_init_vmap_wrapper.py

@@ -0,0 +1,122 @@
+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.typing import Key, PyTree
+from envelope.wrappers.vmap_wrapper import _split_or_keep_key
+from envelope.wrappers.wrapper import WrappedState, Wrapper
+
+
+class PooledInitVmapWrapper(Wrapper):
+    batch_size: int = field(kw_only=True)
+    pool_size: int = field(kw_only=True)
+
+    class PooledInitVmapState(WrappedState):
+        init_key: Key = field()
+        last_final: Info = field()
+
+    @override
+    def init(self, key: Key) -> tuple[WrappedState, Info]:
+        keys = _split_or_keep_key(key, self.batch_size + 1)
+        key_next, keys_pool = keys[0], keys[1:]
+        inner_state, info = jax.vmap(self.env.init)(keys_pool)
+        pholder_info = jax.tree.map(
+            lambda x: jnp.full_like(x, jnp.nan, dtype=jnp.float32), info
+        )
+        state = self.PooledInitVmapState(
+            inner_state=inner_state,
+            init_key=key_next,
+            last_final=pholder_info,
+        )
+        return state, info.update(final=pholder_info)
+
+    @override
+    def reset(self, key: Key, state: WrappedState) -> tuple[WrappedState, Info]:
+        # It's hard to support reset for this wrapper.
+        # We would have to init the state of a pool of unwrapped environments, and then
+        # somehow inject this into the stack of wrapped states. The current data
+        # structure for wrapped states does not make this possible without being super
+        # hacky, and violating the assumption that wrapped states are opaque (we would
+        # likely have to recursively descend by checking if
+        # hasattr(state, "inner_state")).
+        # Since there is currently no use case in which we need to carry state across
+        # episodes before vmapping, we will implement this later.
+        keys = _split_or_keep_key(key, self.batch_size + 1)
+        key_next, keys_pool = keys[0], keys[1:]
+        inner_state, info = jax.vmap(self.env.reset)(keys_pool, state.inner_state)
+        state = state.replace(inner_state=inner_state, init_key=key_next)
+        return state, info.update(final=state.last_final)
+
+    @override
+    def step(self, state: WrappedState, action: PyTree) -> tuple[WrappedState, Info]:
+        inner_state, info = jax.vmap(self.env.step)(state.inner_state, action)
+        done = info.terminated | info.truncated
+
+        # Compute pool_size fresh init states
+        key_pool = jax.random.fold_in(state.init_key, 0)
+        next_init_key = jax.random.fold_in(state.init_key, 1)
+        keys_pool = jax.random.split(key_pool, self.pool_size)
+        inner_states_pool, infos_pool = jax.vmap(self.env.init)(keys_pool)
+
+        # Randomly assign each env a init state from the pool
+        key_idxs = jax.random.fold_in(state.init_key, 2)
+        pool_idxs = jax.random.randint(key_idxs, (self.batch_size,), 0, self.pool_size)
+
+        # Expand pool states to batch_size via indexing
+        mapped_init_state = jax.tree.map(lambda x: x[pool_idxs], inner_states_pool)
+        mapped_init_info = jax.tree.map(lambda x: x[pool_idxs], infos_pool)
+
+        # Select inner_state: init for done envs, continue for others
+        final_inner_state = jax.tree.map(
+            lambda init, curr: jax.vmap(jnp.where)(done, init, curr),
+            mapped_init_state,
+            inner_state,
+        )
+
+        # Select last_final: on done, store terminal info; on continue, keep previous
+        final_last_final = jax.tree.map(
+            lambda curr, prev: jax.vmap(jnp.where)(done, curr, prev),
+            info,
+            state.last_final,
+        )
+
+        # Build final_info with final field
+        # For done envs: obs is new initial obs, final is terminal info
+        # For continue envs: obs is current obs, final is previous last_final
+        final_obs = jax.tree.map(
+            lambda init, curr: jax.vmap(jnp.where)(done, init, curr),
+            mapped_init_info.obs,
+            info.obs,
+        )
+        final_final = jax.tree.map(
+            lambda curr, prev: jax.vmap(jnp.where)(done, curr, prev),
+            info,  # Terminal info snapshot for done envs
+            state.last_final,  # Previous episode's final for continue envs
+        )
+        final_info = info.update(obs=final_obs, final=final_final)
+
+        state = state.replace(
+            inner_state=final_inner_state,
+            init_key=next_init_key,
+            last_final=final_last_final,
+        )
+        return state, final_info
+
+    @override
+    @cached_property
+    def observation_space(self) -> spaces.Space:
+        return spaces.BatchedSpace(
+            space=self.env.observation_space, batch_size=self.batch_size
+        )
+
+    @override
+    @cached_property
+    def action_space(self) -> spaces.Space:
+        return spaces.BatchedSpace(
+            space=self.env.action_space, batch_size=self.batch_size
+        )

envelope/wrappers/state_injection_wrapper.py

@@ -1,3 +1,5 @@
+from typing import override
+
 from envelope.environment import Info, InfoContainer
 from envelope.struct import field
 from envelope.typing import Key, PyTree
@@ -12,7 +14,7 @@ class StateInjectionWrapper(Wrapper):
 
     Usage:
         env = AutoResetWrapper(StateInjectionWrapper(env=base_env))
-        state, info = env.reset(key)
+        state, info = env.init(key)
 
         for outer_iter in range(num_outer_iters):
             # Sample a new task and set it as the reset state
@@ -60,32 +62,26 @@ class StateInjectionWrapper(Wrapper):
 
         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)
+    @override
+    def init(self, key: Key) -> tuple[WrappedState, Info]:
+        inner_state, info = self.env.init(key)
+        state = self.InjectedState(inner_state=inner_state)
+        return state, info
 
-        # If reset state is set, use it
-        elif state.reset_state is not None and state.reset_obs is not None:
+    @override
+    def reset(self, key: Key, state: WrappedState) -> tuple[WrappedState, Info]:
+        # If reset state is set, use it instead of resetting inner env
+        if 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
+        elif state.reset_state is None and state.reset_obs is None:
+            inner_state, info = self.env.reset(key, state.inner_state)
         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
+        return state.replace(inner_state=inner_state), info
 
-    def step(
-        self, state: WrappedState, action: PyTree, **kwargs
-    ) -> tuple[WrappedState, Info]:
-        inner_state, info = self.env.step(state.inner_state, action, **kwargs)
+    @override
+    def step(self, state: WrappedState, action: PyTree) -> tuple[WrappedState, Info]:
+        inner_state, info = self.env.step(state.inner_state, action)
         return state.replace(inner_state=inner_state), info

envelope/wrappers/truncation_wrapper.py

@@ -1,3 +1,5 @@
+from typing import override
+
 import jax.numpy as jnp
 
 from envelope.environment import Info
@@ -12,20 +14,22 @@ class TruncationWrapper(Wrapper):
     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)
+    @override
+    def init(self, key: Key) -> tuple[WrappedState, Info]:
+        inner_state, info = self.env.init(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)
+    @override
+    def reset(self, key: Key, state: WrappedState) -> tuple[WrappedState, Info]:
+        inner_state, info = self.env.reset(key, state.inner_state)
+        state = state.replace(inner_state=inner_state, steps=0)
+        return state, info.update(truncated=self.max_steps <= 0)
+
+    @override
+    def step(self, state: WrappedState, action: PyTree) -> tuple[WrappedState, Info]:
+        next_inner_state, info = self.env.step(state.inner_state, action)
+        steps = state.steps + 1
+        state = self.TruncationState(inner_state=next_inner_state, steps=steps)
+        truncated = jnp.asarray(steps) >= self.max_steps
+        return state, info.update(truncated=truncated)