env-ssl-wrapper 0.5.4__tar.gz → 0.5.6__tar.gz

{env_ssl_wrapper-0.5.4 → env_ssl_wrapper-0.5.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: env-ssl-wrapper
-Version: 0.5.4
+Version: 0.5.6
 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
@@ -180,6 +180,26 @@ uv run test_action_chunk.py --sweep=True                 # PPO over several chun
 uv run test_action_chunk.py --verify=True --sweep=True   # chunked rollouts match per-step reference
 ```
 
+### Success-Guided Sampling
+
+Adaptive curriculum over a bank of reset states, keeping episodes near the policy's frontier of competence — the configurations it currently solves about `target` of the time ([Octi Zhang et al., 2026](https://arxiv.org/abs/2603.15999))
+
+```python
+import gymnasium as gym
+import torch
+from env_ssl_wrapper import StandardizeEnvWrapper, SuccessGuidedSamplingWrapper
+
+reset_states = [torch.zeros(4), torch.ones(4)]  # bank of reset states to sample over
+
+env = StandardizeEnvWrapper(gym.make('CartPole-v1'))
+env = SuccessGuidedSamplingWrapper(env, reset_configs = reset_states, target = 0.5, kappa = 2.)
+
+obs, info = env.reset()
+# info['sgs_config_id'] -> the config this episode started from (success rates update on episode end)
+```
+
+`reset_configs` holds reset states — or callables drawing one (`lambda rng: ...`) — passed through `options`, or built by a `reset_options_fn` / `state_transform`. Every stage of the sampling loop is hookable: score equation (`score_fn`), target schedule, uniform floor (`uniform_prob`), env-native fallback (`default_prob`), success signal (`success_fn`).
+
 ## Wrappers
 
 Pass wrappers as strings (default config) or `(name, dict)` tuples (custom config), in any order.
@@ -195,6 +215,7 @@ Pass wrappers as strings (default config) or `(name, dict)` tuples (custom confi
 | `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))` |
+| `sgs` | Adaptive frontier curriculum over reset configs. `('sgs', dict(reset_configs = configs, target = 0.5, kappa = 2.))` |
 
 Every env emits the same contract: obs `torch.float32`, rewards `torch.float32`, `terminated`/`truncated` `torch.bool`. `env.seed(n)` works on every sim.
 
@@ -312,3 +333,12 @@ uv run pytest
     year    = {2025}
 }
 ```
+
+```bibtex
+@article{zhang2026omnireset,
+    title   = {OmniReset: Solving Dexterous Manipulation from Sparse Rewards via Reset States},
+    author  = {Octi Zhang and others},
+    journal = {arXiv preprint arXiv:2603.15999},
+    year    = {2026}
+}
+```

{env_ssl_wrapper-0.5.4 → env_ssl_wrapper-0.5.6}/README.md

@@ -120,6 +120,26 @@ uv run test_action_chunk.py --sweep=True                 # PPO over several chun
 uv run test_action_chunk.py --verify=True --sweep=True   # chunked rollouts match per-step reference
 ```
 
+### Success-Guided Sampling
+
+Adaptive curriculum over a bank of reset states, keeping episodes near the policy's frontier of competence — the configurations it currently solves about `target` of the time ([Octi Zhang et al., 2026](https://arxiv.org/abs/2603.15999))
+
+```python
+import gymnasium as gym
+import torch
+from env_ssl_wrapper import StandardizeEnvWrapper, SuccessGuidedSamplingWrapper
+
+reset_states = [torch.zeros(4), torch.ones(4)]  # bank of reset states to sample over
+
+env = StandardizeEnvWrapper(gym.make('CartPole-v1'))
+env = SuccessGuidedSamplingWrapper(env, reset_configs = reset_states, target = 0.5, kappa = 2.)
+
+obs, info = env.reset()
+# info['sgs_config_id'] -> the config this episode started from (success rates update on episode end)
+```
+
+`reset_configs` holds reset states — or callables drawing one (`lambda rng: ...`) — passed through `options`, or built by a `reset_options_fn` / `state_transform`. Every stage of the sampling loop is hookable: score equation (`score_fn`), target schedule, uniform floor (`uniform_prob`), env-native fallback (`default_prob`), success signal (`success_fn`).
+
 ## Wrappers
 
 Pass wrappers as strings (default config) or `(name, dict)` tuples (custom config), in any order.
@@ -135,6 +155,7 @@ Pass wrappers as strings (default config) or `(name, dict)` tuples (custom confi
 | `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))` |
+| `sgs` | Adaptive frontier curriculum over reset configs. `('sgs', dict(reset_configs = configs, target = 0.5, kappa = 2.))` |
 
 Every env emits the same contract: obs `torch.float32`, rewards `torch.float32`, `terminated`/`truncated` `torch.bool`. `env.seed(n)` works on every sim.
 
@@ -252,3 +273,12 @@ uv run pytest
     year    = {2025}
 }
 ```
