env-ssl-wrapper 0.4.1__tar.gz → 0.4.3__tar.gz

{env_ssl_wrapper-0.4.1 → env_ssl_wrapper-0.4.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: env-ssl-wrapper
-Version: 0.4.1
+Version: 0.4.3
 Summary: One torch-native interface for any MDP environment
 Project-URL: Homepage, https://pypi.org/project/env-ssl-wrapper/
 Project-URL: Repository, https://codeberg.org/lucidrains/env-ssl-wrapper
@@ -47,14 +47,17 @@ Requires-Dist: fire; extra == 'examples'
 Provides-Extra: test
 Requires-Dist: box2d-py>=2.3.5; extra == 'test'
 Requires-Dist: dm-control; extra == 'test'
+Requires-Dist: gym-pusht; extra == 'test'
 Requires-Dist: gymnasium-robotics; extra == 'test'
 Requires-Dist: gymnasium>=0.29.1; extra == 'test'
 Requires-Dist: imageio; extra == 'test'
+Requires-Dist: mean-conc-beta; extra == 'test'
 Requires-Dist: mujoco; extra == 'test'
 Requires-Dist: pillow; extra == 'test'
 Requires-Dist: populora; extra == 'test'
 Requires-Dist: pybullet; extra == 'test'
 Requires-Dist: pygame; extra == 'test'
+Requires-Dist: pymunk<7; extra == 'test'
 Requires-Dist: pytest; extra == 'test'
 Description-Content-Type: text/markdown
 
@@ -81,7 +84,7 @@ env = StandardizeEnvWrapper(any_env)
 obs, info = env.reset() # torch.float32, batched
 
 while not env.all_done:
-    actions = torch.randint(0, 2, (8,))
+    actions = torch.randint(0, 2, (1,))
     obs, reward, terminated, truncated, info = env.step(actions)
 ```
 
@@ -131,6 +134,53 @@ Run the PPO benchmark on POMDP LunarLander:
 uv run test_memory_trace.py
 ```
 
+### Action Chunk
+
+Open-loop action chunking: one `step` executes a whole chunk of actions and returns only the final state, temporally compressing the environment by the chunk length.
+
+```python
+import gymnasium as gym
+import torch
+from env_ssl_wrapper import StandardizeEnvWrapper, ActionChunkWrapper
+
+env = StandardizeEnvWrapper(gym.make('CartPole-v1'))
+env = ActionChunkWrapper(env, chunk_len = 2, gamma = 0.99)
+
+obs, info = env.reset()
+
+while True:
+    actions = torch.randint(0, 2, (1, 2))  # (num_envs, chunk_len)
+    obs, reward, terminated, truncated, info = env.step(actions)
+
+    # reward                -> discounted sum of substeps (r0 + gamma * r1 + ...)
+    # info['discount']      -> gamma ** chunk_length (macro-step discount for GAE / Bellman target)
+    # info['chunk_length']  -> env steps actually executed (drops below chunk_len on terminal chunk)
+    # info['chunk_rewards'] -> (1, chunk_length) raw per-step rewards
+
+    if bool((terminated | truncated).item()):
+        obs, info = env.reset()
+```
+
+Chunks are shaped `(num_envs, chunk_len, *action_shape)` (`(num_envs, chunk_len)` for discrete actions).
+
+Execution stops early the moment any env terminates or truncates mid-chunk — the terminal state and done flags of that substep are returned and the rest of the chunk is dropped, so a new episode is never silently advanced.
+
+Pass `gamma` (default `1.`) to discount intra-chunk rewards $r = \sum_{i=0}^{L-1} \gamma^i r_i$. The macro-transition discount factor to the next state is provided as `info['discount'] = gamma ** chunk_length`. `reward_mode` can also be `'mean'` or `'last'`.
+
+Can also be passed directly to `StandardizeEnvWrapper` or `compose_env`:
+
+```python
+env = StandardizeEnvWrapper(gym.make('CartPole-v1'), chunk_len = 2, chunk_gamma = 0.99)
+```
+
+Run the chunked PPO benchmark on CartPole, or check chunked rollouts against a per-step reference:
+
+```bash
+uv run test_action_chunk.py --chunk_len=2                # PPO on CartPole
+uv run test_action_chunk.py --sweep=True                 # PPO over several chunk lengths
+uv run test_action_chunk.py --verify=True --sweep=True   # chunked rollouts match per-step reference
+```
+
 ## Wrappers
 
 Pass wrappers as strings (default config) or `(name, dict)` tuples (custom config), in any order.
@@ -145,6 +195,7 @@ Pass wrappers as strings (default config) or `(name, dict)` tuples (custom confi
 | `action_transform` | Rescales actions from a canonical `(0, 1)` range to the env's bounds. |
 | `tensor` | NumPy → torch on a device, torch actions → numpy for the sim. |
 | `flatten_obs` | Flattens dict/tuple observations into a single vector. |
+| `action_chunk` | Executes actions in open-loop chunks of length k. `('action_chunk', dict(chunk_len = 2, gamma = 0.99))` |
 
 Every env emits the same contract: obs `torch.float32`, rewards `torch.float32`, `terminated`/`truncated` `torch.bool`. `env.seed(n)` works on every sim.
 

{env_ssl_wrapper-0.4.1 → env_ssl_wrapper-0.4.3}/README.md

@@ -21,7 +21,7 @@ env = StandardizeEnvWrapper(any_env)
 obs, info = env.reset() # torch.float32, batched
 
 while not env.all_done:
-    actions = torch.randint(0, 2, (8,))
+    actions = torch.randint(0, 2, (1,))
     obs, reward, terminated, truncated, info = env.step(actions)
 ```
 
@@ -71,6 +71,53 @@ Run the PPO benchmark on POMDP LunarLander:
 uv run test_memory_trace.py
 ```
 
+### Action Chunk
+
+Open-loop action chunking: one `step` executes a whole chunk of actions and returns only the final state, temporally compressing the environment by the chunk length.
+
+```python
+import gymnasium as gym
+import torch
+from env_ssl_wrapper import StandardizeEnvWrapper, ActionChunkWrapper
+
+env = StandardizeEnvWrapper(gym.make('CartPole-v1'))
+env = ActionChunkWrapper(env, chunk_len = 2, gamma = 0.99)
+
+obs, info = env.reset()
+
+while True:
+    actions = torch.randint(0, 2, (1, 2))  # (num_envs, chunk_len)
+    obs, reward, terminated, truncated, info = env.step(actions)
+
+    # reward                -> discounted sum of substeps (r0 + gamma * r1 + ...)
+    # info['discount']      -> gamma ** chunk_length (macro-step discount for GAE / Bellman target)
+    # info['chunk_length']  -> env steps actually executed (drops below chunk_len on terminal chunk)
+    # info['chunk_rewards'] -> (1, chunk_length) raw per-step rewards
+
+    if bool((terminated | truncated).item()):
+        obs, info = env.reset()
+```
+
+Chunks are shaped `(num_envs, chunk_len, *action_shape)` (`(num_envs, chunk_len)` for discrete actions).
+
+Execution stops early the moment any env terminates or truncates mid-chunk — the terminal state and done flags of that substep are returned and the rest of the chunk is dropped, so a new episode is never silently advanced.
+
+Pass `gamma` (default `1.`) to discount intra-chunk rewards $r = \sum_{i=0}^{L-1} \gamma^i r_i$. The macro-transition discount factor to the next state is provided as `info['discount'] = gamma ** chunk_length`. `reward_mode` can also be `'mean'` or `'last'`.
+
+Can also be passed directly to `StandardizeEnvWrapper` or `compose_env`:
+
+```python
+env = StandardizeEnvWrapper(gym.make('CartPole-v1'), chunk_len = 2, chunk_gamma = 0.99)
+```
+
+Run the chunked PPO benchmark on CartPole, or check chunked rollouts against a per-step reference:
+
+```bash
+uv run test_action_chunk.py --chunk_len=2                # PPO on CartPole
+uv run test_action_chunk.py --sweep=True                 # PPO over several chunk lengths
+uv run test_action_chunk.py --verify=True --sweep=True   # chunked rollouts match per-step reference
+```
+
 ## Wrappers
 
 Pass wrappers as strings (default config) or `(name, dict)` tuples (custom config), in any order.
@@ -85,6 +132,7 @@ Pass wrappers as strings (default config) or `(name, dict)` tuples (custom confi
 | `action_transform` | Rescales actions from a canonical `(0, 1)` range to the env's bounds. |
 | `tensor` | NumPy → torch on a device, torch actions → numpy for the sim. |
 | `flatten_obs` | Flattens dict/tuple observations into a single vector. |
+| `action_chunk` | Executes actions in open-loop chunks of length k. `('action_chunk', dict(chunk_len = 2, gamma = 0.99))` |
 
 Every env emits the same contract: obs `torch.float32`, rewards `torch.float32`, `terminated`/`truncated` `torch.bool`. `env.seed(n)` works on every sim.
 

{env_ssl_wrapper-0.4.1 → env_ssl_wrapper-0.4.3}/env_ssl_wrapper/__init__.py

@@ -5,6 +5,7 @@ from . import standardize
 from .standardize import *
 
 from .memory_trace import MemoryTraceWrapper
+from .action_chunk import ActionChunkWrapper
 
 # Wire backwards-compatibility aliases in sys.modules and module globals
 # so imports like `from env_ssl_wrapper.done_tracker_wrapper import DoneTrackerWrapper`
@@ -38,6 +39,7 @@ for _name in _STANDARDIZE_SUBMODULES:
 __all__ = [
     *standardize.__all__,
     'MemoryTraceWrapper',
+    'ActionChunkWrapper',
 ]
 
 def __getattr__(name):

env_ssl_wrapper-0.4.3/env_ssl_wrapper/action_chunk.py

@@ -0,0 +1,117 @@
+from __future__ import annotations
+
+import numpy as np
+import torch
+from einops import reduce
+from torch import is_tensor
+
+from .standardize.helpers import (
+    EnvWrapper,
+    exists,
+    default,
+    any_true,
+    dones_of,
+    get_attr,
+)
+
+# helpers
+
+def stack_steps(steps):
+    return torch.stack(steps, dim = -1) if is_tensor(steps[0]) else np.stack(steps, axis = -1)
+
+def chunk_steps(actions, axis):
+    return actions.unbind(dim = axis) if is_tensor(actions) else np.moveaxis(actions, axis, 0)
+
+# wrapper
+
+class ActionChunkWrapper(EnvWrapper):
+    """
+    Open-loop action chunking - executes a chunk of actions in one step,
+    temporally compressing the environment by the chunk length.
+
+    action chunks are shaped (num_envs, chunk_len, *action_shape)
+    """
+
+    def __init__(
+        self,
+        env,
+        chunk_len: int,
+        gamma: float = 1.,
+        discount: float | None = None,
+        reward_mode: str = 'sum'
+    ):
+        super().__init__(env)
+
+        gamma = default(discount, gamma)
+        assert chunk_len >= 1, f'chunk_len must be at least 1, got {chunk_len}'
+        assert 0. <= gamma <= 1., f'gamma must be between 0 and 1, got {gamma}'
+        assert reward_mode in ('sum', 'mean', 'last'), f'unknown reward_mode {reward_mode!r}'
+
+        self.chunk_len = chunk_len
+        self.gamma = float(gamma)
+        self.reward_mode = reward_mode
+
+        action_space = get_attr(env, 'action_space')
+        self.action_shape = tuple(get_attr(action_space, 'shape', ()) or ())
+
+    @property
+    def chunk_action_shape(self):
+        return (self.chunk_len, *self.action_shape)
+
+    def reset(self, **kwargs):
+        return self.env.reset(**kwargs)
+
+    def step(self, actions):
+        if isinstance(actions, dict):
+            raise NotImplementedError('action chunking does not support dict actions yet')
+
+        if not (is_tensor(actions) or isinstance(actions, np.ndarray)):
+            actions = np.asarray(actions)
+
+        chunk_axis = actions.ndim - 1 - len(self.action_shape)
+
+        assert chunk_axis >= 1, (
+            f'action chunks must be shaped (num_envs, chunk_len, *action_shape), got {tuple(actions.shape)}'
+        )
+
+        assert actions.shape[chunk_axis] == self.chunk_len, (
+            f'expected chunk length {self.chunk_len}, got {actions.shape[chunk_axis]}'
+        )
+
+        rewards = []
+        out = None
+
+        for action in chunk_steps(actions, chunk_axis):
+            out = self.env.step(action)
+            rewards.append(out[1])
+
+            if any_true(dones_of(out[2], out[3])):
+                break
+
+        obs, _, terminated, truncated, info = out
+
+        rewards = stack_steps(rewards)
+        executed_len = rewards.shape[-1]
+
+        if self.reward_mode == 'last':
+            reward = rewards[..., -1]
+        elif self.reward_mode == 'mean':
+            reward = reduce(rewards, '... k -> ...', 'mean')
+        elif self.gamma != 1.:
+            if is_tensor(rewards):
+                dtype = rewards.dtype if rewards.is_floating_point() else torch.float32
+                discounts = (self.gamma ** torch.arange(executed_len, device = rewards.device)).to(dtype = dtype)
+            else:
+                dtype = rewards.dtype if np.issubdtype(rewards.dtype, np.floating) else np.float32
+                discounts = (self.gamma ** np.arange(executed_len)).astype(dtype)
+
+            reward = reduce(rewards * discounts, '... k -> ...', 'sum')
+        else:
+            reward = reduce(rewards, '... k -> ...', 'sum')
+
+        info = dict(info) if isinstance(info, dict) else {}
+        info['chunk_length'] = executed_len
+        info['chunk_rewards'] = rewards
+        info['discount'] = self.gamma ** executed_len
+
+        return obs, reward, terminated, truncated, info

{env_ssl_wrapper-0.4.1 → env_ssl_wrapper-0.4.3}/env_ssl_wrapper/standardize/action_transform_wrapper.py

@@ -97,9 +97,14 @@ class ActionTransformWrapper(EnvWrapper):
         if was_scalar:
             t = rearrange(t, '-> 1')
 
-        dim = t.shape[-1]
-        low = np.broadcast_to(low, (dim,))
-        high = np.broadcast_to(high, (dim,))
+        # bounds may be unbatched (single_action_space) or already batched
+        # (vector envs exposing only a batched action_space)
+
+        if low.shape != t.shape and not (low.ndim > 0 and t.shape[-low.ndim:] == low.shape):
+            dim = t.shape[-1]
+            low = np.broadcast_to(low, (dim,))
+            high = np.broadcast_to(high, (dim,))
+
         valid = np.isfinite(low) & np.isfinite(high)
 
         if is_tensor(t):

{env_ssl_wrapper-0.4.1 → env_ssl_wrapper-0.4.3}/env_ssl_wrapper/standardize/adapters.py

@@ -1,32 +1,20 @@
 from __future__ import annotations
 
-import numpy as np
-import torch
-from torch import is_tensor
-
 from .helpers import (
     EnvWrapper,
     default,
     exists,
     first_existing,
     get_attr,
+    is_time_step,
+    normalize_reset_out,
+    normalize_step_out,
     safe_close,
     truthy_attr,
+    zero_like,
 )
 from .spaces import space_from_action_spec
 
-# helpers
-
-def zero_like(x):
-    if is_tensor(x):
-        return torch.zeros_like(x, dtype = torch.bool)
-
-    arr = np.asarray(x)
-    return np.zeros_like(arr, dtype = bool) if arr.ndim > 0 else False
-
-def is_time_step(out):
-    return exists(get_attr(out, 'step_type')) and exists(get_attr(out, 'observation'))
-
 # base adapter
 
 class BaseEnvAdapter:
@@ -38,27 +26,10 @@ class BaseEnvAdapter:
         self.env = env
 
     def step(self, action) -> tuple:
-        out = self.env.step(action)
-        if is_time_step(out):
-            last = out.last() if callable(get_attr(out, 'last')) else out.step_type == 2
-            return out.observation, out.reward, last, False, dict(discount = out.discount)
-        if len(out) == 5:
-            obs, reward, term, trunc, info = out
-            return obs, reward, term, trunc, info if isinstance(info, dict) else {}
-        if len(out) in (3, 4):
-            obs, reward, done, *rest = out
-            info = rest[0] if rest and isinstance(rest[0], dict) else {}
-            return obs, reward, done, zero_like(done), info
-        raise ValueError(f'cannot standardize step output of length {len(out)}')
+        return normalize_step_out(self.env.step(action))
 
     def reset(self, **kwargs) -> tuple:
-        out = self.env.reset(**kwargs)
-        if is_time_step(out):
-            return out.observation, {}
-        if isinstance(out, tuple) and len(out) == 2:
-            obs, info = out
-            return obs, {} if info is None else (info if isinstance(info, dict) else {})
-        return out, {}
+        return normalize_reset_out(self.env.reset(**kwargs))
 
     def seed(self, seed: int):
         if callable(get_attr(self.env, 'seed')):
@@ -381,15 +352,16 @@ class GymnasiumAdapter(BaseEnvAdapter):
 
     @property
     def autoresets(self) -> bool:
-        if truthy_attr(first_existing(self.env, 'autoreset', 'autoresets', 'autoreset_mode')):
-            return True
+        mode = first_existing(self.env, 'autoreset', 'autoresets', 'autoreset_mode')
+
+        if exists(mode):
+            return truthy_attr(mode)
+
         try:
             from gymnasium.vector import VectorEnv
-            if isinstance(self.env, VectorEnv):
-                return True
+            return isinstance(self.env, VectorEnv)
         except ImportError:
-            pass
-        return False
+            return False
 
     def seed(self, seed: int):
         try:

{env_ssl_wrapper-0.4.1 → env_ssl_wrapper-0.4.3}/env_ssl_wrapper/standardize/done_tracker_wrapper.py

@@ -2,8 +2,6 @@ from __future__ import annotations
 
 import numpy as np
 
-from torch.utils._pytree import tree_flatten
-
 from .auto_batched_wrapper import AutoBatchedWrapper
 from .helpers import (
     EnvWrapper,
@@ -11,27 +9,11 @@ from .helpers import (
     env_autoresets,
     env_num_envs,
     exists,
-    get_attr,
-    is_array,
+    get_batch_size,
     is_vectorized,
     to_numpy,
 )
 
-# helper functions
-
-def get_batch_size(tree) -> int | None:
-    leaves, _ = tree_flatten(tree)
-
-    if not leaves:
-        return None
-
-    first = leaves[0]
-
-    if is_array(first):
-        return len(first) if first.ndim > 0 else None
-
-    return len(first) if exists(get_attr(first, '__len__')) else None
-
 # classes
 
 class DoneTrackerWrapper(EnvWrapper):
@@ -113,6 +95,6 @@ class DoneTrackerWrapper(EnvWrapper):
             info['episode_lengths'] = self.episode_lengths.copy()
 
             if self.all_done:
-                info.update(needs_reset = True, all_done = True)
+                info.update(needs_reset = self.needs_reset, all_done = True)
 
         return obs, reward, terminated, truncated, info

{env_ssl_wrapper-0.4.1 → env_ssl_wrapper-0.4.3}/env_ssl_wrapper/standardize/episode_padding_wrapper.py

@@ -21,15 +21,18 @@ from .helpers import (
 
 # helpers
 
-def zero_mask(x, mask, fill_scalar = None):
-    if is_tensor(x):
-        m = torch.as_tensor(mask, device = x.device, dtype = torch.bool)
-        diff = x.ndim - m.ndim
+def broadcast_mask(mask, x):
+    if not is_tensor(x):
+        return mask
+
+    m = torch.as_tensor(mask, device = x.device, dtype = torch.bool)
+    diff = x.ndim - m.ndim
 
-        if diff > 0:
-            m = rearrange(m, f'... -> ... {" ".join(["1"] * diff)}')
+    return rearrange(m, f'... -> ... {" ".join(["1"] * diff)}') if diff > 0 else m
 
-        return torch.where(m, torch.zeros_like(x), x)
+def zero_mask(x, mask, fill_scalar = None):
+    if is_tensor(x):
+        return torch.where(broadcast_mask(mask, x), torch.zeros_like(x), x)
 
     arr = np.asarray(x)
 
@@ -49,13 +52,7 @@ def back_to_mask_type(dones, newly):
 
 def merge_final(current, value, mask):
     if is_tensor(current):
-        m = torch.as_tensor(mask, device = current.device, dtype = torch.bool)
-        diff = current.ndim - m.ndim
-
-        if diff > 0:
-            m = rearrange(m, f'... -> ... {" ".join(["1"] * diff)}')
-
-        return torch.where(m, value, current)
+        return torch.where(broadcast_mask(mask, current), value, current)
 
     curr = np.asarray(current)
 
@@ -97,15 +94,15 @@ class EpisodePaddingWrapper(EnvWrapper):
             if self._is_done is None or len(self._is_done) != len(mask):
                 self._is_done = np.zeros(len(mask), dtype = bool)
 
-            if self.autoreset:
-                self._is_done &= mask
+            # autoreset envs revive every done slot, so each done is a new
+            # terminal transition; non-autoreset slots stay done until reset
+
+            newly = mask if self.autoreset else mask & ~self._is_done
+            self._is_done |= mask
 
             if mask.any():
                 assert exists(self._last_obs), 'environment needs reset before calling step. call env.reset() first'
 
-                newly = mask & ~self._is_done
-                self._is_done |= mask
-
                 if newly.any():
                     final_val = first_existing(info, 'final_observation', 'final_obs')
                     value = final_val if exists(final_val) else self._last_obs

{env_ssl_wrapper-0.4.1 → env_ssl_wrapper-0.4.3}/env_ssl_wrapper/standardize/helpers.py

@@ -22,10 +22,15 @@ def get_attr(obj, name, default = None):
         return default
 
 def truthy_attr(value):
-    # flags arrive as None, methods, numpy scalars — only honest truths count
+    # flags arrive as None, methods, numpy scalars — only honest truths count.
+    # gymnasium's AutoresetMode.DISABLED is a truthy enum that means "off"
+
     if not exists(value) or callable(value):
         return False
 
+    if getattr(value, 'name', None) == 'DISABLED':
+        return False
+
     try:
         return bool(value)
     except Exception:
@@ -46,6 +51,19 @@ def first_existing(obj, *names):
 def is_scalar(v):
     return isinstance(v, (int, float, bool, np.number, np.bool_))
 
+def get_batch_size(tree) -> int | None:
+    leaves, _ = tree_flatten(tree)
+
+    if not leaves:
+        return None
+
+    first = leaves[0]
+
+    if is_array(first):
+        return len(first) if first.ndim > 0 else None
+
+    return len(first) if exists(get_attr(first, '__len__')) else None
+
 def is_array(v):
     return is_tensor(v) or isinstance(v, np.ndarray)
 
@@ -71,6 +89,43 @@ def dones_of(terminated, truncated):
         return terminated | truncated
     return tree_map(lambda a, b: a | b, terminated, truncated)
 
+# sim step / reset normalization
+
+def is_time_step(out):
+    return exists(get_attr(out, 'step_type')) and exists(get_attr(out, 'observation'))
+
+def zero_like(x):
+    if is_tensor(x):
+        return torch.zeros_like(x, dtype = torch.bool)
+
+    arr = np.asarray(x)
+    return np.zeros_like(arr, dtype = bool) if arr.ndim > 0 else False
+
+def normalize_reset_out(out):
+    if is_time_step(out):
+        return out.observation, {}
+
+    if isinstance(out, tuple) and len(out) == 2:
+        obs, info = out
+        return obs, {} if info is None else (info if isinstance(info, dict) else {})
+
+    return out, {}
+
+def normalize_step_out(out):
+    if is_time_step(out):
+        last = out.last() if callable(get_attr(out, 'last')) else out.step_type == 2
+        return out.observation, out.reward, last, False, dict(discount = out.discount)
+
+    if len(out) == 5:
+        return out
+
+    if len(out) in (3, 4):
+        obs, reward, done, *rest = out
+        info = rest[0] if rest and isinstance(rest[0], dict) else {}
+        return obs, reward, done, zero_like(done), info
+
+    raise ValueError(f'could not standardize step output of length {len(out)}')
+
 def _zero_leaf(x):
     if is_tensor(x):
         return torch.zeros_like(x)
@@ -272,16 +327,11 @@ class TransformObservationWrapper(EnvWrapper):
         self.takes_done = accepts_done_param(self.transform_obs)
 
     def transform_obs(self, obs, done = None):
-        if hasattr(self, 'observation'):
-            return self.observation(obs)
         return obs
 
     def transform(self, obs, done = None):
         return self.transform_obs(obs, done = done) if self.takes_done else self.transform_obs(obs)
 
-    def observation(self, obs):
-        return self.transform_obs(obs)
-
     def reset(self, **kwargs):
         obs, info = self.env.reset(**kwargs)
         obs = self.transform_obs(obs)

{env_ssl_wrapper-0.4.1 → env_ssl_wrapper-0.4.3}/env_ssl_wrapper/standardize/image_wrapper.py

@@ -6,8 +6,15 @@ import numpy as np
 
 from einops import rearrange
 
-from .helpers import EnvWrapper, env_render, env_render_mode, exists
-from .standardize_wrapper import normalize_reset_out, normalize_step_out
+from .helpers import (
+    FINAL_OBSERVATION_KEYS,
+    TransformObservationWrapper,
+    env_render,
+    env_render_mode,
+    exists,
+    normalize_reset_out,
+    normalize_step_out,
+)
 
 # helper functions
 
@@ -32,12 +39,19 @@ def render_frame(env, image_size = (64, 64), camera = None):
 
         img = env.render()
 
+    if isinstance(img, (list, tuple)):
+        img = np.ascontiguousarray(img)
+
     if isinstance(img, torch.Tensor):
         img = img.detach().cpu().numpy()
 
     img = np.ascontiguousarray(img)
     img = torch.from_numpy(img)
-    return rearrange(img, 'h w c -> 1 c h w')
+
+    if img.ndim == 4:
+        return rearrange(img, 'b h w c -> b c h w'), True
+
+    return rearrange(img, 'h w c -> 1 c h w'), False
 
 def process_image(
     img,
@@ -61,7 +75,7 @@ def process_image(
 
 # class
 
-class ImageObservationWrapper(EnvWrapper):
+class ImageObservationWrapper(TransformObservationWrapper):
     def __init__(
         self,
         env,
@@ -81,18 +95,21 @@ class ImageObservationWrapper(EnvWrapper):
         self.normalize_divisor = normalize_divisor
 
     def render_frame(self):
-        img = render_frame(self.env, image_size = self.image_size, camera = self.camera)
-        return process_image(
+        img, is_batched = render_frame(self.env, image_size = self.image_size, camera = self.camera)
+        processed = process_image(
             img,
             image_size = self.image_size,
             mode = self.mode,
             normalize = self.normalize,
             normalize_divisor = self.normalize_divisor
         )
+        return processed, is_batched
 
     def observation(self, obs):
-        img = self.render_frame()
-        img = rearrange(img, '1 c h w -> c h w')
+        img, is_batched = self.render_frame()
+
+        if not is_batched:
+            img = rearrange(img, '1 c h w -> c h w')
 
         if not isinstance(obs, dict):
             return dict(state = obs, **{self.image_key: img})
@@ -102,10 +119,23 @@ class ImageObservationWrapper(EnvWrapper):
 
         return {**obs, self.image_key: img}
 
+    def transform_obs(self, obs, done = None):
+        return self.observation(obs)
+
     def reset(self, **kwargs):
         obs, info = normalize_reset_out(self.env.reset(**kwargs))
         return self.observation(obs), info
 
     def step(self, action):
         obs, reward, terminated, truncated, info = normalize_step_out(self.env.step(action))
-        return self.observation(obs), reward, terminated, truncated, info
+        out = self.observation(obs)
+
+        # terminal obs are augmented in place, preserving whatever the sim
+        # or the padding wrapper froze as the true terminal observation
+
+        if isinstance(info, dict) and not self.autoresets:
+            for key in FINAL_OBSERVATION_KEYS:
+                if key in info and not (isinstance(info[key], dict) and self.image_key in info[key]):
+                    info[key] = self.observation(info[key])
+
+        return out, reward, terminated, truncated, info

{env_ssl_wrapper-0.4.1 → env_ssl_wrapper-0.4.3}/env_ssl_wrapper/standardize/standardize_env_wrapper.py

@@ -18,11 +18,14 @@ class StandardizeEnvWrapper(EnvWrapper):
         pad_episodes: bool = True,
         done_tracker: bool = True,
         flatten_obs: bool = False,
-        action_transform: bool = False,
+        action_transform: bool | dict = False,
         max_timesteps: int | None = None,
-        image_size: tuple[int, int] | None = None,
+        image_size: int | tuple[int, int] | None = None,
         lambdas: tuple[float, ...] | list[float] | None = None,
-        **kwargs
+        keys: str | tuple[str, ...] | None = None,
+        chunk_len: int | None = None,
+        chunk_gamma: float = 1.,
+        chunk_reward_mode: str = 'sum',
     ):
         wrappers = []
 
@@ -30,7 +33,8 @@ class StandardizeEnvWrapper(EnvWrapper):
             wrappers.append(('image', dict(image_size = image_size)))
 
         if action_transform:
-            wrappers.append('action_transform')
+            transform_kwargs = dict(auto = True) if action_transform is True else action_transform
+            wrappers.append(('action_transform', transform_kwargs))
 
         if exists(max_timesteps):
             wrappers.append(('time_limit', dict(max_timesteps = max_timesteps)))
@@ -45,7 +49,10 @@ class StandardizeEnvWrapper(EnvWrapper):
             wrappers.append('done_tracker')
 
         if exists(lambdas):
-            wrappers.append(('memory_trace', dict(lambdas = lambdas, **kwargs)))
+            wrappers.append(('memory_trace', dict(lambdas = lambdas, keys = keys)))
+
+        if exists(chunk_len):
+            wrappers.append(('action_chunk', dict(chunk_len = chunk_len, gamma = chunk_gamma, reward_mode = chunk_reward_mode)))
 
         if flatten_obs:
             wrappers.append('flatten_obs')

{env_ssl_wrapper-0.4.1 → env_ssl_wrapper-0.4.3}/env_ssl_wrapper/standardize/standardize_wrapper.py

@@ -1,48 +1,16 @@
 from __future__ import annotations
 
-import numpy as np
-import torch
-from torch import is_tensor
-
-from .adapters import get_adapter, is_time_step, zero_like
+from .adapters import get_adapter
 from .helpers import (
     EnvWrapper,
     default,
     dones_of,
     exists,
-    get_attr,
     instantiate_env,
     mark_terminal_obs,
 )
 from .spaces import infer_observation_space, space_from_action_spec
 
-# helpers
-
-def normalize_reset_out(out):
-    if is_time_step(out):
-        return out.observation, {}
-
-    if isinstance(out, tuple) and len(out) == 2:
-        obs, info = out
-        return obs, {} if info is None else (info if isinstance(info, dict) else {})
-
-    return out, {}
-
-def normalize_step_out(out):
-    if is_time_step(out):
-        last = out.last() if callable(get_attr(out, 'last')) else out.step_type == 2
-        return out.observation, out.reward, last, False, dict(discount = out.discount)
-
-    if len(out) == 5:
-        return out
-
-    if len(out) in (3, 4):
-        obs, reward, done, *rest = out
-        info = rest[0] if rest and isinstance(rest[0], dict) else {}
-        return obs, reward, done, zero_like(done), info
-
-    raise ValueError(f'could not standardize step output of length {len(out)}')
-
 # class
 
 class StandardizeWrapper(EnvWrapper):

{env_ssl_wrapper-0.4.1 → env_ssl_wrapper-0.4.3}/env_ssl_wrapper/standardize/time_limit_wrapper.py

@@ -4,8 +4,18 @@ import numpy as np
 import torch
 from torch import is_tensor
 
-from .helpers import EnvWrapper, dones_of, env_num_envs, is_vectorized, mark_terminal_obs, to_numpy
-from .standardize_wrapper import normalize_reset_out, normalize_step_out
+from .helpers import (
+    EnvWrapper,
+    dones_of,
+    env_num_envs,
+    exists,
+    get_batch_size,
+    is_vectorized,
+    mark_terminal_obs,
+    normalize_reset_out,
+    normalize_step_out,
+    to_numpy,
+)
 
 def back_to_like(t, numpy_arr):
     if is_tensor(t):
@@ -44,8 +54,13 @@ class TimeLimitWrapper(EnvWrapper):
         self.t = np.zeros(self.num_envs, dtype = int)
 
     def reset(self, **kwargs):
+        obs, info = normalize_reset_out(self.env.reset(**kwargs))
+
+        if self.is_vector and exists(batch_size := get_batch_size(obs)):
+            self.num_envs = batch_size
+
         self.t = np.zeros(self.num_envs, dtype = int)
-        return normalize_reset_out(self.env.reset(**kwargs))
+        return obs, info
 
     def step(self, action):
         obs, reward, terminated, truncated, info = normalize_step_out(self.env.step(action))

{env_ssl_wrapper-0.4.1 → env_ssl_wrapper-0.4.3}/env_ssl_wrapper/standardize/utils.py

@@ -25,29 +25,31 @@ WRAPPERS = dict(
     pad_episodes = EpisodePaddingWrapper
 )
 
+def get_wrapper(name):
+    if name in ('standardize_env', 'master'):
+        from .standardize_env_wrapper import StandardizeEnvWrapper
+        return StandardizeEnvWrapper
+
+    if name == 'memory_trace':
+        from ..memory_trace import MemoryTraceWrapper
+        return MemoryTraceWrapper
+
+    if name in ('action_chunk', 'chunk'):
+        from ..action_chunk import ActionChunkWrapper
+        return ActionChunkWrapper
+
+    if name in WRAPPERS:
+        return WRAPPERS[name]
+
+    raise ValueError(f'unknown wrapper {name!r} — choose from {sorted([*WRAPPERS, "memory_trace", "action_chunk", "standardize_env"])}')
+
 def parse_wrapper(wrapper):
     if isinstance(wrapper, str):
-        if wrapper == 'memory_trace':
-            from ..memory_trace import MemoryTraceWrapper
-            wrapper = MemoryTraceWrapper
-        elif wrapper in ('standardize_env', 'master'):
-            from .standardize_env_wrapper import StandardizeEnvWrapper
-            wrapper = StandardizeEnvWrapper
-        elif wrapper not in WRAPPERS:
-            raise ValueError(f'unknown wrapper {wrapper!r} — choose from {sorted([*WRAPPERS, "memory_trace", "standardize_env"])}')
-        else:
-            wrapper = WRAPPERS[wrapper]
+        wrapper = get_wrapper(wrapper)
 
     if isinstance(wrapper, tuple):
         name, kwargs = wrapper
-        if name == 'memory_trace':
-            from ..memory_trace import MemoryTraceWrapper
-            wrapper = partial(MemoryTraceWrapper, **kwargs)
-        elif name in ('standardize_env', 'master'):
-            from .standardize_env_wrapper import StandardizeEnvWrapper
-            wrapper = partial(StandardizeEnvWrapper, **kwargs)
-        else:
-            wrapper = partial(WRAPPERS.get(name, name), **kwargs)
+        wrapper = partial(get_wrapper(name) if isinstance(name, str) else name, **kwargs)
 
     elif isinstance(wrapper, dict):
         raise ValueError("wrapper kwargs must be passed as (name, kwargs), e.g. ('tensor', dict(device = 'cpu'))")
@@ -89,6 +91,14 @@ def compose_env(env, *wrappers, pad_episodes: bool = True):
             funcs.insert(idx_ten + 1, f)
             classes.insert(idx_ten + 1, c)
 
+    # action chunking changes the step signature, so it always goes outermost
+
+    from ..action_chunk import ActionChunkWrapper
+    if ActionChunkWrapper in classes:
+        idx_chunk = classes.index(ActionChunkWrapper)
+        funcs.append(funcs.pop(idx_chunk))
+        classes.append(classes.pop(idx_chunk))
+
     assert len(set(classes)) == len(classes), 'duplicate wrappers found'
 
     for func in funcs:

{env_ssl_wrapper-0.4.1 → env_ssl_wrapper-0.4.3}/env_ssl_wrapper/standardize/vector.py

@@ -22,8 +22,6 @@ from .standardize_wrapper import StandardizeWrapper
 
 # leaf helpers
 
-_stack_trees = stack_trees
-
 def _merge_infos(infos):
     infos = [info if isinstance(info, dict) else {} for info in infos]
 
@@ -38,7 +36,7 @@ def _merge_infos(infos):
     for key in keys:
         vals = [info[key] for info in infos]
         try:
-            out[key] = _stack_trees(vals)
+            out[key] = stack_trees(vals)
         except Exception:
             out[key] = np.array(vals, dtype = object)
 
@@ -276,7 +274,7 @@ class MultiprocessingVecEnv:
             _safe_send(conn, ('reset', kwargs))
 
         results = _recv_all(self._conns)
-        return _stack_trees([obs for obs, _ in results]), _merge_infos([info for _, info in results])
+        return stack_trees([obs for obs, _ in results]), _merge_infos([info for _, info in results])
 
     def step(self, actions):
         actions = _split_actions(actions, self.num_envs)
@@ -302,10 +300,10 @@ class MultiprocessingVecEnv:
                 for final, obs_i in zip(final_obs, obs)
             ]
 
-            info['final_observation'] = _stack_trees(final_obs)
+            info['final_observation'] = stack_trees(final_obs)
             info['_final_observation'] = done
 
-        return _stack_trees(obs), reward, terminated, truncated, info
+        return stack_trees(obs), reward, terminated, truncated, info
 
     def close(self):
         _shutdown(self._conns, self._procs)

{env_ssl_wrapper-0.4.1 → env_ssl_wrapper-0.4.3}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "env-ssl-wrapper"
-version = "0.4.1"
+version = "0.4.3"
 description = "One torch-native interface for any MDP environment"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }
@@ -54,6 +54,9 @@ test = [
     "pillow",
     "imageio",
     "populora",
+    "gym-pusht",
+    "pymunk<7",
+    "mean-conc-beta",
 ]
 
 [tool.pytest.ini_options]