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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: env-ssl-wrapper
-Version: 0.4.7
+Version: 0.5.1
 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
@@ -37,6 +37,8 @@ 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: torch>=2.5

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

@@ -6,6 +6,7 @@ from .standardize import *
 
 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`
@@ -40,6 +41,8 @@ __all__ = [
     *standardize.__all__,
     'MemoryTraceWrapper',
     'ActionChunkWrapper',
+    'evaluate_actor',
+    'EpisodeStats',
 ]
 
 def __getattr__(name):

{env_ssl_wrapper-0.4.7 → env_ssl_wrapper-0.5.1}/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.1/env_ssl_wrapper/evaluate.py

@@ -0,0 +1,340 @@
+from __future__ import annotations
+
+from pathlib import Path
+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 .standardize.helpers import (
+    EnvWrapper,
+    env_autoresets,
+    env_num_envs,
+    env_render,
+    exists,
+    first_existing,
+    get_adapter,
+    get_attr,
+    is_array,
+    is_vectorized,
+    normalize_reset_out,
+    normalize_step_out,
+    to_numpy,
+)
+
+# results
+
+class EpisodeStats(NamedTuple):
+    returns: torch.Tensor
+    lengths: torch.Tensor
+    videos: tuple[str, ...] = ()
+
+    @property
+    def num_episodes(self):
+        return self.returns.numel()
+
+    @property
+    def mean(self):
+        return self.returns.mean().item()
+
+    @property
+    def std(self):
+        return self.returns.std(unbiased = False).item()
+
+    @property
+    def min(self):
+        return self.returns.min().item()
+
+    @property
+    def max(self):
+        return self.returns.max().item()
+
+    @property
+    def video_path(self):
+        return self.videos[0] if self.videos else None
+
+    def success_rate(self, threshold):
+        return (self.returns >= threshold).float().mean().item()
+
+    def summary(self):
+        return dict(episodes = self.num_episodes, mean = self.mean, std = self.std, min = self.min, max = self.max)
+
+# actor output -> action
+
+def resolve_actor(actor):
+    # a callable `.dist` takes precedence, sb3 style
+
+    module = actor if isinstance(actor, nn.Module) else None
+    policy = get_attr(actor, 'dist')
+
+    if not callable(policy):
+        policy = actor
+
+    assert callable(policy), 'actor must be callable, or expose a callable `.dist`'
+
+    return policy, module
+
+def is_distribution(out):
+    if isinstance(out, Distribution):
+        return True
+
+    if is_tensor(out):
+        return False
+
+    return callable(get_attr(out, 'sample')) and (exists(get_attr(out, 'mean')) or exists(get_attr(out, 'probs')))
+
+def to_action(env, out, deterministic = True):
+    # distributions collapse to their mode, raw floating logits over a discrete space to their argmax
+
+    if is_distribution(out):
+        if not deterministic:
+            return out.sample()
+
+        return out.probs.argmax(dim = -1) if exists(get_attr(out, 'probs')) else 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')
+
+        if exists(num_actions) and out.shape[-1] == num_actions:
+            return out.argmax(dim = -1) if deterministic else Categorical(logits = out).sample()
+
+    return out
+
+# states and actions
+
+def to_actor_obs(obs, device, batched):
+    def convert(x):
+        if not (is_array(x) or callable(get_attr(x, '__array__'))):
+            return x
+
+        tensor = x if is_tensor(x) else torch.as_tensor(np.asarray(x))
+
+        if tensor.is_floating_point():
+            tensor = tensor.float()
+
+        tensor = tensor.to(device)
+
+        return tensor if batched or tensor.ndim == 0 else tensor[None]
+
+    return tree_map(convert, obs)
+
+def to_env_action(env, action):
+    if not is_tensor(action):
+        action = torch.as_tensor(action)
+
+    device = get_attr(env, 'device')
+    action = action.to(device) if exists(device) else action
+
+    # 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]
+
+        return action.item() if action.ndim == 0 else action
+
+    # wrappers take batched torch actions - pad up to the declared rank
+
+    chunk_shape = get_attr(env, 'chunk_action_shape')
+
+    if exists(chunk_shape):
+        rank = len(chunk_shape) + 1
+    else:
+        space = first_existing(env, 'single_action_space', 'action_space')
+        rank = len(get_attr(space, 'shape', ())) + 1
+
+    while action.ndim < rank:
+        action = action[None]
+
+    return action
+
+
+# video
+
+def capture_frame(env):
+    render = get_attr(env, 'render')
+    frame = render() if callable(render) else None
+
+    if not exists(frame):
+        try:
+            frame = env_render(env, 256, 256)
+        except Exception:
+            return None
+
+    if is_tensor(frame):
+        frame = frame.detach().cpu().numpy()
+
+    frame = np.asarray(frame)
+
+    if frame.ndim == 4:
+        frame = frame[0]
+
+    if frame.ndim == 3 and frame.shape[0] in (3, 4) and frame.shape[-1] not in (3, 4):
+        frame = np.moveaxis(frame, 0, -1)
+
+    if frame.ndim == 3 and frame.shape[-1] == 4:
+        frame = frame[..., :3]
+
+    if frame.dtype != np.uint8:
+        scale = 255. if frame.max() <= 1. else 1.
+        frame = np.clip(frame * scale, 0., 255.).astype(np.uint8)
+
+    return frame
+
+def save_video(frames, path, fps = 30):
+    import imageio
+
+    path = Path(path)
+    path.parent.mkdir(parents = True, exist_ok = True)
+    imageio.mimsave(path, frames, fps = fps, macro_block_size = None)
+
+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}'))
+
+# evaluate
+
+@torch.no_grad()
+def evaluate_actor(
+    actor,
+    env,
+    episodes = 10,
+    max_steps = 1000,
+    seed = None,
+    goal = None,
+    device = None,
+    deterministic = True,
+    video_path = None,
+    video_episodes = 1,
+    fps = 30,
+    verbose = False
+):
+    # seed fixes the episode seeds, so checkpoints are compared on identical episodes
+    # video_path records the first video_episodes episodes
+
+    policy, module = resolve_actor(actor)
+
+    if not exists(device):
+        params = next(module.parameters(), None) if exists(module) else None
+        device = get_attr(params, 'device', torch.device('cpu'))
+
+    was_training = get_attr(module, 'training')
+
+    if exists(module):
+        module.eval()
+
+    if exists(goal):
+        goal = torch.as_tensor(goal) if not is_tensor(goal) else goal
+        goal = goal.to(device)
+        goal = goal[None] if goal.ndim == 1 else goal
+
+    num_envs = env_num_envs(env)
+    record = exists(video_path)
+
+    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())
+    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
+
+    returns = torch.zeros(num_envs)
+    lengths = torch.zeros(num_envs)
+    steps = torch.zeros(num_envs, dtype = torch.long)
+
+    done_returns, done_lengths, videos, frames = [], [], [], []
+
+    try:
+        while len(done_returns) < episodes:
+            if record and len(done_returns) < video_episodes:
+                frame = capture_frame(env)
+                if exists(frame):
+                    frames.append(frame)
+
+            state = to_actor_obs(obs, device, batched)
+            out = policy(state, goal) if exists(goal) else policy(state)
+            action = to_env_action(env, to_action(env, out, deterministic))
+
+            obs, reward, terminated, truncated, info = normalize_step_out(env.step(action))
+
+            # reward chunks collapse to one scalar per env
+
+            reward = torch.as_tensor(to_numpy(reward), dtype = torch.float32)
+
+            while reward.ndim > 1:
+                reward = reward.sum(dim = -1)
+
+            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)
+
+            lengths += info.get('chunk_length', 1) if isinstance(info, dict) else 1
+            steps += 1
+
+            done = to_numpy(terminated | truncated).astype(bool).reshape(-1)
+
+            if exists(max_steps):
+                done |= to_numpy(steps >= max_steps).reshape(-1)
+
+            for i in np.where(done)[0]:
+                if len(done_returns) >= episodes:
+                    break
+
+                episode = len(done_returns)
+                done_returns.append(returns[i].item())
+                done_lengths.append(lengths[i].item())
+
+                if verbose:
+                    print(f'episode {episode:3d} | return {done_returns[-1]:8.1f} | length {done_lengths[-1]:7.0f}', flush = True)
+
+                if record and episode < video_episodes:
+                    frame = capture_frame(env)
+                    if exists(frame):
+                        frames.append(frame)
+
+                    if frames:
+                        path = video_path_at(video_path, episode)
+                        save_video(frames, path, fps)
+                        videos.append(path)
+
+                returns[i] = 0.
+                lengths[i] = 0.
+                steps[i] = 0
+
+            if done.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))):
+                    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())
+    finally:
+        if exists(module) and exists(was_training):
+            module.train(was_training)
+
+    return EpisodeStats(
+        returns = torch.tensor(done_returns, dtype = torch.float32),
+        lengths = torch.tensor(done_lengths, dtype = torch.float32),
+        videos = tuple(videos)
+    )
+
+if __name__ == '__main__':
+    import gymnasium as gym
+    from env_ssl_wrapper import StandardizeEnvWrapper
+
+    env = StandardizeEnvWrapper(gym.make('CartPole-v1'))
+
+    stats = evaluate_actor(lambda obs: torch.zeros(obs.shape[0], 2), env, episodes = 10, seed = 0, verbose = True)
+
+    print(stats.summary())

{env_ssl_wrapper-0.4.7 → env_ssl_wrapper-0.5.1}/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,
@@ -121,6 +123,7 @@ __all__ = [
     'BaseEnvAdapter',
     'get_adapter',
     'register_adapter',
+    'MjlabAdapter',
     'MujocoWarpAdapter',
     'IsaacAdapter',
     'PyBulletAdapter',
@@ -149,6 +152,7 @@ __all__ = [
     'AutoresetVectorMockEnv',
     'PufferVectorMockEnv',
     'PufferTensorMockEnv',
+    'MjlabMockEnv',
     'ManiSkillMockEnv',
     'BraxMockEnv',
     'MjxMockEnv',

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

@@ -129,8 +129,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 +139,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.4.7 → env_ssl_wrapper-0.5.1}/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.4.7 → env_ssl_wrapper-0.5.1}/env_ssl_wrapper/standardize/auto_batched_wrapper.py

@@ -19,7 +19,6 @@ from .helpers import (
     is_tensor,
     is_vectorized,
 )
-from .spaces import space_from_action_spec
 
 # helper functions
 

{env_ssl_wrapper-0.4.7 → env_ssl_wrapper-0.5.1}/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.4.7 → env_ssl_wrapper-0.5.1}/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.4.7 → env_ssl_wrapper-0.5.1}/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.4.7 → env_ssl_wrapper-0.5.1}/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,
@@ -96,6 +97,11 @@ 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):
@@ -126,7 +132,9 @@ class TensorWrapper(EnvWrapper):
         return contract(self.cast(t), to_float = to_float)
 
     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.4.7 → env_ssl_wrapper-0.5.1}/env_ssl_wrapper/standardize/vector.py

@@ -4,7 +4,7 @@ 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,

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

@@ -1,6 +1,6 @@
 [project]
 name = "env-ssl-wrapper"
-version = "0.4.7"
+version = "0.5.1"
 description = "One torch-native interface for any MDP environment"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }
@@ -27,6 +27,8 @@ dependencies = [
     "discrete-continuous-embed-readout",
     "einx>=0.3.0",
     "einops>=0.8.1",
+    "imageio",
+    "imageio-ffmpeg",
     "memmap-replay-buffer",
     "torch>=2.5",
     "torch-einops-utils>=0.0.29",