+
+```bibtex
+@article{zhang2026omnireset,
+    title   = {OmniReset: Solving Dexterous Manipulation from Sparse Rewards via Reset States},
+    author  = {Octi Zhang and others},
+    journal = {arXiv preprint arXiv:2603.15999},
+    year    = {2026}
+}
+```

{env_ssl_wrapper-0.5.4 → env_ssl_wrapper-0.5.6}/env_ssl_wrapper/__init__.py

@@ -7,6 +7,14 @@ from .standardize import *  # noqa: F403
 
 from .memory_trace import MemoryTraceWrapper
 from .action_chunk import ActionChunkWrapper
+from .sgs import (
+    SuccessGuidedSamplingWrapper,
+    SGSWrapper,
+    SGSSampler,
+    sgs_score_fn,
+    variance_score_fn,
+    uniform_score_fn,
+)
 from .evaluate import evaluate_actor, EpisodeStats
 
 # backwards compatibility — the standardize wrappers used to live at the top
@@ -38,6 +46,12 @@ __all__ = [
     *standardize.__all__,
     'MemoryTraceWrapper',
     'ActionChunkWrapper',
+    'SuccessGuidedSamplingWrapper',
+    'SGSWrapper',
+    'SGSSampler',
+    'sgs_score_fn',
+    'variance_score_fn',
+    'uniform_score_fn',
     'evaluate_actor',
     'EpisodeStats',
 ]

{env_ssl_wrapper-0.5.4 → env_ssl_wrapper-0.5.6}/env_ssl_wrapper/action_chunk.py

@@ -7,15 +7,17 @@ from torch import is_tensor
 
 from .standardize.helpers import (
     EnvWrapper,
+    any_true,
     copy_leaf,
     default,
-    any_true,
     dones_of,
     env_takes_torch,
     first_existing,
     get_adapter,
     get_attr,
     is_vectorized,
+    normalize_reset_out,
+    normalize_step_out,
     to_numpy,
 )
 
@@ -72,7 +74,7 @@ class ActionChunkWrapper(EnvWrapper):
 
         # vectorized gym envs expose a batched action space - the single env space defines the chunk layout
 
-        action_space = first_existing(env, 'single_action_space', 'action_space')
+        action_space = default(first_existing(env, 'single_action_space', 'action_space'), get_adapter(env).action_space)
         self.action_shape = tuple(get_attr(action_space, 'shape', ()) or ())
 
         self.expects_batch = is_vectorized(env)
@@ -83,7 +85,7 @@ class ActionChunkWrapper(EnvWrapper):
         return (self.chunk_len, *self.action_shape)
 
     def reset(self, **kwargs):
-        return self.env.reset(**kwargs)
+        return normalize_reset_out(self.env.reset(**kwargs))
 
     def step(self, actions):
         if isinstance(actions, dict):
@@ -114,7 +116,7 @@ class ActionChunkWrapper(EnvWrapper):
             if not self.expects_batch:
                 action = unbatch_step(action)
 
-            out = self.env.step(action)
+            out = normalize_step_out(self.env.step(action))
             rewards.append(copy_leaf(out[1]))
 
             if any_true(dones_of(out[2], out[3])):

{env_ssl_wrapper-0.5.4 → env_ssl_wrapper-0.5.6}/env_ssl_wrapper/evaluate.py

@@ -6,12 +6,11 @@ from typing import NamedTuple
 import numpy as np
 import torch
 from torch import nn, is_tensor
