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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: env-ssl-wrapper
-Version: 0.5.7
+Version: 0.5.8
 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
@@ -182,7 +182,7 @@ uv run test_action_chunk.py --verify=True --sweep=True   # chunked rollouts matc
 
 ### 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))
+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/2610.12465))
 
 ```python
 import gymnasium as gym
@@ -202,9 +202,9 @@ obs, info = env.reset()
 
 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 sampling distribution is swappable too: the default follows the paper's Eq. (2) — the beta shape the score produces is floored at `eps`, logits are formed and softmaxed at `temperature` (floored at 1; the paper uses 2.) — and `sgs_sample_fn` draws i.i.d. from those `probs`. A custom `sample_fn` is handed, by name, whatever it declares — `p`, `counts`, `history`, `configs`, `weights` (the raw kernel), `probs` (floored and tempered), `rng`, `size`, `step`, `temperature`, 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`).
+The rest of the loop is hookable as well: target and temperature schedules, 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:
 
@@ -360,10 +360,10 @@ uv run pytest
 ```
 
 ```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},
+@article{zhang2026balanced,
+    title   = {A Balanced Data Diet: Addressing the Exploration Bottleneck in Mega-Scale RL for Robot Control},
+    author  = {Octi Zhang and Mateo Guaman Castro and Patrick Yin and Ignacio Dagnino and Abhishek Gupta and Rosario Scalise and Byron Boots},
+    journal = {arXiv preprint arXiv:2610.12465},
     year    = {2026}
 }
 ```

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

@@ -122,7 +122,7 @@ uv run test_action_chunk.py --verify=True --sweep=True   # chunked rollouts matc
 
 ### 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))
+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/2610.12465))
 
 ```python
 import gymnasium as gym
@@ -142,9 +142,9 @@ obs, info = env.reset()
 
 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 sampling distribution is swappable too: the default follows the paper's Eq. (2) — the beta shape the score produces is floored at `eps`, logits are formed and softmaxed at `temperature` (floored at 1; the paper uses 2.) — and `sgs_sample_fn` draws i.i.d. from those `probs`. A custom `sample_fn` is handed, by name, whatever it declares — `p`, `counts`, `history`, `configs`, `weights` (the raw kernel), `probs` (floored and tempered), `rng`, `size`, `step`, `temperature`, 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`).
+The rest of the loop is hookable as well: target and temperature schedules, 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:
 
@@ -300,10 +300,10 @@ uv run pytest
 ```
 
 ```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},
+@article{zhang2026balanced,
+    title   = {A Balanced Data Diet: Addressing the Exploration Bottleneck in Mega-Scale RL for Robot Control},
+    author  = {Octi Zhang and Mateo Guaman Castro and Patrick Yin and Ignacio Dagnino and Abhishek Gupta and Rosario Scalise and Byron Boots},
+    journal = {arXiv preprint arXiv:2610.12465},
     year    = {2026}
 }
 ```

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

@@ -107,8 +107,8 @@ def try_reset_slot(env, idx: int, options: dict):
     return None
 
 # curriculum score functions
-# Success-Guided Sampling (SGS) / OmniReset - Octi Zhang et al.
-# https://arxiv.org/abs/2603.15999
+# Success-Guided Sampling (SGS) - Octi Zhang et al.
+# https://arxiv.org/abs/2610.12465
 
 def sgs_score_fn(
     p: torch.Tensor,
@@ -167,7 +167,7 @@ def sgs_sample_fn(
     rng: np.random.Generator,
     size: int = 1
 ) -> np.ndarray:
-    # the paper's default - categorical draw over the beta shaped probs (see `sgs_score_fn`)
+    # categorical draw over the handed probs - the eps floor and temperature are applied upstream
 
     probs_np = to_numpy(probs).astype(np.float64)
 
@@ -178,13 +178,14 @@ def sgs_sample_fn(
 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).
+    Proposed by Octi Zhang et al. (https://arxiv.org/abs/2610.12465).
 
-    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.
+    The default sampling distribution is the paper's Eq. (2): the beta kernel around `target`
+    evaluated at the success rates (`sgs_score_fn`), floored at `eps`, logits formed, and
+    softmaxed at `temperature` (floored at 1, the paper uses 2.), 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).
@@ -198,6 +199,7 @@ class SGSSampler:
         eps: float = 0.05,
         window: int = 20,
         unvisited_weight: float = 1.,
+        temperature: float | Callable[[int], float] = 1.,
         uniform_prob: float | Callable[[int], float] = 0.,
         default_prob: float | Callable[[int], float] = 0.,
         score_fn: Callable | None = None,
@@ -218,6 +220,7 @@ class SGSSampler:
         self.eps = float(eps)
         self.window = int(window)
         self.unvisited_weight = float(unvisited_weight)
+        self.temperature = temperature
         self.uniform_prob = uniform_prob
         self.default_prob = default_prob
 
@@ -272,7 +275,19 @@ class SGSSampler:
 
     @property
     def probs(self) -> torch.Tensor:
-        return self._mixed_probs(self.weights)
+        return self._mixed_probs(self._sampling_weights(self.weights))
+
+    def _effective_temperature(self) -> float:
+        # the paper floors the temperature at one: T_eff = max(T, 1)
+
+        return max(resolve_schedule(self.temperature, self.step_count), 1.)
+
+    def _sampling_weights(self, weights: torch.Tensor) -> torch.Tensor:
+        # paper's Eq. (2) - floor the kernel at eps, form logits, softmax at T_eff
+
+        eps = self.eps
+
+        return (weights.clamp(min = eps) + eps) ** (1. / self._effective_temperature())
 
     def _mixed_probs(self, weights: torch.Tensor) -> torch.Tensor:
         # normalize, then mix in the uniform floor - what the default sampler draws from
@@ -313,10 +328,11 @@ class SGSSampler:
             **self._bank_args(step),
             configs = self.configs,
             weights = weights,
-            probs = self._mixed_probs(weights),
+            probs = self._mixed_probs(self._sampling_weights(weights)),
             rng = self.rng,
             size = size,
             uniform_prob = resolve_schedule(self.uniform_prob, step),
+            temperature = self._effective_temperature(),
         )
 
         indices = to_numpy(sample_fn(**inject_args(sample_fn, available))).astype(int).reshape(-1)
@@ -388,7 +404,7 @@ class SGSSampler:
 
 class SuccessGuidedSamplingWrapper(EnvWrapper):
     """
-    Success-Guided Sampling (SGS) environment wrapper (Octi Zhang et al., https://arxiv.org/abs/2603.15999).
+    Success-Guided Sampling (SGS) environment wrapper (Octi Zhang et al., https://arxiv.org/abs/2610.12465).
     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`).
@@ -405,6 +421,7 @@ class SuccessGuidedSamplingWrapper(EnvWrapper):
         eps: float = 0.05,
         window: int = 20,
         unvisited_weight: float = 1.,
+        temperature: float | Callable[[int], float] = 1.,
         uniform_prob: float | Callable[[int], float] = 0.,
         default_prob: float | Callable[[int], float] = 0.,
         score_fn: Callable | None = None,
@@ -432,6 +449,7 @@ class SuccessGuidedSamplingWrapper(EnvWrapper):
                 eps = eps,
                 window = window,
                 unvisited_weight = unvisited_weight,
+                temperature = temperature,
                 uniform_prob = uniform_prob,
                 default_prob = default_prob,
                 score_fn = score_fn,

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

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