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

{env_ssl_wrapper-0.5.6 → env_ssl_wrapper-0.5.7}/.gitignore

@@ -207,3 +207,6 @@ cython_debug/
 marimo/_static/
 marimo/_lsp/
 __marimo__/
+
+# Experiment artifacts (charts, checkpoints, progression logs)
+experiments/

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: env-ssl-wrapper
-Version: 0.5.6
+Version: 0.5.7
 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
@@ -198,7 +198,32 @@ 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`).
+`reset_configs` holds reset states — or callables drawing one (`lambda rng: ...`) — passed through `options`, or built by a `reset_options_fn` / `state_transform`.
+
+The score equation is swappable: the default beta shaped kernel around `target` (`sgs_score_fn`) is simply one rule of the bank's state. Whatever you pass as `score_fn` is handed, by name, the arguments it declares — `p` (per state success rates), `counts`, `history` (the raw episode outcomes), `step`, `window`, and the hparams `target`, `kappa`, `eps`, `unvisited_weight`.
+
+The sampling distribution is swappable too: the default (`sgs_sample_fn`) draws from the beta shape the score produces, while a custom `sample_fn` is handed, by name, whatever it declares — `p`, `counts`, `history`, `configs`, `weights`, `probs`, `rng`, `size`, `step`, and the hparams — so any distribution of the bank can be plugged in.
+
+The rest of the loop is hookable as well: target schedule, uniform floor (`uniform_prob`), env-native fallback (`default_prob`), success signal (`success_fn`).
+
+Track the curriculum as it runs — one snapshot of the bank per credited episode, kept in memory and appended to a csv when `log_path` is set:
+
+```python
+env = SuccessGuidedSamplingWrapper(
+    StandardizeEnvWrapper(gym.make('CartPole-v1')),
+    reset_configs = reset_states,
+    log_every = 1,                  # snapshot every episode (None turns logging off)
+    log_path = 'sgs_progression.csv',
+)
+
+env.trace                           # [dict(step, success_rates, probs), ...]
+```
+
+LunarLander is the worked example — a difficulty ladder of reset states, three arms (native spawns, uniform bank resets, SGS), paired evaluation and the speedup chart:
+
+```bash
+uv run train_lunar_sgs.py --steps=200000 --seeds=0,1,2,3,4   # writes experiments/lunar_sgs/
+```
 
 ## Wrappers
 

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

@@ -138,7 +138,32 @@ 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`).
+`reset_configs` holds reset states — or callables drawing one (`lambda rng: ...`) — passed through `options`, or built by a `reset_options_fn` / `state_transform`.
+
+The score equation is swappable: the default beta shaped kernel around `target` (`sgs_score_fn`) is simply one rule of the bank's state. Whatever you pass as `score_fn` is handed, by name, the arguments it declares — `p` (per state success rates), `counts`, `history` (the raw episode outcomes), `step`, `window`, and the hparams `target`, `kappa`, `eps`, `unvisited_weight`.
+
+The sampling distribution is swappable too: the default (`sgs_sample_fn`) draws from the beta shape the score produces, while a custom `sample_fn` is handed, by name, whatever it declares — `p`, `counts`, `history`, `configs`, `weights`, `probs`, `rng`, `size`, `step`, and the hparams — so any distribution of the bank can be plugged in.
+
+The rest of the loop is hookable as well: target schedule, uniform floor (`uniform_prob`), env-native fallback (`default_prob`), success signal (`success_fn`).
+
+Track the curriculum as it runs — one snapshot of the bank per credited episode, kept in memory and appended to a csv when `log_path` is set:
+
+```python
+env = SuccessGuidedSamplingWrapper(
+    StandardizeEnvWrapper(gym.make('CartPole-v1')),
+    reset_configs = reset_states,
+    log_every = 1,                  # snapshot every episode (None turns logging off)
+    log_path = 'sgs_progression.csv',
+)
+
+env.trace                           # [dict(step, success_rates, probs), ...]
+```
+
+LunarLander is the worked example — a difficulty ladder of reset states, three arms (native spawns, uniform bank resets, SGS), paired evaluation and the speedup chart:
+
+```bash
+uv run train_lunar_sgs.py --steps=200000 --seeds=0,1,2,3,4   # writes experiments/lunar_sgs/
+```
 
 ## Wrappers
 

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

