env-ssl-wrapper 0.5.1__tar.gz → 0.5.2__tar.gz

{env_ssl_wrapper-0.5.1 → env_ssl_wrapper-0.5.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: env-ssl-wrapper
-Version: 0.5.1
+Version: 0.5.2
 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
@@ -34,16 +34,11 @@ Classifier: License :: OSI Approved :: MIT License
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Requires-Python: >=3.10
-Requires-Dist: discrete-continuous-embed-readout
 Requires-Dist: einops>=0.8.1
-Requires-Dist: einx>=0.3.0
 Requires-Dist: imageio
 Requires-Dist: imageio-ffmpeg
-Requires-Dist: memmap-replay-buffer
-Requires-Dist: torch-einops-utils>=0.0.29
+Requires-Dist: numpy>=1.24
 Requires-Dist: torch>=2.5
-Requires-Dist: x-mlps-pytorch
-Requires-Dist: x-transformers
 Provides-Extra: examples
 Requires-Dist: fire; extra == 'examples'
 Provides-Extra: test
@@ -277,6 +272,29 @@ with MultiprocessingVecEnv('CartPole-v1', num_envs = 8) as env:
     obs, info = env.reset()
 ```
 
+## Evaluation
+
+`evaluate_actor` rolls a policy through any standardized env and returns per-episode stats, optionally recording videos.
+
+```python
+import gymnasium as gym
+from env_ssl_wrapper import StandardizeEnvWrapper, evaluate_actor
+
+env = StandardizeEnvWrapper(gym.make('CartPole-v1'))
+
+stats = evaluate_actor(policy, env, episodes = 10, seed = 0)
+
+stats.mean          # mean return
+stats.std           # std of returns
+stats.min, stats.max
+stats.success_rate(475.)          # fraction of episodes above a threshold
+stats.summary()                   # dict(episodes, mean, std, min, max)
+
+# deterministic = False samples from distributions instead of collapsing to the mode
+# an actor exposing `.dist` takes precedence over calling it directly
+# video_path = 'eval.mp4' records the first video_episodes episodes (single env only)
+```
+
 ## Tests
 
 ```bash

{env_ssl_wrapper-0.5.1 → env_ssl_wrapper-0.5.2}/README.md

@@ -212,6 +212,29 @@ with MultiprocessingVecEnv('CartPole-v1', num_envs = 8) as env:
     obs, info = env.reset()
 ```
 
+## Evaluation
+
+`evaluate_actor` rolls a policy through any standardized env and returns per-episode stats, optionally recording videos.
+
+```python
+import gymnasium as gym
+from env_ssl_wrapper import StandardizeEnvWrapper, evaluate_actor
+
+env = StandardizeEnvWrapper(gym.make('CartPole-v1'))
+
+stats = evaluate_actor(policy, env, episodes = 10, seed = 0)
+
+stats.mean          # mean return
+stats.std           # std of returns
+stats.min, stats.max
+stats.success_rate(475.)          # fraction of episodes above a threshold
+stats.summary()                   # dict(episodes, mean, std, min, max)
+
+# deterministic = False samples from distributions instead of collapsing to the mode
+# an actor exposing `.dist` takes precedence over calling it directly
+# video_path = 'eval.mp4' records the first video_episodes episodes (single env only)
+```
+
 ## Tests
 
 ```bash

{env_ssl_wrapper-0.5.1 → env_ssl_wrapper-0.5.2}/env_ssl_wrapper/__init__.py

@@ -1,18 +1,18 @@
 from __future__ import annotations
 
 import sys
+
 from . import standardize
-from .standardize import *
+from .standardize import *  # noqa: F403
 
 from .memory_trace import MemoryTraceWrapper
 from .action_chunk import ActionChunkWrapper
 from .evaluate import evaluate_actor, EpisodeStats
 
-# Wire backwards-compatibility aliases in sys.modules and module globals
-# so imports like `from env_ssl_wrapper.done_tracker_wrapper import DoneTrackerWrapper`
-# or `import env_ssl_wrapper.mocks` continue to work without breaking.
+# backwards compatibility — the standardize wrappers used to live at the top
+# level, so keep the old module paths (e.g. `env_ssl_wrapper.mocks`) importable
 
-_STANDARDIZE_SUBMODULES = (
+_LEGACY_MODULES = (
     'adapters',
     'auto_batched_wrapper',
     'action_transform_wrapper',
@@ -31,11 +31,8 @@ _STANDARDIZE_SUBMODULES = (
     'vector',
 )
 
-for _name in _STANDARDIZE_SUBMODULES:
-    _mod = getattr(standardize, _name, None)
-    if _mod is not None:
-        sys.modules[f'{__name__}.{_name}'] = _mod
-        globals()[_name] = _mod
+for _name in _LEGACY_MODULES:
+    sys.modules[f'{__name__}.{_name}'] = getattr(standardize, _name)
 
 __all__ = [
     *standardize.__all__,
@@ -49,4 +46,3 @@ def __getattr__(name):
     if hasattr(standardize, name):
         return getattr(standardize, name)
     raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
-

{env_ssl_wrapper-0.5.1 → env_ssl_wrapper-0.5.2}/env_ssl_wrapper/evaluate.py

@@ -9,6 +9,7 @@ from torch import nn, is_tensor
 from torch.distributions import Categorical, Distribution
 from torch.utils._pytree import tree_flatten, tree_map
 
+from .standardize.spaces import action_space_is_discrete
 from .standardize.helpers import (
     EnvWrapper,
     env_autoresets,
@@ -93,7 +94,19 @@ def to_action(env, out, deterministic = True):
         if not deterministic:
             return out.sample()
 
-        return out.probs.argmax(dim = -1) if exists(get_attr(out, 'probs')) else out.mean
+        probs = get_attr(out, 'probs')
+
+        # categorical-style probs carry a trailing action axis to argmax over
+
+        if exists(probs) and probs.ndim >= 2 and probs.shape[-1] > 1:
+            return probs.argmax(dim = -1)
+
+        # bernoulli-style probabilities threshold at 0.5
+
+        if exists(probs):
+            return (probs > 0.5).long()
+
+        return out.mean
 
     if is_tensor(out) and out.is_floating_point() and out.ndim > 0:
         num_actions = get_attr(get_adapter(env).action_space, 'n')
@@ -136,10 +149,21 @@ def to_env_action(env, action):
 
         action = to_numpy(action)
 
-        if action.ndim > 1 and action.shape[0] == 1:
-            action = action[0]
+        # single raw sims get the leading batch dim stripped, while discrete
+        # actions additionally collapse to scalars
+
+        if not is_vectorized(env):
+            if action.ndim > 1 and action.shape[0] == 1:
+                action = action[0]
+
+            if action_space_is_discrete(get_adapter(env).action_space):
+                if action.ndim > 0 and action.shape[0] == 1:
+                    action = action[0]
+
+                if action.ndim == 0:
+                    action = action.item()
 
-        return action.item() if action.ndim == 0 else action
+        return action
 
     # wrappers take batched torch actions - pad up to the declared rank
 
@@ -169,6 +193,9 @@ def capture_frame(env):
         except Exception:
             return None
 
+    if not exists(frame):
+        return None
+
     if is_tensor(frame):
         frame = frame.detach().cpu().numpy()
 
@@ -200,6 +227,17 @@ def video_path_at(path, index):
     path = Path(path)
     return str(path if index == 0 else path.with_name(f'{path.stem}_ep{index}{path.suffix}'))
 
+def safe_reset(env, seed = None):
+    # gymnasium resets accept seed; legacy sims reset without it
+
+    if exists(seed):
+        try:
+            return normalize_reset_out(env.reset(seed = seed))
+        except TypeError:
+            pass
+
+    return normalize_reset_out(env.reset())
+
 # evaluate
 
 @torch.no_grad()
@@ -241,7 +279,7 @@ def evaluate_actor(
 
     assert not (record and num_envs > 1), 'video recording requires a single env'
 
-    obs, info = normalize_reset_out(env.reset(seed = seed) if exists(seed) else env.reset())
+    obs, info = safe_reset(env, seed)
     batched = is_vectorized(env) or all(map(is_tensor, tree_flatten(obs)[0]))
 
     goal = goal.expand(num_envs, -1) if exists(goal) and goal.shape[0] == 1 and num_envs > 1 else goal
@@ -249,6 +287,7 @@ def evaluate_actor(
     returns = torch.zeros(num_envs)
     lengths = torch.zeros(num_envs)
     steps = torch.zeros(num_envs, dtype = torch.long)
+    active = np.ones(num_envs, dtype = bool)
 
     done_returns, done_lengths, videos, frames = [], [], [], []
 
@@ -285,7 +324,12 @@ def evaluate_actor(
             if exists(max_steps):
                 done |= to_numpy(steps >= max_steps).reshape(-1)
 
-            for i in np.where(done)[0]:
+            # a slot is only recorded once per episode — non-autoreset slots
+            # latch their done flag until the whole vector env resets
+
+            newly_done = done & active
+
+            for i in np.where(newly_done)[0]:
                 if len(done_returns) >= episodes:
                     break
 
@@ -309,6 +353,7 @@ def evaluate_actor(
                 returns[i] = 0.
                 lengths[i] = 0.
                 steps[i] = 0
+                active[i] = False
 
             if done.any() and len(done_returns) < episodes:
                 if record:
@@ -318,7 +363,10 @@ def evaluate_actor(
 
                 if not env_autoresets(env) and (num_envs == 1 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 = normalize_reset_out(env.reset(seed = reset_seed) if exists(reset_seed) else env.reset())
+                    obs, info = safe_reset(env, reset_seed)
+                    active[:] = True
+                elif env_autoresets(env):
+                    active[newly_done] = True
     finally:
         if exists(module) and exists(was_training):
             module.train(was_training)

{env_ssl_wrapper-0.5.1 → env_ssl_wrapper-0.5.2}/env_ssl_wrapper/memory_trace.py

@@ -3,7 +3,6 @@ from typing import Sequence
 
 import torch
 from torch import is_tensor
-from torch_einops_utils import pad_right_ndim_to
 
 from .standardize.helpers import (
     TransformObservationWrapper,
@@ -27,7 +26,11 @@ def calc_lerp_weight(lam, done, x):
         done = torch.as_tensor(done, device = x.device)
 
     weight = torch.where(done, 1., 1. - lam).to(x)
-    return pad_right_ndim_to(weight, x.ndim)
+
+    if x.ndim <= weight.ndim:
+        return weight
+
+    return weight.reshape(*weight.shape, *((1,) * (x.ndim - weight.ndim)))
 
 # Partially Observable Reinforcement Learning with Memory Traces - Eberhard et al.
 # https://arxiv.org/abs/2503.15200

{env_ssl_wrapper-0.5.1 → env_ssl_wrapper-0.5.2}/env_ssl_wrapper/standardize/__init__.py

@@ -120,6 +120,7 @@ __all__ = [
     'MultiprocessingVecEnv',
     'StandardizeEnvWrapper',
     'StandardizeEnv',
+    'StandardizedEnv',
     'BaseEnvAdapter',
     'get_adapter',
     'register_adapter',

{env_ssl_wrapper-0.5.1 → env_ssl_wrapper-0.5.2}/env_ssl_wrapper/standardize/action_transform_wrapper.py

@@ -5,7 +5,6 @@ import torch
 from torch import is_tensor
 from torch.utils._pytree import tree_map
 from functools import partial
-from einops import rearrange
 
 from .helpers import EnvWrapper, copy_leaf, default, exists, first_existing, get_attr
 from .spaces import space_from_action_spec
@@ -95,7 +94,7 @@ class ActionTransformWrapper(EnvWrapper):
         was_scalar = t.ndim == 0
 
         if was_scalar:
-            t = rearrange(t, '-> 1')
+            t = t[None]
 
         # bounds may be unbatched (single_action_space) or already batched
         # (vector envs exposing only a batched action_space)
@@ -105,22 +104,28 @@ class ActionTransformWrapper(EnvWrapper):
             low = np.broadcast_to(low, (dim,))
             high = np.broadcast_to(high, (dim,))
 
+        # invalid (unbounded) dims pass through — neutral bounds keep inf out of the arithmetic
+
         valid = np.isfinite(low) & np.isfinite(high)
 
         if is_tensor(t):
-            low = torch.tensor(low, device = t.device, dtype = t.dtype)
-            high = torch.tensor(high, device = t.device, dtype = t.dtype)
-            valid = torch.tensor(valid, device = t.device)
+            valid_t = torch.as_tensor(valid, device = t.device)
+            low = torch.as_tensor(np.where(valid, low, 0.), device = t.device, dtype = t.dtype)
+            high = torch.as_tensor(np.where(valid, high, 1.), device = t.device, dtype = t.dtype)
+        else:
+            valid_t = valid
+            low = np.where(valid, low, 0.)
+            high = np.where(valid, high, 1.)
 
         rescaled = rescale(t, self.from_range, (low, high))
 
         if is_tensor(t):
-            rescaled = torch.where(valid, torch.clamp(rescaled, low, high), t)
+            rescaled = torch.where(valid_t, torch.clamp(rescaled, low, high), t)
         else:
-            rescaled = np.where(valid, np.clip(rescaled, low, high), t)
+            rescaled = np.where(valid_t, np.clip(rescaled, low, high), t)
 
         if was_scalar:
-            rescaled = rearrange(rescaled, '1 ->')
+            rescaled = rescaled[0]
 
         return rescaled
 

{env_ssl_wrapper-0.5.1 → env_ssl_wrapper-0.5.2}/env_ssl_wrapper/standardize/auto_batched_wrapper.py

@@ -5,7 +5,6 @@ from functools import partial
 import numpy as np
 import torch
 from torch.utils._pytree import tree_map
-from einops import rearrange
 
 from .helpers import (
     FINAL_OBSERVATION_KEYS,
@@ -40,22 +39,12 @@ def maybe_expand_dim(x):
 
     def _expand(t):
         arr = to_numeric_array(t)
+        return arr[None] if is_array(arr) else t
 
-        if not is_array(arr):
-            return t
+    if is_array(x):
+        return x[None]
 
-        if arr.ndim == 0:
-            return rearrange(arr, '-> 1')
-
-        return rearrange(arr, '... -> 1 ...')
-
-    if is_tensor(x):
-        return rearrange(x, '-> 1') if x.ndim == 0 else rearrange(x, '... -> 1 ...')
-
-    if isinstance(x, np.ndarray) and x.dtype.kind in 'biufc':
-        return rearrange(x, '-> 1') if x.ndim == 0 else rearrange(x, '... -> 1 ...')
-
-    return tree_map(_expand, x)
+    return tree_map(_expand, x) if isinstance(x, (dict, list, tuple)) else _expand(x)
 
 def is_integer_dtype(t):
     if is_tensor(t):
@@ -120,11 +109,11 @@ def heuristic_leaf(t, is_vector = False):
         return arr
 
     if not is_vector and arr.ndim > 1 and arr.shape[0] == 1:
-        arr = rearrange(arr, '1 ... -> ...')
+        arr = arr.reshape(*arr.shape[1:])
 
     if is_integer_dtype(arr):
         while arr.ndim > 1 and arr.shape[-1] == 1:
-            arr = rearrange(arr, '... 1 -> ...')
+            arr = arr.reshape(*arr.shape[:-1])
 
         if not is_vector and (arr.numel() if is_tensor(arr) else arr.size) == 1:
             return arr.item()
@@ -206,7 +195,8 @@ class AutoBatchedWrapper(EnvWrapper):
         if self.is_vector:
             return out
 
-        obs, reward, terminated, truncated, info = *maybe_expand_dim(out[:4]), out[4]
+        obs, reward, terminated, truncated, info = out
+        obs, reward, terminated, truncated = (maybe_expand_dim(t) for t in (obs, reward, terminated, truncated))
 
         if isinstance(info, dict):
             for key in FINAL_OBSERVATION_KEYS:

{env_ssl_wrapper-0.5.1 → env_ssl_wrapper-0.5.2}/env_ssl_wrapper/standardize/episode_padding_wrapper.py

@@ -6,7 +6,6 @@ import numpy as np
 import torch
 from torch import is_tensor
 from torch.utils._pytree import tree_map
-from einops import rearrange
 
 from .helpers import (
     EnvWrapper,
@@ -28,7 +27,7 @@ def broadcast_mask(mask, x):
     m = torch.as_tensor(mask, device = x.device, dtype = torch.bool)
     diff = x.ndim - m.ndim
 
-    return rearrange(m, f'... -> ... {" ".join(["1"] * diff)}') if diff > 0 else m
+    return m.reshape(*m.shape, *((1,) * diff)) if diff > 0 else m
 
 def zero_mask(x, mask, fill_scalar = None):
     if is_tensor(x):

{env_ssl_wrapper-0.5.1 → env_ssl_wrapper-0.5.2}/env_ssl_wrapper/standardize/flatten_obs_wrapper.py

@@ -4,7 +4,6 @@ import numpy as np
 import torch
 from torch import is_tensor
 from torch.utils._pytree import tree_flatten
-from einops import rearrange
 
 from .helpers import (
     TransformObservationWrapper,
@@ -30,17 +29,9 @@ def flatten_leaf(t, is_vector = False):
         t = np.asarray(t)
 
     if not is_vector:
-        if t.ndim == 0:
-            return rearrange(t, '-> 1')
-        return rearrange(t, '... -> (...)')
+        return t.reshape(1) if t.ndim == 0 else t.reshape(-1)
 
-    if t.ndim == 0:
-        return rearrange(t, '-> 1 1')
-
-    if t.ndim == 1:
-        return rearrange(t, 'b -> b 1')
-
-    return rearrange(t, 'b ... -> b (...)')
+    return t.reshape(1, 1) if t.ndim == 0 else t.reshape(t.shape[0], -1)
 
 def concat_leaves(leaves, is_vector = False):
     axis = -1 if is_vector else 0

{env_ssl_wrapper-0.5.1 → env_ssl_wrapper-0.5.2}/env_ssl_wrapper/standardize/helpers.py

@@ -84,9 +84,6 @@ def copy_leaf(x):
 
     return x
 
-def copy_tree(tree):
-    return tree_map(copy_leaf, tree)
-
 def dones_of(terminated, truncated):
     if not isinstance(terminated, (dict, list, tuple)):
         return terminated | truncated
@@ -138,6 +135,8 @@ def _zero_leaf(x):
         return False
     if isinstance(x, (int, float, np.number)):
         return np.zeros_like(x)
+    if exists(get_attr(x, '__array__')):
+        return np.zeros_like(np.asarray(x))
     return 0
 
 def _stack_leaves(leaves):

{env_ssl_wrapper-0.5.1 → env_ssl_wrapper-0.5.2}/env_ssl_wrapper/standardize/image_wrapper.py

@@ -39,14 +39,10 @@ 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)
+    img = torch.from_numpy(np.ascontiguousarray(img))
 
     if img.ndim == 4:
         return rearrange(img, 'b h w c -> b c h w'), True

{env_ssl_wrapper-0.5.1 → env_ssl_wrapper-0.5.2}/env_ssl_wrapper/standardize/tensor_wrapper.py

@@ -32,6 +32,10 @@ def to_torch_leaf(t, device, cast_obs_to_float = True):
             return t
 
     dtype = torch.float32 if cast_obs_to_float and t.dtype != torch.bool else t.dtype
+
+    if t.dtype == dtype and t.device == device:
+        return t
+
     return t.to(device = device, dtype = dtype)
 
 def numpy_to_torch(x, device, cast_obs_to_float = True):
@@ -72,7 +76,11 @@ def torch_to_numpy(x):
 def contract(t, to_float = False):
     if not is_tensor(t):
         return t
-    return t.float() if to_float else t.bool()
+
+    if to_float:
+        return t if t.dtype == torch.float32 else t.float()
+
+    return t if t.dtype == torch.bool else t.bool()
 
 # class
 
@@ -105,7 +113,7 @@ class TensorWrapper(EnvWrapper):
         self.cast = partial(numpy_to_torch, device = self.device, cast_obs_to_float = self.cast_obs_to_float)
 
     def cast_info(self, info):
-        if not isinstance(info, dict):
+        if not isinstance(info, dict) or not info:
             return
 
         for key in FINAL_OBSERVATION_KEYS:
@@ -126,10 +134,14 @@ class TensorWrapper(EnvWrapper):
         return obs, info
 
     def to_contract(self, t, to_float = False):
-        if not isinstance(t, (dict, list, tuple)):
-            leaf = to_torch_leaf(t, self.device, cast_obs_to_float = False)
+        def convert(leaf):
+            leaf = to_torch_leaf(leaf, self.device, cast_obs_to_float = False)
             return contract(leaf, to_float = to_float)
-        return contract(self.cast(t), to_float = to_float)
+
+        if not isinstance(t, (dict, list, tuple)):
+            return convert(t)
+
+        return tree_map(convert, t)
 
     def step(self, action):
         if self.convert_in and not self.torch_native:

{env_ssl_wrapper-0.5.1 → env_ssl_wrapper-0.5.2}/env_ssl_wrapper/standardize/vector.py

@@ -2,7 +2,6 @@ from __future__ import annotations
 
 from multiprocessing import get_context
 import numpy as np
-import torch
 from torch import is_tensor
 from torch.utils._pytree import tree_map
 
@@ -119,6 +118,9 @@ class _InlineConn:
     def recv(self):
         return self._result
 
+    def close(self):
+        pass
+
 def _safe_send(conn, msg):
     """Send that swallows broken-pipe — for broadcast patterns where _recv_all reports the error."""
     try:

{env_ssl_wrapper-0.5.1 → env_ssl_wrapper-0.5.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "env-ssl-wrapper"
-version = "0.5.1"
+version = "0.5.2"
 description = "One torch-native interface for any MDP environment"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }
@@ -24,16 +24,11 @@ classifiers=[
 ]
 
 dependencies = [
-    "discrete-continuous-embed-readout",
-    "einx>=0.3.0",
     "einops>=0.8.1",
     "imageio",
     "imageio-ffmpeg",
-    "memmap-replay-buffer",
+    "numpy>=1.24",
     "torch>=2.5",
-    "torch-einops-utils>=0.0.29",
-    "x-transformers",
-    "x-mlps-pytorch",
 ]
 
 [project.urls]