-from torch.distributions import Categorical, Distribution
-from torch.utils._pytree import tree_flatten, tree_map
+from torch.distributions import Bernoulli, Categorical, Distribution
+from torch.utils._pytree import tree_map
 
 from .standardize.spaces import action_space_is_discrete
 from .standardize.helpers import (
-    EnvWrapper,
     dones_of,
     env_autoresets,
     env_num_envs,
@@ -97,6 +96,12 @@ def to_action(env, out, deterministic = True):
             if not deterministic:
                 return val.sample()
 
+            if isinstance(val, Bernoulli):
+                return (val.probs > 0.5).long()
+
+            if isinstance(val, Categorical):
+                return val.probs.argmax(dim = -1)
+
             probs = get_attr(val, 'probs')
 
             # categorical-style probs carry a trailing action axis to argmax over
@@ -171,24 +176,26 @@ def to_env_action(env, action):
     if not env_takes_torch(env):
         action = tree_map(to_numpy, action)
 
-        # single raw sims get the leading batch dim stripped, while discrete
-        # actions additionally collapse to scalars
+    # single non-vectorized envs get the leading batch dim stripped, while discrete
+    # actions additionally collapse to scalars
 
-        if not is_vectorized(env) and not exists(chunk_shape):
-            def strip_leaf(a):
-                if a.ndim > 1 and a.shape[0] == 1:
-                    a = a[0]
+    if not is_vectorized(env) and not exists(chunk_shape):
+        is_discrete = action_space_is_discrete(get_adapter(env).action_space)
 
-                if action_space_is_discrete(get_adapter(env).action_space):
-                    if a.ndim > 0 and a.shape[0] == 1:
-                        a = a[0]
+        def strip_leaf(a):
+            if a.ndim > 1 and a.shape[0] == 1:
+                a = a[0]
 
-                    if a.ndim == 0:
-                        a = a.item()
+            if is_discrete:
+                if a.ndim > 0 and a.shape[0] == 1:
+                    a = a[0]
+
+                if a.ndim == 0:
+                    a = a.item()
 
-                return a
+            return a
 
-            action = tree_map(strip_leaf, action)
+        action = tree_map(strip_leaf, action)
 
     return action
 
@@ -292,7 +299,7 @@ def evaluate_actor(
     assert not (record and num_envs > 1), 'video recording requires a single env'
 
     obs, info = safe_reset(env, seed)
-    batched = is_vectorized(env) or all(map(is_tensor, tree_flatten(obs)[0]))
+    batched = is_vectorized(env)
 
     goal = goal.expand(num_envs, -1) if exists(goal) and goal.shape[0] == 1 and num_envs > 1 else goal
 
@@ -326,10 +333,12 @@ def evaluate_actor(
             if reward.numel() > num_envs:
                 reward = reward.reshape(num_envs, -1).sum(dim = -1)
 
-            returns += reward if reward.numel() == num_envs else reward.expand(num_envs)
+            reward = reward if reward.numel() == num_envs else reward.expand(num_envs)
+            active_t = torch.from_numpy(active)
 
-            lengths += info.get('chunk_length', 1) if isinstance(info, dict) else 1
-            steps += 1
+            returns += torch.where(active_t, reward, 0.)
+            lengths += torch.where(active_t, float(info.get('chunk_length', 1) if isinstance(info, dict) else 1), 0.)
+            steps += torch.where(active_t, 1, 0)
 
             done = to_numpy(dones_of(terminated, truncated)).astype(bool).reshape(-1)
 
@@ -367,15 +376,18 @@ def evaluate_actor(
                 steps[i] = 0
                 active[i] = False
 
-            if done.any() and len(done_returns) < episodes:
+            if (done.any() or not active.any()) and len(done_returns) < episodes:
                 if record:
                     frames = []
 
                 # autoresetting envs reset themselves, everything else resets here
 
-                if not env_autoresets(env) and (num_envs == 1 or done.all() or bool(get_attr(env, 'needs_reset', False))):
+                if not env_autoresets(env) and (num_envs == 1 or not active.any() or done.all() or bool(get_attr(env, 'needs_reset', False))):
                     reset_seed = seed + len(done_returns) if exists(seed) and num_envs == 1 else None
                     obs, info = safe_reset(env, reset_seed)
+                    returns.zero_()
+                    lengths.zero_()
+                    steps.zero_()
                     active[:] = True
                 elif env_autoresets(env):
                     active[newly_done] = True

{env_ssl_wrapper-0.5.4 → env_ssl_wrapper-0.5.6}/env_ssl_wrapper/memory_trace.py

@@ -1,4 +1,5 @@
 from __future__ import annotations
+from functools import partial
 from typing import Sequence
 
 import torch
@@ -9,9 +10,12 @@ from .standardize.helpers import (
     any_true,
     cast_tuple,
     default,
+    dones_of,
     exists,
     is_array_like,
     is_scalar,
+    maybe_transform_final_observation,
+    normalize_step_out,
 )
 
 # helper functions
@@ -63,7 +67,7 @@ class MemoryTraceWrapper(TransformObservationWrapper):
         prefix = self.trace_key if key == self.obs_key else f'{key}_{self.trace_key}'
         return f'{prefix}_{lam}' if len(self.lambdas) > 1 else prefix
 
-    def transform_obs(self, obs, done = None):
+    def transform_obs(self, obs, done = None, update_state: bool = True):
         if isinstance(obs, (tuple, list)):
             out = {f'{self.obs_key}_{i}': elem for i, elem in enumerate(obs)}
         elif isinstance(obs, dict):
@@ -96,7 +100,8 @@ class MemoryTraceWrapper(TransformObservationWrapper):
             else:
                 traces = [prev.lerp(val_float, calc_lerp_weight(lam, done, val_float)) for lam, prev in zip(self.lambdas, self.traces[key])]
 
-            self.traces[key] = traces
+            if update_state:
+                self.traces[key] = traces
 
             for lam, trace in zip(self.lambdas, traces):
                 out[self.trace_key_for(key, lam)] = trace
@@ -106,3 +111,17 @@ class MemoryTraceWrapper(TransformObservationWrapper):
     def reset(self, **kwargs):
         self.traces = dict()
         return super().reset(**kwargs)
+
+    def step(self, action):
+        obs, reward, terminated, truncated, info = normalize_step_out(self.env.step(action))
+        done = dones_of(terminated, truncated)
+
+        # terminal observations are transformed against pre-reset traces without mutating state
+
+        maybe_transform_final_observation(
+            info,
+            partial(self.transform_obs, done = None, update_state = False)
+        )
+
+        out = self.transform_obs(obs, done = done, update_state = True)
+        return out, reward, terminated, truncated, info