@@ -12,6 +12,7 @@ from .sgs import (
     SGSWrapper,
     SGSSampler,
     sgs_score_fn,
+    sgs_sample_fn,
     variance_score_fn,
     uniform_score_fn,
 )
@@ -50,6 +51,7 @@ __all__ = [
     'SGSWrapper',
     'SGSSampler',
     'sgs_score_fn',
+    'sgs_sample_fn',
     'variance_score_fn',
     'uniform_score_fn',
     'evaluate_actor',

{env_ssl_wrapper-0.5.6 → env_ssl_wrapper-0.5.7}/env_ssl_wrapper/sgs.py

@@ -2,6 +2,8 @@ from __future__ import annotations
 
 from collections import deque
 import inspect
+import os
+from pathlib import Path
 from typing import Callable, Sequence
 
 import numpy as np
@@ -14,6 +16,7 @@ from .standardize.helpers import (
     any_true,
     copy_leaf,
     default,
+    divisible_by,
     dones_of,
     env_autoresets,
     env_num_envs,
@@ -48,6 +51,19 @@ def call_hook(fn: Callable, *args):
 def resolve_schedule(val: float | Callable[[int], float], step: int) -> float:
     return float(val(step)) if callable(val) else float(val)
 
+def inject_args(fn: Callable, available: dict) -> dict:
+    # dependency injection - hand the callable exactly the arguments it declares
+
+    try:
+        params = inspect.signature(fn).parameters
+    except (TypeError, ValueError):
+        return dict(available)
+
+    if any(param.kind == inspect.Parameter.VAR_KEYWORD for param in params.values()):
+        return dict(available)
+
+    return {name: value for name, value in available.items() if name in params}
+
 def default_success_fn(obs, reward, terminated, truncated, info):
     if isinstance(info, dict):
         succ = first_existing(info, 'is_success', 'success')
@@ -144,12 +160,34 @@ def uniform_score_fn(
 ) -> torch.Tensor:
     return torch.ones_like(p, dtype = torch.float32)
 
+# default sampling function
+
+def sgs_sample_fn(
+    probs: torch.Tensor,
+    rng: np.random.Generator,
+    size: int = 1
+) -> np.ndarray:
+    # the paper's default - categorical draw over the beta shaped probs (see `sgs_score_fn`)
+
+    probs_np = to_numpy(probs).astype(np.float64)
+
+    return rng.choice(len(probs_np), size = size, p = probs_np / probs_np.sum())
+
 # sampler
 
 class SGSSampler:
     """
     Success-Guided Sampling (SGS) curriculum sampler over a bank of reset configurations.
     Proposed by Octi Zhang et al. in OmniReset (https://arxiv.org/abs/2603.15999).
+
+    The default sampling distribution is the beta kernel around `target` evaluated at the
+    success rates (`sgs_score_fn`), drawn by `sgs_sample_fn`. Swap either hook - `score_fn`
+    reshapes the scores, `sample_fn` picks the distribution - each is handed, by name,
+    whatever it declares of `p`, `counts`, `history`, `configs`, `weights`, `probs`, `rng`,
+    `size`, `step`, `window` and the hparams.
+
+    Set `log_every` to keep snapshots of the bank in `self.trace` as episodes are credited,
+    and `log_path` to also append them to a csv (step, success rates, sampling probs).
     """
 
     def __init__(
@@ -163,13 +201,17 @@ class SGSSampler:
         uniform_prob: float | Callable[[int], float] = 0.,
         default_prob: float | Callable[[int], float] = 0.,
         score_fn: Callable | None = None,
+        sample_fn: Callable | None = None,
         state_transform: Callable | None = None,
+        log_every: int | None = None,
+        log_path: str | os.PathLike | None = None,
         seed: int | None = None,
     ):
         self.configs = [c for c in configs]
         assert len(self.configs) > 0, 'configs must not be empty'
         assert window >= 1, f'window must be at least 1, got {window}'
         assert eps > 0., f'eps must be positive, got {eps}'
+        assert not exists(log_every) or log_every >= 1, f'log_every must be at least 1, got {log_every}'
 
         self.target = target
         self.kappa = kappa
@@ -180,8 +222,13 @@ class SGSSampler:
         self.default_prob = default_prob
 
         self.score_fn = score_fn
+        self.sample_fn = sample_fn
         self.state_transform = state_transform
 
+        self.log_every = log_every
+        self.log_path = log_path
+        self.trace: list[dict] = []  # progression of the bank, one snapshot per log_every episodes
+
         self.history: list[deque[float]] = [deque(maxlen = self.window) for _ in range(len(self.configs))]
         self.counts = torch.zeros(len(self.configs), dtype = torch.long)
         self.step_count = 0
@@ -196,27 +243,42 @@ class SGSSampler:
         rates = [float(np.mean(h)) if len(h) else float('nan') for h in self.history]
         return torch.tensor(rates, dtype = torch.float32)
 
+    def _bank_args(self, step: int) -> dict:
+        # the common bank state both hooks are handed - the score fn sees this, the sampler
+        # all of it
+
+        return dict(
+            p = self.success_rates,
+            counts = self.counts,
+            history = self.history,
+            step = step,
+            window = self.window,
+            target = resolve_schedule(self.target, step),
+            kappa = resolve_schedule(self.kappa, step),
+            eps = self.eps,
+            unvisited_weight = self.unvisited_weight,
+        )
+
     @property
     def weights(self) -> torch.Tensor:
-        p = self.success_rates
+        # score_fn is swappable - handed, by name, whatever it declares (see `inject_args`),
+        # sgs_score_fn by default
 
-        if exists(self.score_fn):
-            return torch.as_tensor(call_hook(self.score_fn, p, self.counts, self.step_count), dtype = torch.float32).clamp(min = 0.)
+        score_fn = self.score_fn if exists(self.score_fn) else sgs_score_fn
 
-        return sgs_score_fn(
-            p,
-            self.counts,
-            target = resolve_schedule(self.target, self.step_count),
-            kappa = resolve_schedule(self.kappa, self.step_count),
-            eps = self.eps,
-            unvisited_weight = self.unvisited_weight
-        )
+        weights = score_fn(**inject_args(score_fn, self._bank_args(self.step_count)))
+
+        return torch.as_tensor(weights, dtype = torch.float32).clamp(min = 0.)
 
     @property
     def probs(self) -> torch.Tensor:
-        w = self.weights
-        total = w.sum()
-        probs = w / total if total > 0. else torch.full_like(w, 1. / len(w))
+        return self._mixed_probs(self.weights)
+
+    def _mixed_probs(self, weights: torch.Tensor) -> torch.Tensor:
+        # normalize, then mix in the uniform floor - what the default sampler draws from
+
+        total = weights.sum()
+        probs = weights / total if total > 0. else torch.full_like(weights, 1. / len(weights))
 
         uniform_prob = resolve_schedule(self.uniform_prob, self.step_count)
         assert 0. <= uniform_prob <= 1., f'uniform_prob must be within [0, 1], got {uniform_prob}'
@@ -239,15 +301,30 @@ class SGSSampler:
 
     def sample(self, batch_size: int | None = None):
         size = 1 if not exists(batch_size) else int(batch_size)
+        step = self.step_count
+
+        # sample_fn is swappable, handed whatever it declares (see `inject_args`) - the
+        # default draws from the beta shaped `probs`, a custom one from any distribution
+
+        sample_fn = self.sample_fn if exists(self.sample_fn) else sgs_sample_fn
+
+        weights = self.weights
+        available = dict(
+            **self._bank_args(step),
+            configs = self.configs,
+            weights = weights,
+            probs = self._mixed_probs(weights),
+            rng = self.rng,
+            size = size,
+            uniform_prob = resolve_schedule(self.uniform_prob, step),
+        )
 
-        probs_np = self.probs.detach().cpu().numpy()
-        probs_np = probs_np / probs_np.sum()
+        indices = to_numpy(sample_fn(**inject_args(sample_fn, available))).astype(int).reshape(-1)
+        assert len(indices) == size, f'sample_fn must return {size} indices, got {len(indices)}'
 
-        default_prob = resolve_schedule(self.default_prob, self.step_count)
+        default_prob = resolve_schedule(self.default_prob, step)
         assert 0. <= default_prob <= 1., f'default_prob must be within [0, 1], got {default_prob}'
 
-        indices = self.rng.choice(len(self.configs), size = size, p = probs_np)
-
         # a fraction of the draws can fall back to the environment's own reset distribution
 
         if default_prob > 0.:
@@ -275,6 +352,31 @@ class SGSSampler:
                 self.counts[idx] += 1
                 self.step_count += 1
 
+        self.log()
+
+    def log(self):
+        # snapshot in memory, and append to the csv when log_path is set
+
+        if not exists(self.log_every) or not divisible_by(self.step_count, self.log_every):
+            return
+
+        snapshot = dict(step = self.step_count, success_rates = self.success_rates, probs = self.probs)
+        self.trace.append(snapshot)
+
+        if not exists(self.log_path):
+            return
+
+        path = Path(self.log_path)
+        path.parent.mkdir(parents = True, exist_ok = True)
+
+        with path.open('a') as log:
+            if log.tell() == 0:  # empty file - write the header
+                columns = [f'success_rate_{i}' for i in range(len(self))] + [f'prob_{i}' for i in range(len(self))]
+                log.write(','.join(['step', *columns]) + '\n')
+
+            values = [snapshot['step'], *to_numpy(snapshot['success_rates']), *to_numpy(snapshot['probs'])]
+            log.write(','.join(f'{value:.4f}' for value in values) + '\n')
+
     def add_config(self, config, init_p: float | None = None) -> int:
         idx = len(self.configs)
         self.configs.append(copy_leaf(config))
@@ -289,6 +391,7 @@ class SuccessGuidedSamplingWrapper(EnvWrapper):
     Success-Guided Sampling (SGS) environment wrapper (Octi Zhang et al., https://arxiv.org/abs/2603.15999).
     Wraps on top of StandardizeEnvWrapper (or any single / vectorized environment) to adaptively sample
     initial reset states at the policy's frontier of competence and track per-config success rates.
+    Pass `log_every` / `log_path` to record the progression of the bank (see `SGSSampler`).
     """
 
     priority = 110
@@ -305,12 +408,15 @@ class SuccessGuidedSamplingWrapper(EnvWrapper):
         uniform_prob: float | Callable[[int], float] = 0.,
         default_prob: float | Callable[[int], float] = 0.,
         score_fn: Callable | None = None,
+        sample_fn: Callable | None = None,
         state_transform: Callable | None = None,
         success_fn: Callable | None = None,
         harvest_fn: Callable | None = None,
         reset_options_fn: Callable | None = None,
         set_state_fn: Callable | None = None,
         options_key: str = 'reset_to_state',
+        log_every: int | None = None,
+        log_path: str | os.PathLike | None = None,
         enabled: bool = True,
         seed: int | None = None,
     ):
@@ -329,7 +435,10 @@ class SuccessGuidedSamplingWrapper(EnvWrapper):
                 uniform_prob = uniform_prob,
                 default_prob = default_prob,
                 score_fn = score_fn,
+                sample_fn = sample_fn,
                 state_transform = state_transform,
+                log_every = log_every,
+                log_path = log_path,
                 seed = seed,
             )
 
@@ -364,6 +473,10 @@ class SuccessGuidedSamplingWrapper(EnvWrapper):
     def counts(self) -> torch.Tensor:
         return self.sampler.counts
 
+    @property
+    def trace(self) -> list[dict]:
+        return self.sampler.trace
+
     def _build_options(self, state, cfg_id) -> dict:
         if exists(self.reset_options_fn):
             return call_hook(self.reset_options_fn, state, cfg_id)
@@ -407,13 +520,6 @@ class SuccessGuidedSamplingWrapper(EnvWrapper):
 
         return obs, info
 
-        self._done_since_reset = np.zeros(self.num_envs, dtype = bool)
-
-        info = dict(info) if isinstance(info, dict) else {}
-        info['sgs_config_id'] = self._format_config_id_info()
-
-        return obs, info
-
     def step(self, action):
         obs, reward, terminated, truncated, info = normalize_step_out(self.env.step(action))
         info = dict(info) if isinstance(info, dict) else {}

{env_ssl_wrapper-0.5.6 → env_ssl_wrapper-0.5.7}/env_ssl_wrapper/standardize/helpers.py

@@ -14,6 +14,9 @@ def exists(v):
 def default(v, d):
     return v if exists(v) else (d() if callable(d) else d)
 
+def divisible_by(num, den):
+    return (num % den) == 0
+
 def cast_tuple(val, length = 1):
     if isinstance(val, tuple):
         return val

{env_ssl_wrapper-0.5.6 → env_ssl_wrapper-0.5.7}/pyproject.toml

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