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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: env-ssl-wrapper
-Version: 0.5.2
+Version: 0.5.3
 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

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

@@ -19,7 +19,7 @@ from .standardize.helpers import (
     first_existing,
     get_adapter,
     get_attr,
-    is_array,
+    is_array_like,
     is_vectorized,
     normalize_reset_out,
     normalize_step_out,
@@ -120,7 +120,7 @@ def to_action(env, out, deterministic = True):
 
 def to_actor_obs(obs, device, batched):
     def convert(x):
-        if not (is_array(x) or callable(get_attr(x, '__array__'))):
+        if not is_array_like(x):
             return x
 
         tensor = x if is_tensor(x) else torch.as_tensor(np.asarray(x))

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

@@ -77,6 +77,16 @@ from .mocks import (
     TrifingerMockEnv,
     HabitatMockEnv,
     TupleObsMockEnv,
+    RLBenchObservation,
+    RLBenchMockEnv,
+    OmniGibsonMockEnv,
+    RoboMimicMockEnv,
+    MyoSuiteMockEnv,
+    GymnasiumRoboticsMockEnv,
+    LeRobotMockEnv,
+    CalvinMockEnv,
+    PushTMockEnv,
+    DroneAviaryMockEnv,
     JaxArray,
 )
 
@@ -84,6 +94,12 @@ from .utils import wrap_env, compose_env
 from .helpers import (
     TransformObservationWrapper,
     ObservationWrapper,
+    is_array,
+    is_foreign_array,
+    is_array_like,
+    normalize_obs,
+    normalize_reset_out,
+    normalize_step_out,
     has_final_observation,
     maybe_get_final_observation,
     get_final_observation,
@@ -143,6 +159,12 @@ __all__ = [
     'action_space_bounds',
     'action_dim_of',
     'obs_dim_of',
+    'is_array',
+    'is_foreign_array',
+    'is_array_like',
+    'normalize_obs',
+    'normalize_reset_out',
+    'normalize_step_out',
     'MockEnv',
     'GymnasiumMockEnv',
     'GymnasiumDiscreteMockEnv',
@@ -161,6 +183,16 @@ __all__ = [
     'TrifingerMockEnv',
     'HabitatMockEnv',
     'TupleObsMockEnv',
+    'RLBenchObservation',
+    'RLBenchMockEnv',
+    'OmniGibsonMockEnv',
+    'RoboMimicMockEnv',
+    'MyoSuiteMockEnv',
+    'GymnasiumRoboticsMockEnv',
+    'LeRobotMockEnv',
+    'CalvinMockEnv',
+    'PushTMockEnv',
+    'DroneAviaryMockEnv',
     'JaxArray',
     'wrap_env',
     'compose_env',

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

@@ -84,7 +84,7 @@ class ActionTransformWrapper(EnvWrapper):
         # rescale canonical range (default (0, 1), beta-friendly) to per-env
         # bounds, leaving unbounded dimensions untouched
 
-        if self.bounds is None:
+        if not exists(self.bounds):
             return t
 
         low, high = self.bounds

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

@@ -13,6 +13,31 @@ from .helpers import (
 )
 from .spaces import space_from_action_spec
 
+# helpers
+
+def render_with_dims(render, height, width, camera = None):
+    # modern render(height, width, camera) vs legacy bare render()
+
+    try:
+        return render(height = height, width = width, camera = camera)
+    except TypeError:
+        return render()
+
+def render_sim(sim, height, width, camera = None):
+    # sims vary — robosuite-style (camera_name) vs bare no-kwargs render
+
+    render = get_attr(sim, 'render')
+
+    if not callable(render):
+        return None
+
+    kwargs = dict(camera_name = camera) if exists(camera) else {}
+
+    try:
+        return render(height = height, width = width, **kwargs)
+    except TypeError:
+        return None
+
 # base adapter
 
 class BaseEnvAdapter:
@@ -58,13 +83,14 @@ class BaseEnvAdapter:
             _, _, rgba, _, _ = client.getCameraImage(width, height, renderer = renderer)
             return rgba[..., :3]
 
-        sim = get_attr(self.env, 'sim')
-        if exists(sim) and callable(get_attr(sim, 'render')):
-            kwargs = dict(camera_name = camera) if exists(camera) else {}
-            return sim.render(height = height, width = width, **kwargs)
+        frame = render_sim(get_attr(self.env, 'sim'), height, width, camera)
+        if exists(frame):
+            return frame
 
-        if callable(get_attr(self.env, 'render')):
-            return self.env.render()
+        render = get_attr(self.env, 'render')
+
+        if callable(render):
+            return render_with_dims(render, height, width, camera)
 
         return None
 
@@ -127,8 +153,11 @@ class WrapperAdapter(BaseEnvAdapter):
         self.inner_adapter.seed(seed)
 
     def render(self, height: int, width: int, camera = None):
-        if callable(get_attr(self.env, 'render')):
-            return self.env.render()
+        render = get_attr(self.env, 'render')
+
+        if callable(render):
+            return render_with_dims(render, height, width, camera)
+
         return self.inner_adapter.render(height, width, camera)
 
     def close(self):
@@ -296,25 +325,37 @@ class PufferLibAdapter(BaseEnvAdapter):
     def autoresets(self) -> bool:
         return truthy_attr(first_existing(self.env, 'autoreset', 'autoresets', 'autoreset_mode'))
 
-# robotics adapter (robosuite, maniskill, metaworld, trifinger, habitat)
+# robotics adapter (robosuite, maniskill, metaworld, trifinger, habitat, rlbench,
+# omnigibson, robomimic, myosuite, gymnasium-robotics, lerobot, ...)
+# sits before dm_control — several robotics sims expose `physics` but are not dm_control
+
+ROBOTICS_KEYWORDS = (
+    'robosuite', 'mani_skill', 'maniskill', 'metaworld', 'trifinger', 'habitat',
+    'omnigibson', 'igibson', 'robomimic', 'myosuite', 'gymnasiumrobotics',
+    'gymnasium_robotics', 'gym_robotics', 'goalenv', 'lerobot',
+    'gym_aloha', 'gym_pusht', 'aloha', 'pusht', 'sapien', 'rlbench', 'pyrep',
+    'calvin', 'language_table', 'drone', 'aviary', 'softgym', 'plasticinelab'
+)
 
 class RoboticsAdapter(BaseEnvAdapter):
     @classmethod
     def matches(cls, env):
-        mod = getattr(type(env), '__module__', '')
-        name = type(env).__name__
-        robotics_names = ('Robosuite', 'ManiSkill', 'MetaWorld', 'Trifinger', 'Habitat')
-        robotics_mods = ('robosuite', 'mani_skill', 'maniskill', 'metaworld', 'habitat')
+        mod = getattr(type(env), '__module__', '').lower()
+        name = type(env).__name__.lower()
         has_sim_render = exists(get_attr(env, 'sim')) and callable(get_attr(get_attr(env, 'sim'), 'render'))
-        return any(r in name for r in robotics_names) or any(r in mod.lower() for r in robotics_mods) or has_sim_render
+        return any(k in name or k in mod for k in ROBOTICS_KEYWORDS) or has_sim_render
 
     def render(self, height: int, width: int, camera = None):
-        sim = get_attr(self.env, 'sim')
-        if exists(sim) and callable(get_attr(sim, 'render')):
-            kwargs = dict(camera_name = camera) if exists(camera) else {}
-            return sim.render(height = height, width = width, **kwargs)
-        if callable(get_attr(self.env, 'render')):
-            return self.env.render()
+        frame = render_sim(get_attr(self.env, 'sim'), height, width, camera)
+
+        if exists(frame):
+            return frame
+
+        render = get_attr(self.env, 'render')
+
+        if callable(render):
+            return render_with_dims(render, height, width, camera)
+
         return super().render(height, width, camera)
 
     @property
@@ -426,9 +467,9 @@ ADAPTER_REGISTRY: list[type[BaseEnvAdapter]] = [
     MujocoWarpAdapter,
     IsaacAdapter,
     PyBulletAdapter,
+    RoboticsAdapter,
     DMControlAdapter,
     PufferLibAdapter,
-    RoboticsAdapter,
     GymnasiumAdapter,
     LegacyGymAdapter,
     DefaultAdapter,

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

@@ -65,17 +65,18 @@ def get_action_space(env):
     return get_adapter(env).action_space, False
 
 def action_shape_tree(space):
-    # canonical shapes parallel to action structure
+    # canonical shapes parallel to action structure — raw dict / list / tuple
+    # action spaces are handled as well as gymnasium's composite spaces
 
     if not exists(space):
         return None
 
-    subspaces = get_attr(space, 'spaces')
+    subspaces = space if isinstance(space, (dict, list, tuple)) else get_attr(space, 'spaces')
 
-    if exists(subspaces):
-        if isinstance(subspaces, dict):
-            return {key: action_shape_tree(subspace) for key, subspace in subspaces.items()}
+    if isinstance(subspaces, dict):
+        return {key: action_shape_tree(subspace) for key, subspace in subspaces.items()}
 
+    if exists(subspaces):
         return [action_shape_tree(subspace) for subspace in subspaces]
 
     return get_attr(space, 'shape')

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

@@ -36,7 +36,7 @@ def zero_mask(x, mask, fill_scalar = None):
     arr = np.asarray(x)
 
     if arr.ndim == 0:
-        if fill_scalar is not None and bool(np.asarray(mask).any()):
+        if exists(fill_scalar) and bool(np.asarray(mask).any()):
             return fill_scalar
         return x
 
@@ -51,6 +51,7 @@ def back_to_mask_type(dones, newly):
 
 def merge_final(current, value, mask):
     if is_tensor(current):
+        value = torch.as_tensor(value, device = current.device, dtype = current.dtype)
         return torch.where(broadcast_mask(mask, current), value, current)
 
     curr = np.asarray(current)
@@ -59,7 +60,7 @@ def merge_final(current, value, mask):
         return current
 
     out = curr.copy()
-    out[mask] = np.asarray(value)[mask]
+    out[mask] = to_numpy(value)[mask]
     return out
 
 # class
@@ -90,7 +91,7 @@ class EpisodePaddingWrapper(EnvWrapper):
             dones = dones_of(terminated, truncated)
             mask = to_numpy(dones).astype(bool)
 
-            if self._is_done is None or len(self._is_done) != len(mask):
+            if not exists(self._is_done) or len(self._is_done) != len(mask):
                 self._is_done = np.zeros(len(mask), dtype = bool)
 
             # autoreset envs revive every done slot, so each done is a new
@@ -106,7 +107,7 @@ class EpisodePaddingWrapper(EnvWrapper):
                     final_val = first_existing(info, 'final_observation', 'final_obs')
                     value = final_val if exists(final_val) else self._last_obs
 
-                    if self._final_obs is None:
+                    if not exists(self._final_obs):
                         self._final_obs = tree_map(copy_leaf, value)
                     else:
                         self._final_obs = tree_map(partial(merge_final, mask = newly), self._final_obs, value)

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

@@ -59,14 +59,27 @@ def get_batch_size(tree) -> int | None:
 
     first = leaves[0]
 
-    if is_array(first):
-        return len(first) if first.ndim > 0 else None
+    if isinstance(first, (str, bytes)):
+        return None
+
+    ndim = get_attr(first, 'ndim')
+
+    if exists(ndim):
+        return len(first) if ndim > 0 else None
 
     return len(first) if exists(get_attr(first, '__len__')) else None
 
 def is_array(v):
     return is_tensor(v) or isinstance(v, np.ndarray)
 
+def is_foreign_array(v):
+    # array-like types from foreign frameworks (e.g. jax.Array, pil Image)
+    # that implement numpy's __array__ conversion protocol
+    return not is_array(v) and callable(get_attr(v, '__array__'))
+
+def is_array_like(v):
+    return is_array(v) or is_foreign_array(v)
+
 def to_numpy(t):
     return t.detach().cpu().numpy() if is_tensor(t) else np.asarray(t)
 
@@ -101,28 +114,97 @@ def zero_like(x):
     arr = np.asarray(x)
     return np.zeros_like(arr, dtype = bool) if arr.ndim > 0 else False
 
+def normalize_obs(obs):
+    # raw arrays / scalars / text / containers pass through — custom observation
+    # objects (rlbench, ...) flatten via get_low_dim_data or unpack their
+    # public array, scalar, or container attributes
+
+    if is_array_like(obs) or is_scalar(obs) or isinstance(obs, (str, bytes, dict, list, tuple)):
+        return obs
+
+    low_dim = get_attr(obs, 'get_low_dim_data')
+
+    if callable(low_dim):
+        try:
+            flattened = low_dim()
+            if exists(flattened):
+                return flattened
+        except Exception:
+            pass
+
+    out = {}
+    for key in dir(obs):
+        if key.startswith('_'):
+            continue
+
+        value = get_attr(obs, key)
+
+        if callable(value):
+            continue
+
+        if is_array_like(value) or is_scalar(value) or isinstance(value, (dict, list, tuple)):
+            out[key] = value
+
+    return out if out else obs
+
+def contains_array(x):
+    # deep array-leaf check over nested containers — used to tell an obs tree
+    # apart from a plain info dict when language instructions come first
+
+    if isinstance(x, dict):
+        return any(contains_array(v) for v in x.values())
+
+    if isinstance(x, (list, tuple)):
+        return any(contains_array(v) for v in x)
+
+    return is_array_like(x)
+
+def is_instruction(x):
+    if isinstance(x, str):
+        return True
+
+    return isinstance(x, (list, tuple)) and len(x) > 0 and all(isinstance(d, str) for d in x)
+
+def as_descriptions(x):
+    return [x] if isinstance(x, str) else list(x)
+
 def normalize_reset_out(out):
     if is_time_step(out):
-        return out.observation, {}
+        return normalize_obs(out.observation), {}
+
+    if isinstance(out, tuple):
+        if len(out) == 2:
+            first, second = out
+
+            # language-conditioned robotics (e.g. rlbench, calvin, language-table):
+            # (instruction(s), obs) — a plain info dict in second means text observation
 
-    if isinstance(out, tuple) and len(out) == 2:
-        obs, info = out
-        return obs, {} if info is None else (info if isinstance(info, dict) else {})
+            if is_instruction(first) and (not isinstance(second, dict) or contains_array(second)):
+                return normalize_obs(second), dict(descriptions = as_descriptions(first))
 
-    return out, {}
+            info = second if isinstance(second, dict) else {}
+            return normalize_obs(first), info
+
+        if len(out) == 3 and is_instruction(out[0]):
+            obs = normalize_obs(out[1])
+            info = out[2] if isinstance(out[2], dict) else {}
+            return obs, {**info, 'descriptions': as_descriptions(out[0])}
+
+    return normalize_obs(out), {}
 
 def normalize_step_out(out):
     if is_time_step(out):
         last = out.last() if callable(get_attr(out, 'last')) else out.step_type == 2
-        return out.observation, out.reward, last, False, dict(discount = out.discount)
+        return normalize_obs(out.observation), out.reward, last, False, dict(discount = out.discount)
 
     if len(out) == 5:
-        return out
+        obs, reward, terminated, truncated, info = out
+        return normalize_obs(obs), reward, terminated, truncated, info if isinstance(info, dict) else {}
 
     if len(out) in (3, 4):
         obs, reward, done, *rest = out
         info = rest[0] if rest and isinstance(rest[0], dict) else {}
-        return obs, reward, done, zero_like(done), info
+        return normalize_obs(obs), reward, done, zero_like(done), info
 
     raise ValueError(f'could not standardize step output of length {len(out)}')
 
@@ -135,7 +217,7 @@ def _zero_leaf(x):
         return False
     if isinstance(x, (int, float, np.number)):
         return np.zeros_like(x)
-    if exists(get_attr(x, '__array__')):
+    if is_foreign_array(x):
         return np.zeros_like(np.asarray(x))
     return 0
 
@@ -173,11 +255,11 @@ def unpack_vector_observations(arr):
     if isinstance(arr, np.ndarray) and arr.dtype != object:
         return arr
 
-    sample = next((x for x in arr if x is not None), None)
-    if sample is None:
+    sample = next((x for x in arr if exists(x)), None)
+    if not exists(sample):
         return arr
 
-    trees = [tree_map(_zero_leaf, sample) if x is None else x for x in arr]
+    trees = [tree_map(_zero_leaf, sample) if not exists(x) else x for x in arr]
     return stack_trees(trees)
 
 # environment probes
@@ -336,12 +418,12 @@ class TransformObservationWrapper(EnvWrapper):
 
     def reset(self, **kwargs):
         obs, info = self.env.reset(**kwargs)
-        obs = self.transform_obs(obs)
+        obs = self.transform(obs, done = None)
 
         if isinstance(info, dict):
             for key in FINAL_OBSERVATION_KEYS:
                 if key in info:
-                    info[key] = self.transform_obs(info[key])
+                    info[key] = self.transform(info[key], done = None)
 
         return obs, info
 
@@ -349,17 +431,16 @@ class TransformObservationWrapper(EnvWrapper):
         obs, reward, terminated, truncated, info = self.env.step(action)
         done = dones_of(terminated, truncated) if self.autoresets else None
 
-        out = self.transform_obs(obs, done = done) if self.takes_done else self.transform_obs(obs)
+        out = self.transform(obs, done = done)
 
         if isinstance(info, dict):
             for key in FINAL_OBSERVATION_KEYS:
                 if key in info:
                     if not self.takes_done:
-                        info[key] = self.transform_obs(info[key])
+                        info[key] = self.transform(info[key])
                     elif not self.autoresets:
                         info[key] = out
 
         return out, reward, terminated, truncated, info
 
 ObservationWrapper = TransformObservationWrapper
-

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

@@ -31,7 +31,7 @@ def render_frame(env, image_size = (64, 64), camera = None):
     img = env_render(env, height, width, camera)
 
     if not exists(img):
-        if env_render_mode(env) is None:
+        if not exists(env_render_mode(env)):
             raise ValueError(
                 'env must be created with render_mode = "rgb_array", '
                 'e.g. gym.make(id, render_mode = "rgb_array")'

{env_ssl_wrapper-0.5.2 → env_ssl_wrapper-0.5.3}/env_ssl_wrapper/standardize/mocks.py

@@ -7,6 +7,8 @@ import numpy as np
 import torch
 from torch import is_tensor
 
+from .helpers import default, exists
+
 # mock sims — stand-ins for real-world environments with varying MDP interfaces
 # all share the underlying dynamics; each emulates the quirks of a real sim
 
@@ -60,6 +62,10 @@ class DiscreteSpace:
     def sample(self):
         return np.random.randint(self.n)
 
+class DictSpace:
+    def __init__(self, **spaces):
+        self.spaces = spaces
+
 # sim-specific render surfaces — emulating how each sim actually produces images
 
 class ActionSpec:
@@ -125,16 +131,20 @@ class MockEnv:
         return False
 
     def seed(self, seed = 0):
+        seed = default(seed, 0)
         self.rng = np.random.default_rng(seed)
         torch.manual_seed(seed)
         self.reset_state()
 
     def reset(self, seed = None, **kwargs):
-        if seed is not None:
+        self.restart(seed)
+        return self.obs(), {}
+
+    def restart(self, seed = None):
+        if exists(seed):
             self.seed(seed)
         else:
             self.reset_state()
-        return self.obs(), {}
 
     def reset_state(self):
         self.t = np.zeros(self.num_envs, dtype = int) if self.is_vector else 0
@@ -406,10 +416,7 @@ class MjlabMockEnv(MockEnv):
         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()
+        self.restart(seed)
         return self.obs(), {}
 
     def obs(self):
@@ -487,10 +494,7 @@ class ManiSkillMockEnv(MockEnv):
         return dict(state = state) if self.obs_mode == 'state_dict' else state
 
     def reset(self, seed = None, options = None):
-        if seed is not None:
-            self.seed(seed)
-        else:
-            self.reset_state()
+        self.restart(seed)
         return self.obs(), dict(reconfigure = False)
 
     def step(self, action):
@@ -515,6 +519,9 @@ class JaxArray:
     def __array__(self, dtype = None):
         return np.asarray(self._arr, dtype = dtype)
 
+    def __getitem__(self, item):
+        return JaxArray(self._arr[item])
+
     @property
     def dtype(self):
         return self._arr.dtype
@@ -550,19 +557,15 @@ class BraxMockEnv(MockEnv):
         self.advance(action)
         return JaxArray(self.obs()), JaxArray(np.ones(self.num_envs)), JaxArray(self.is_done()), {}
 
-# mujoco-mjx — gymnasium-style 5-tuple step, (obs, info) reset, single env, and
-# every stream (obs / reward / dones) is a jax array: the dialect of envs built
-# on the mjx jax physics backend
+# mujoco-mjx — gymnasium-style 5-tuple step, (obs, info) reset, single env,
+# every stream (obs / reward / dones) a jax array
 
 class MjxMockEnv(MockEnv):
     obs_dim = 8
     action_dim = 6
 
     def reset(self, seed = None, **kwargs):
-        if seed is not None:
-            self.seed(seed)
-        else:
-            self.reset_state()
+        self.restart(seed)
         return JaxArray(self.obs()), {}
 
     def step(self, action):
@@ -652,3 +655,232 @@ class TupleObsMockEnv(MockEnv):
     def step(self, action):
         self.advance(action)
         return self.obs(), np.ones(self.num_envs), self.is_done(), np.zeros(self.num_envs, dtype = bool), {}
+
+# rlbench — language-conditioned manipulation: reset -> (descriptions, obs),
+# step -> (obs, reward, terminate); obs unpacks via get_low_dim_data()
+
+class RLBenchObservation:
+    def __init__(self, joint_pos, gripper, wrist_rgb):
+        self.joint_positions = joint_pos
+        self.gripper_open = gripper
+        self.wrist_rgb = wrist_rgb
+
+    def get_low_dim_data(self):
+        return np.concatenate([self.joint_positions, np.array([self.gripper_open])])
+
+class RLBenchMockEnv(MockEnv):
+    obs_dim = 8
+    action_dim = 8
+    max_steps = 40
+
+    def __init__(self, task_name = 'open_drawer', seed = 0):
+        self.task_name = task_name
+        self.descriptions = [f'{task_name.replace("_", " ")} task']
+        super().__init__(seed)
+
+    def rlbench_obs(self):
+        return RLBenchObservation(
+            joint_pos = self.obs()[:7],
+            gripper = 1.0,
+            wrist_rgb = np.zeros((64, 64, 3), dtype = np.uint8)
+        )
+
+    def reset(self, seed = None, **kwargs):
+        self.restart(seed)
+        return self.descriptions, self.rlbench_obs()
+
+    def step(self, action):
+        self.advance(action)
+        return self.rlbench_obs(), 1.0, self.is_done()
+
+# omnigibson (behavior-1k) — obs-dict-only reset, 5-tuple step, sim render
+
+class OmniGibsonMockEnv(MockEnv):
+    obs_dim = 14
+    action_dim = 6
+    max_steps = 100
+
+    def __init__(self, seed = 0):
+        super().__init__(seed)
+        self.sim = FakeSim()
+
+    def obs_dict(self):
+        return dict(
+            robot0_proprio = self.obs(),
+            rgb = np.zeros((64, 64, 3), dtype = np.uint8)
+        )
+
+    def reset(self, seed = None, **kwargs):
+        self.restart(seed)
+        return self.obs_dict()
+
+    def step(self, action):
+        self.advance(action)
+        return self.obs_dict(), 1.0, self.is_done(), False, {}
+
+# robomimic — structured obs dict, obs-only reset, 4-tuple step
+
+class RoboMimicMockEnv(MockEnv):
+    obs_dim = 10
+    action_dim = 7
+    max_steps = 50
+
+    def obs_dict(self):
+        return dict(
+            robot0_eef_pos = self.obs()[:3],
+            robot0_eef_quat = np.array([0., 0., 0., 1.]),
+            object = self.obs()[3:]
+        )
+
+    def reset(self, seed = None, **kwargs):
+        self.restart(seed)
+        return self.obs_dict()
+
+    def step(self, action):
+        self.advance(action)
+        return self.obs_dict(), 1.0, self.is_done(), {}
+
+# myosuite — musculoskeletal mujoco robotics: high-dim muscle states,
+# physics render, 5-tuple step
+
+class MyoSuiteMockEnv(MockEnv):
+    obs_dim = 36
+    action_dim = 12
+    max_steps = 80
+
+    def __init__(self, seed = 0):
+        super().__init__(seed)
+        self.physics = FakePhysics()
+
+    def step(self, action):
+        self.advance(action)
+        return self.obs(), 1.0, self.is_done(), False, {}
+
+# gymnasium-robotics (farama goal envs: fetch, shadow hand, franka kitchen) —
+# dict obs (observation / achieved_goal / desired_goal), exposes compute_reward
+
+class GymnasiumRoboticsMockEnv(MockEnv):
+    obs_dim = 10
+    goal_dim = 3
+    action_dim = 4
+    max_steps = 50
+
+    @property
+    def observation_space(self):
+        return DictSpace(
+            observation = Space((self.obs_dim,)),
+            achieved_goal = Space((self.goal_dim,)),
+            desired_goal = Space((self.goal_dim,)),
+        )
+
+    def goal_dict(self):
+        return dict(
+            observation = self.obs(),
+            achieved_goal = self.obs()[:self.goal_dim],
+            desired_goal = np.zeros(self.goal_dim)
+        )
+
+    def reset(self, seed = None, **kwargs):
+        self.restart(seed)
+        return self.goal_dict(), {}
+
+    def compute_reward(self, achieved_goal, desired_goal, info):
+        return -np.linalg.norm(achieved_goal - desired_goal, axis = -1)
+
+    def step(self, action):
+        self.advance(action)
+        obs = self.goal_dict()
+        reward = float(self.compute_reward(obs['achieved_goal'], obs['desired_goal'], {}))
+        return obs, reward, self.is_done(), False, {}
+
+# lerobot (hugging face aloha / pusht) — bimanual multi-camera robotics:
+# dict obs with visual pixels + robot state, 5-tuple step
+
+class LeRobotMockEnv(MockEnv):
+    obs_dim = 14
+    action_dim = 14
+    max_steps = 60
+
+    def obs_dict(self):
+        return dict(
+            agent_pos = self.obs(),
+            pixels = dict(
+                top = np.zeros((64, 64, 3), dtype = np.uint8),
+                wrist = np.zeros((64, 64, 3), dtype = np.uint8)
+            )
+        )
+
+    def reset(self, seed = None, **kwargs):
+        self.restart(seed)
+        return self.obs_dict(), {}
+
+    def step(self, action):
+        self.advance(action)
+        return self.obs_dict(), 1.0, self.is_done(), False, {}
+
+# calvin (language-conditioned multi-task 3d manipulation) — string instruction
+# reset with structured vision + proprio dict, 5-tuple step
+
+class CalvinMockEnv(MockEnv):
+    obs_dim = 15
+    action_dim = 7
+    max_steps = 40
+
+    def __init__(self, task_instruction = 'open drawer and push red block', seed = 0):
+        super().__init__(seed = seed)
+        self.task_instruction = task_instruction
+
+    def obs_dict(self):
+        return dict(
+            robot_obs = self.obs(),
+            rgb_obs = dict(
+                rgb_static = np.zeros((64, 64, 3), dtype = np.uint8),
+                rgb_gripper = np.zeros((64, 64, 3), dtype = np.uint8)
+            )
+        )
+
+    def reset(self, seed = None, **kwargs):
+        self.restart(seed)
+        return self.task_instruction, self.obs_dict()
+
+    def step(self, action):
+        self.advance(action)
+        return self.obs_dict(), 1.0, self.is_done(), False, {}
+
+# pusht (diffusion policy planar pushing benchmark) — 2-dim eef actions in (0, 512),
+# dict observation with agent_pos and block_pose
+
+class PushTMockEnv(MockEnv):
+    obs_dim = 2
+    action_dim = 2
+    max_steps = 50
+
+    @property
+    def action_space(self):
+        return Space((self.action_dim,), 0., 512.)
+
+    def obs_dict(self):
+        return dict(
+            agent_pos = self.obs(),
+            block_pose = np.zeros(3)
+        )
+
+    def reset(self, seed = None, **kwargs):
+        self.restart(seed)
+        return self.obs_dict(), {}
+
+    def step(self, action):
+        self.advance(action)
+        return self.obs_dict(), 1.0, self.is_done(), False, {}
+
+# gym-pybullet-drones / aerial robotics — 4-dim motor rpm action space [-1, 1],
+# 12-dim kinematics observation (pos, rpy, vel, ang_vel)
+
+class DroneAviaryMockEnv(MockEnv):
+    obs_dim = 12
+    action_dim = 4
+    max_steps = 60
+
+    def step(self, action):
+        self.advance(action)
+        return self.obs(), 1.0, self.is_done(), False, {}

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

@@ -43,6 +43,9 @@ def space_from_action_spec(env):
 
 # action-space duck-typing across gymnasium, dm_control, and custom envs
 
+def shape_dim(shape) -> int:
+    return int(np.prod(shape)) if len(shape) > 0 else 1
+
 def action_space_dim(space) -> int | None:
     if not exists(space):
         return None
@@ -55,11 +58,30 @@ def action_space_dim(space) -> int | None:
     if exists(nvec):
         return int(np.prod(nvec))
 
+    shape = get_attr(space, 'shape')
+    return shape_dim(shape) if exists(shape) else None
+
+def space_dim(space) -> int | None:
+    # flat dim of an observation space tree — composite spaces sum their children
+
+    if not exists(space):
+        return None
+
     shape = get_attr(space, 'shape')
     if exists(shape):
-        return int(np.prod(shape)) if len(shape) > 0 else 1
+        return shape_dim(shape)
+
+    subspaces = get_attr(space, 'spaces')
+    if not exists(subspaces):
+        return None
+
+    child_spaces = subspaces.values() if isinstance(subspaces, dict) else subspaces
+    child_dims = [space_dim(child) for child in child_spaces]
+
+    if not child_dims or not all(map(exists, child_dims)):
+        return None
 
-    return None
+    return sum(child_dims)
 
 def action_space_is_discrete(space) -> bool:
     return exists(get_attr(space, 'n'))
@@ -92,9 +114,4 @@ def action_dim_of(env) -> int | None:
 def obs_dim_of(env) -> int | None:
     from .helpers import get_adapter
     space = get_adapter(env).observation_space
-    shape = get_attr(space, 'shape')
-
-    if exists(shape):
-        return int(np.prod(shape)) if len(shape) > 0 else 1
-
-    return get_attr(env, 'obs_dim')
+    return default(space_dim(space), get_attr(env, 'obs_dim'))

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

@@ -13,7 +13,7 @@ from .helpers import (
     FINAL_OBSERVATION_MASK_KEYS,
     exists,
     get_adapter,
-    get_attr,
+    is_foreign_array,
     is_scalar,
     unpack_vector_observations,
 )
@@ -26,7 +26,7 @@ def to_torch_leaf(t, device, cast_obs_to_float = True):
             t = from_numpy(t.copy())
         elif is_scalar(t):
             t = tensor(t)
-        elif exists(get_attr(t, '__array__')):
+        elif is_foreign_array(t):
             t = from_numpy(np.asarray(t))
         else:
             return t
@@ -98,7 +98,7 @@ class TensorWrapper(EnvWrapper):
 
         # deprecated alias, kept for backwards compatibility
 
-        if cast_float64_to_float32 is not None:
+        if exists(cast_float64_to_float32):
             cast_obs_to_float = cast_float64_to_float32
 
         self.device = torch_device(device)
@@ -144,8 +144,12 @@ class TensorWrapper(EnvWrapper):
         return tree_map(convert, t)
 
     def step(self, action):
-        if self.convert_in and not self.torch_native:
-            action = torch_to_numpy(action)
+        if self.convert_in:
+            if not self.torch_native:
+                action = torch_to_numpy(action)
+            elif not is_tensor(action):
+                # torch-native sims take torch on device — lift numpy / foreign actions
+                action = numpy_to_torch(action, device = self.device, cast_obs_to_float = False)
 
         obs, reward, terminated, truncated, info = self.env.step(action)
 

{env_ssl_wrapper-0.5.2 → env_ssl_wrapper-0.5.3}/env_ssl_wrapper/standardize/time_limit_wrapper.py

@@ -18,18 +18,15 @@ from .helpers import (
 )
 
 def back_to_like(t, numpy_arr):
-    if is_tensor(t):
-        return torch.from_numpy(numpy_arr).to(t.device)
+    # scalar dones stay scalar
 
-    if isinstance(t, np.ndarray):
-        return numpy_arr
+    if is_tensor(t) or isinstance(t, np.ndarray):
+        res = torch.from_numpy(numpy_arr).to(t.device) if is_tensor(t) else numpy_arr
+        return res.squeeze() if t.ndim == 0 else res
 
     # python scalars (single-env dones) collapse to a scalar; foreign
     # array-likes (jax) keep the numpy array so the batch dim survives
 
-    if numpy_arr.ndim == 0:
-        return bool(numpy_arr)
-
     if numpy_arr.size == 1:
         return bool(numpy_arr.reshape(-1)[0])
 

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

@@ -16,7 +16,7 @@ from .helpers import (
     truthy_attr,
     _zero_leaf,
 )
-from .spaces import action_dim_of
+from .spaces import action_dim_of, obs_dim_of
 from .standardize_wrapper import StandardizeWrapper
 
 # leaf helpers
@@ -216,6 +216,7 @@ class MultiprocessingVecEnv:
         self.action_space = self._get_attr('action_space')
         self.observation_space = self._get_attr('observation_space')
         self.action_dim = action_dim_of(self)
+        self.obs_dim = obs_dim_of(self)
 
         self.seed(seed)
 

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

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