env-ssl-wrapper 0.4.2__tar.gz → 0.4.4__tar.gz

{env_ssl_wrapper-0.4.2 → env_ssl_wrapper-0.4.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: env-ssl-wrapper
-Version: 0.4.2
+Version: 0.4.4
 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
@@ -134,6 +134,55 @@ 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'`.
+
+Set `reward_mode = 'chunk'` to get the raw per-step rewards back as a reward chunk of shape `(num_envs, chunk_length)` (`(chunk_length,)` for a single env) instead of a reduced scalar — useful when a learner wants to do per-step credit assignment within the chunk. The chunk is truncated to the number of substeps actually executed (`info['chunk_length']`) and is undiscounted, so `gamma` only affects the aggregated modes.
+
+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.
@@ -148,6 +197,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.2 → env_ssl_wrapper-0.4.4}/README.md

@@ -71,6 +71,55 @@ 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'`.
+
+Set `reward_mode = 'chunk'` to get the raw per-step rewards back as a reward chunk of shape `(num_envs, chunk_length)` (`(chunk_length,)` for a single env) instead of a reduced scalar — useful when a learner wants to do per-step credit assignment within the chunk. The chunk is truncated to the number of substeps actually executed (`info['chunk_length']`) and is undiscounted, so `gamma` only affects the aggregated modes.
+
+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 +134,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.2 → env_ssl_wrapper-0.4.4}/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.4/env_ssl_wrapper/action_chunk.py

@@ -0,0 +1,119 @@
+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', 'chunk'), 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 == 'chunk':
+            reward = rewards
+        elif 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.2 → env_ssl_wrapper-0.4.4}/env_ssl_wrapper/standardize/standardize_env_wrapper.py

@@ -23,6 +23,9 @@ class StandardizeEnvWrapper(EnvWrapper):
         image_size: int | tuple[int, int] | None = None,
         lambdas: tuple[float, ...] | list[float] | None = None,
         keys: str | tuple[str, ...] | None = None,
+        chunk_len: int | None = None,
+        chunk_gamma: float = 1.,
+        chunk_reward_mode: str = 'sum',
     ):
         wrappers = []
 
@@ -48,6 +51,9 @@ class StandardizeEnvWrapper(EnvWrapper):
         if exists(lambdas):
             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.2 → env_ssl_wrapper-0.4.4}/env_ssl_wrapper/standardize/utils.py

@@ -34,10 +34,14 @@ def get_wrapper(name):
         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", "standardize_env"])}')
+    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):
@@ -87,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.2 → env_ssl_wrapper-0.4.4}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "env-ssl-wrapper"
-version = "0.4.2"
+version = "0.4.4"
 description = "One torch-native interface for any MDP environment"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }