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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: env-ssl-wrapper
-Version: 0.5.0
+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.0 → 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.0 → 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.0 → env_ssl_wrapper-0.5.2}/env_ssl_wrapper/action_chunk.py

@@ -12,6 +12,7 @@ from .standardize.helpers import (
     any_true,
     dones_of,
     first_existing,
+    get_adapter,
     get_attr,
     is_vectorized,
     to_numpy,
@@ -70,7 +71,7 @@ class ActionChunkWrapper(EnvWrapper):
         self.action_shape = tuple(get_attr(action_space, 'shape', ()) or ())
 
         self.expects_batch = is_vectorized(env)
-        self.raw_env = not isinstance(env, EnvWrapper)
+        self.convert_in = not isinstance(env, EnvWrapper) and not get_adapter(env).torch_native
 
     @property
     def chunk_action_shape(self):
@@ -86,9 +87,9 @@ class ActionChunkWrapper(EnvWrapper):
         if not (is_tensor(actions) or isinstance(actions, np.ndarray)):
             actions = np.asarray(actions)
 
-        # raw simulators only understand numpy - wrapper chains from this package handle torch actions
+        # raw simulators only understand numpy — wrapper chains and torch-native sims take torch
 
-        if self.raw_env and is_tensor(actions):
+        if is_tensor(actions) and self.convert_in:
             actions = to_numpy(actions)
 
         chunk_axis = actions.ndim - 1 - len(self.action_shape)

{env_ssl_wrapper-0.5.0 → 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')
@@ -128,15 +141,29 @@ def to_env_action(env, action):
     device = get_attr(env, 'device')
     action = action.to(device) if exists(device) else action
 
-    # raw simulators take unbatched numpy
+    # raw simulators take unbatched numpy — torch-native sims take torch as-is (zero-copy)
 
     if not isinstance(env, EnvWrapper):
+        if get_adapter(env).torch_native:
+            return 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
 
@@ -166,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()
 
@@ -197,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()
@@ -238,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
@@ -246,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 = [], [], [], []
 
@@ -282,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
 
@@ -306,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:
@@ -315,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.0 → 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.0 → env_ssl_wrapper-0.5.2}/env_ssl_wrapper/standardize/__init__.py

@@ -34,6 +34,7 @@ from .adapters import (
     BaseEnvAdapter,
     get_adapter,
     register_adapter,
+    MjlabAdapter,
     MujocoWarpAdapter,
     IsaacAdapter,
     PyBulletAdapter,
@@ -68,6 +69,7 @@ from .mocks import (
     AutoresetVectorMockEnv,
     PufferVectorMockEnv,
     PufferTensorMockEnv,
+    MjlabMockEnv,
     ManiSkillMockEnv,
     BraxMockEnv,
     MjxMockEnv,
@@ -118,9 +120,11 @@ __all__ = [
     'MultiprocessingVecEnv',
     'StandardizeEnvWrapper',
     'StandardizeEnv',
+    'StandardizedEnv',
     'BaseEnvAdapter',
     'get_adapter',
     'register_adapter',
+    'MjlabAdapter',
     'MujocoWarpAdapter',
     'IsaacAdapter',
     'PyBulletAdapter',
@@ -149,6 +153,7 @@ __all__ = [
     'AutoresetVectorMockEnv',
     'PufferVectorMockEnv',
     'PufferTensorMockEnv',
+    'MjlabMockEnv',
     'ManiSkillMockEnv',
     'BraxMockEnv',
     'MjxMockEnv',

{env_ssl_wrapper-0.5.0 → 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
 
@@ -129,8 +134,6 @@ class ActionTransformWrapper(EnvWrapper):
             if not is_float_dtype(t):
                 return t
 
-            t = copy_leaf(t)
-
             for ind, transform in enumerate(self.transforms):
                 indices = transform.get('indices', ind if len(self.transforms) > 1 else None)
 
@@ -141,6 +144,9 @@ class ActionTransformWrapper(EnvWrapper):
                 fn = partial(rescale, from_range = from_range, to_range = to_range)
 
                 if exists(indices):
+                    # copy only when writing in place, so the caller's action is never mutated
+
+                    t = copy_leaf(t)
                     t[..., indices] = fn(t[..., indices])
                 else:
                     t = fn(t)

{env_ssl_wrapper-0.5.0 → env_ssl_wrapper-0.5.2}/env_ssl_wrapper/standardize/adapters.py

@@ -6,18 +6,20 @@ from .helpers import (
     exists,
     first_existing,
     get_attr,
-    is_time_step,
     normalize_reset_out,
     normalize_step_out,
     safe_close,
     truthy_attr,
-    zero_like,
 )
 from .spaces import space_from_action_spec
 
 # base adapter
 
 class BaseEnvAdapter:
+    # torch-native sims consume torch actions directly on device — no numpy roundtrip
+
+    torch_native = False
+
     @classmethod
     def matches(cls, env) -> bool:
         return False
@@ -148,6 +150,10 @@ class WrapperAdapter(BaseEnvAdapter):
         val = first_existing(self.env, 'autoreset', 'autoresets', 'autoreset_mode')
         return truthy_attr(val) if exists(val) else self.inner_adapter.autoresets
 
+    @property
+    def torch_native(self) -> bool:
+        return self.inner_adapter.torch_native
+
     @property
     def action_space(self):
         return default(
@@ -194,13 +200,8 @@ class DMControlAdapter(BaseEnvAdapter):
     def action_space(self):
         return default(super().action_space, space_from_action_spec(self.env))
 
-    @property
-    def is_vectorized(self) -> bool:
-        return False
-
-    @property
-    def autoresets(self) -> bool:
-        return False
+    is_vectorized = False
+    autoresets = False
 
 # pybullet adapter
 
@@ -227,17 +228,32 @@ class PyBulletAdapter(BaseEnvAdapter):
             return
         super().seed(seed)
 
-    @property
-    def is_vectorized(self) -> bool:
-        return False
+    is_vectorized = False
+    autoresets = False
+
+# mjlab adapter (manager-based rl envs on mujoco warp)
+
+class MjlabAdapter(BaseEnvAdapter):
+    torch_native = True
+    is_vectorized = True
+
+    @classmethod
+    def matches(cls, env):
+        mod = getattr(type(env), '__module__', '')
+        name = type(env).__name__
+        return 'mjlab' in mod.lower() or name == 'ManagerBasedRlEnv'
 
     @property
     def autoresets(self) -> bool:
-        return False
+        return truthy_attr(get_attr(get_attr(self.env, 'cfg'), 'auto_reset', True))
 
 # isaac sim adapter (isaac gym / isaac lab / omniverse)
 
 class IsaacAdapter(BaseEnvAdapter):
+    torch_native = True
+    is_vectorized = True
+    autoresets = True
+
     @classmethod
     def matches(cls, env):
         mod = getattr(type(env), '__module__', '')
@@ -245,14 +261,6 @@ class IsaacAdapter(BaseEnvAdapter):
         isaac_kw = ('isaac', 'omni.isaac', 'isaacgym', 'isaaclab')
         return any(k in mod.lower() for k in isaac_kw) or 'Isaac' in name or exists(get_attr(env, 'sim_device')) or exists(get_attr(env, 'physics_sim_view'))
 
-    @property
-    def is_vectorized(self) -> bool:
-        return True
-
-    @property
-    def autoresets(self) -> bool:
-        return True
-
 # mujoco warp / warp / brax / mjx adapter
 
 class MujocoWarpAdapter(BaseEnvAdapter):
@@ -276,16 +284,14 @@ class MujocoWarpAdapter(BaseEnvAdapter):
 # pufferlib adapter
 
 class PufferLibAdapter(BaseEnvAdapter):
+    is_vectorized = True
+
     @classmethod
     def matches(cls, env):
         mod = getattr(type(env), '__module__', '')
         name = type(env).__name__
         return 'Puffer' in name or 'pufferlib' in mod or exists(get_attr(env, 'puffer_env'))
 
-    @property
-    def is_vectorized(self) -> bool:
-        return True
-
     @property
     def autoresets(self) -> bool:
         return truthy_attr(first_existing(self.env, 'autoreset', 'autoresets', 'autoreset_mode'))
@@ -350,6 +356,18 @@ class GymnasiumAdapter(BaseEnvAdapter):
             pass
         return self.num_envs > 1
 
+    def unwrapped_adapter(self):
+        # foreign gym wrappers (TimeLimit, RecordEpisodeStatistics, ...) inherit
+        # the unwrapped sim's adapter traits
+
+        inner = get_attr(self.env, 'unwrapped')
+        return get_adapter(inner) if exists(inner) and inner is not self.env else None
+
+    @property
+    def torch_native(self) -> bool:
+        inner = self.unwrapped_adapter()
+        return exists(inner) and inner.torch_native
+
     @property
     def autoresets(self) -> bool:
         mode = first_existing(self.env, 'autoreset', 'autoresets', 'autoreset_mode')
@@ -359,9 +377,13 @@ class GymnasiumAdapter(BaseEnvAdapter):
 
         try:
             from gymnasium.vector import VectorEnv
-            return isinstance(self.env, VectorEnv)
+            if isinstance(self.env, VectorEnv):
+                return True
         except ImportError:
-            return False
+            pass
+
+        inner = self.unwrapped_adapter()
+        return exists(inner) and inner.autoresets
 
     def seed(self, seed: int):
         try:
@@ -400,6 +422,7 @@ class DefaultAdapter(BaseEnvAdapter):
 
 ADAPTER_REGISTRY: list[type[BaseEnvAdapter]] = [
     WrapperAdapter,
+    MjlabAdapter,
     MujocoWarpAdapter,
     IsaacAdapter,
     PyBulletAdapter,

{env_ssl_wrapper-0.5.0 → 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,
@@ -19,7 +18,6 @@ from .helpers import (
     is_tensor,
     is_vectorized,
 )
-from .spaces import space_from_action_spec
 
 # helper functions
 
@@ -41,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):
@@ -121,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()
@@ -207,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.0 → 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.0 → 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.0 → 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.0 → 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.0 → env_ssl_wrapper-0.5.2}/env_ssl_wrapper/standardize/mocks.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from types import SimpleNamespace
 from typing import NamedTuple
 
 import numpy as np
@@ -391,6 +392,52 @@ class IsaacLabMockEnv(MockEnv):
 
         return self.obs(), torch.randn(self.num_envs), dones, torch.zeros(self.num_envs, dtype = torch.bool), {}
 
+# mjlab — manager-based RL on mujoco warp: always vectorized, torch-native
+# (step takes torch tensors only), autoresets within step, dict obs of torch tensors
+
+class MjlabMockEnv(MockEnv):
+    num_envs = 4
+    is_vector = True
+    is_vector_env = True
+
+    def __init__(self, num_envs = 4, auto_reset = True, seed = 0):
+        self.num_envs = num_envs
+        self.cfg = SimpleNamespace(auto_reset = auto_reset)
+        super().__init__(seed)
+
+    def reset(self, *, seed = None, env_ids = None, options = None):
+        if seed is not None:
+            self.seed(seed)
+        else:
+            self.reset_state()
+        return self.obs(), {}
+
+    def obs(self):
+        return dict(
+            policy = torch.randn(self.num_envs, self.obs_dim),
+            critic = torch.randn(self.num_envs, self.obs_dim * 2)
+        )
+
+    def step(self, action):
+        if not is_tensor(action):
+            raise TypeError(f'mjlab envs step on torch tensors, got {type(action).__name__}')
+
+        self.advance(action)
+        self.last_action = action
+
+        dones = torch.from_numpy(self.is_done())
+
+        if self.cfg.auto_reset:
+            for ind in np.where(dones.numpy())[0]:
+                self.t[ind] = 0
+                self.state[ind] = 0
+
+        return self.obs(), torch.ones(self.num_envs), dones, torch.zeros(self.num_envs, dtype = torch.bool), {}
+
+# live in the real module path so the adapter recognizes the mock
+
+MjlabMockEnv.__module__ = 'mjlab.envs.manager_based_rl_env'
+
 # maniskill — gymnasium-compliant, but always batched even at num_envs = 1:
 # obs / reward / dones are torch tensors with a leading batch dim, and the
 # env exposes single_action_space / single_observation_space (unbatched)

{env_ssl_wrapper-0.5.0 → env_ssl_wrapper-0.5.2}/env_ssl_wrapper/standardize/spaces.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 import numpy as np
 from torch.utils._pytree import tree_map
 
-from .helpers import default, exists, first_existing, get_attr, to_numpy
+from .helpers import default, exists, get_attr, to_numpy
 
 # stand-in for environments that do not expose gymnasium spaces
 

{env_ssl_wrapper-0.5.0 → env_ssl_wrapper-0.5.2}/env_ssl_wrapper/standardize/standardize_env_wrapper.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import torch
-from .helpers import EnvWrapper, exists, default, is_vectorized
+from .helpers import EnvWrapper, exists, is_vectorized
 from .utils import compose_env
 
 class StandardizeEnvWrapper(EnvWrapper):

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

@@ -12,6 +12,7 @@ from .helpers import (
     FINAL_OBSERVATION_KEYS,
     FINAL_OBSERVATION_MASK_KEYS,
     exists,
+    get_adapter,
     get_attr,
     is_scalar,
     unpack_vector_observations,
@@ -31,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):
@@ -71,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
 
@@ -96,10 +105,15 @@ class TensorWrapper(EnvWrapper):
         self.convert_in = convert_in
         self.convert_out = convert_out
         self.cast_obs_to_float = cast_obs_to_float
+
+        # torch-native sims (mjlab, isaac) take torch actions as-is — zero-copy passthrough
+
+        self.torch_native = get_adapter(env).torch_native
+
         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:
@@ -120,13 +134,19 @@ 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):
-        action = torch_to_numpy(action) if self.convert_in else action
+        if self.convert_in and not self.torch_native:
+            action = torch_to_numpy(action)
+
         obs, reward, terminated, truncated, info = self.env.step(action)
 
         if not self.convert_out:

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

@@ -2,9 +2,8 @@ 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_flatten, tree_map, tree_structure, tree_unflatten
+from torch.utils._pytree import tree_map
 
 from .helpers import (
     any_true,
@@ -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.0 → env_ssl_wrapper-0.5.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "env-ssl-wrapper"
-version = "0.5.0"
+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]