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

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

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

@@ -11,6 +11,7 @@ from .standardize.helpers import (
     default,
     any_true,
     dones_of,
+    env_takes_torch,
     first_existing,
     get_adapter,
     get_attr,
@@ -29,16 +30,20 @@ def chunk_steps(actions, axis):
 def unbatch_step(action):
     # chunks always carry a num_envs axis - drop it for plain single envs
     if is_tensor(action):
-        return action.squeeze(0)
+        action = action.squeeze(0)
+        return action.item() if action.ndim == 0 and not action.is_floating_point() else action
 
     if isinstance(action, np.ndarray):
-        return np.squeeze(action, axis = 0) if action.ndim > 0 else action
+        action = np.squeeze(action, axis = 0) if action.ndim > 0 else action
+        return action.item() if action.ndim == 0 and np.issubdtype(action.dtype, np.integer) else action
 
     return action
 
 # wrapper
 
 class ActionChunkWrapper(EnvWrapper):
+    priority = 100
+
     """
     Open-loop action chunking - executes a chunk of actions in one step,
     temporally compressing the environment by the chunk length.
@@ -71,7 +76,7 @@ class ActionChunkWrapper(EnvWrapper):
         self.action_shape = tuple(get_attr(action_space, 'shape', ()) or ())
 
         self.expects_batch = is_vectorized(env)
-        self.convert_in = not isinstance(env, EnvWrapper) and not get_adapter(env).torch_native
+        self.convert_in = not env_takes_torch(env)
 
     @property
     def chunk_action_shape(self):

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

@@ -12,14 +12,16 @@ from torch.utils._pytree import tree_flatten, tree_map
 from .standardize.spaces import action_space_is_discrete
 from .standardize.helpers import (
     EnvWrapper,
+    dones_of,
     env_autoresets,
     env_num_envs,
     env_render,
+    env_takes_torch,
     exists,
     first_existing,
     get_adapter,
     get_attr,
-    is_array,
+    is_array_like,
     is_vectorized,
     normalize_reset_out,
     normalize_step_out,
@@ -90,37 +92,40 @@ def is_distribution(out):
 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()
+    def convert_leaf(val):
+        if is_distribution(val):
+            if not deterministic:
+                return val.sample()
 
-        probs = get_attr(out, 'probs')
+            probs = get_attr(val, 'probs')
 
-        # categorical-style probs carry a trailing action axis to argmax over
+            # 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)
+            if exists(probs) and probs.ndim >= 2 and probs.shape[-1] > 1:
+                return probs.argmax(dim = -1)
 
-        # bernoulli-style probabilities threshold at 0.5
+            # bernoulli-style probabilities threshold at 0.5
 
-        if exists(probs):
-            return (probs > 0.5).long()
+            if exists(probs):
+                return (probs > 0.5).long()
 
-        return out.mean
+            return val.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 is_tensor(val) and val.is_floating_point() and val.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()
+            if exists(num_actions) and val.shape[-1] == num_actions:
+                return val.argmax(dim = -1) if deterministic else Categorical(logits = val).sample()
 
-    return out
+        return val
+
+    return tree_map(convert_leaf, 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__'))):
+        if not is_array_like(x):
             return x
 
         tensor = x if is_tensor(x) else torch.as_tensor(np.asarray(x))
@@ -135,48 +140,55 @@ def to_actor_obs(obs, device, batched):
     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)
+    def to_torch(a):
+        if not is_tensor(a):
+            a = torch.as_tensor(a)
+        return a.to(device) if exists(device) else a
 
-    if not isinstance(env, EnvWrapper):
-        if get_adapter(env).torch_native:
-            return action
+    action = tree_map(to_torch, action)
 
-        action = to_numpy(action)
+    chunk_shape = get_attr(env, 'chunk_action_shape')
 
-        # single raw sims get the leading batch dim stripped, while discrete
-        # actions additionally collapse to scalars
+    if exists(chunk_shape):
+        rank = len(chunk_shape) + 1
+    else:
+        space = first_existing(env, 'single_action_space', 'action_space')
+        shape = get_attr(space, 'shape')
+        rank = len(shape) + 1 if exists(shape) else None
 
-        if not is_vectorized(env):
-            if action.ndim > 1 and action.shape[0] == 1:
-                action = action[0]
+    if exists(rank):
+        def pad_rank(a):
+            while a.ndim < rank:
+                a = a[None]
+            return a
 
-            if action_space_is_discrete(get_adapter(env).action_space):
-                if action.ndim > 0 and action.shape[0] == 1:
-                    action = action[0]
+        action = tree_map(pad_rank, action)
 
-                if action.ndim == 0:
-                    action = action.item()
+    # raw simulators take unbatched numpy — torch-native sims take torch as-is (zero-copy)
 
-        return action
+    if not env_takes_torch(env):
+        action = tree_map(to_numpy, action)
 
-    # wrappers take batched torch actions - pad up to the declared rank
+        # single raw sims get the leading batch dim stripped, while discrete
+        # actions additionally collapse to scalars
 
-    chunk_shape = get_attr(env, 'chunk_action_shape')
+        if not is_vectorized(env) and not exists(chunk_shape):
+            def strip_leaf(a):
+                if a.ndim > 1 and a.shape[0] == 1:
+                    a = a[0]
 
-    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
+                if action_space_is_discrete(get_adapter(env).action_space):
+                    if a.ndim > 0 and a.shape[0] == 1:
+                        a = a[0]
+
+                    if a.ndim == 0:
+                        a = a.item()
+
+                return a
 
-    while action.ndim < rank:
-        action = action[None]
+            action = tree_map(strip_leaf, action)
 
     return action
 
@@ -319,7 +331,7 @@ def evaluate_actor(
             lengths += info.get('chunk_length', 1) if isinstance(info, dict) else 1
             steps += 1
 
-            done = to_numpy(terminated | truncated).astype(bool).reshape(-1)
+            done = to_numpy(dones_of(terminated, truncated)).astype(bool).reshape(-1)
 
             if exists(max_steps):
                 done |= to_numpy(steps >= max_steps).reshape(-1)

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

@@ -6,16 +6,16 @@ from torch import is_tensor
 
 from .standardize.helpers import (
     TransformObservationWrapper,
-    exists,
-    default,
     any_true,
+    cast_tuple,
+    default,
+    exists,
+    is_array_like,
+    is_scalar,
 )
 
 # helper functions
 
-def cast_tuple(val):
-    return val if isinstance(val, (tuple, list)) else (val,)
-
 def calc_lerp_weight(lam, done, x):
     if not exists(done) or not any_true(done):
         return 1. - lam
@@ -23,7 +23,9 @@ def calc_lerp_weight(lam, done, x):
     # reset trace on done, else decay
 
     if not is_tensor(done):
-        done = torch.as_tensor(done, device = x.device)
+        done = torch.as_tensor(done, device = x.device, dtype = torch.bool)
+    else:
+        done = done.to(device = x.device, dtype = torch.bool)
 
     weight = torch.where(done, 1., 1. - lam).to(x)
 
@@ -36,6 +38,7 @@ def calc_lerp_weight(lam, done, x):
 # https://arxiv.org/abs/2503.15200
 
 class MemoryTraceWrapper(TransformObservationWrapper):
+    priority = 80
 
     def __init__(
         self,
@@ -61,17 +64,26 @@ class MemoryTraceWrapper(TransformObservationWrapper):
         return f'{prefix}_{lam}' if len(self.lambdas) > 1 else prefix
 
     def transform_obs(self, obs, done = None):
-        out = dict(obs) if isinstance(obs, dict) else {self.obs_key: obs}
+        if isinstance(obs, (tuple, list)):
+            out = {f'{self.obs_key}_{i}': elem for i, elem in enumerate(obs)}
+        elif isinstance(obs, dict):
+            out = dict(obs)
+        else:
+            out = {self.obs_key: obs}
+
         target_keys = tuple(k for k in default(self.keys, tuple(out.keys())) if k in out)
 
         for key in target_keys:
             val = out.get(key)
 
-            if not exists(val):
+            if not exists(val) or not (is_array_like(val) or is_scalar(val)):
                 continue
 
             if not is_tensor(val):
-                val = torch.as_tensor(val)
+                try:
+                    val = torch.as_tensor(val)
+                except (TypeError, ValueError):
+                    continue
                 out[key] = val
 
             val_float = val.float() if not val.is_floating_point() else val

{env_ssl_wrapper-0.5.2 → env_ssl_wrapper-0.5.4}/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.4}/env_ssl_wrapper/standardize/action_transform_wrapper.py

@@ -57,6 +57,8 @@ def action_bounds(env):
 # wrapper
 
 class ActionTransformWrapper(EnvWrapper):
+    priority = 30
+
     def __init__(
         self,
         env,
@@ -84,7 +86,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.4}/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.4}/env_ssl_wrapper/standardize/auto_batched_wrapper.py

@@ -4,7 +4,7 @@ from functools import partial
 
 import numpy as np
 import torch
-from torch.utils._pytree import tree_map
+from torch.utils._pytree import tree_flatten, tree_map, tree_structure, tree_unflatten
 
 from .helpers import (
     FINAL_OBSERVATION_KEYS,
@@ -35,16 +35,11 @@ def to_numeric_array(t):
     return arr if arr.dtype.kind in 'biufc' else t
 
 def maybe_expand_dim(x):
-    # add leading batch dim to every numeric leaf
-
     def _expand(t):
         arr = to_numeric_array(t)
         return arr[None] if is_array(arr) else t
 
-    if is_array(x):
-        return x[None]
-
-    return tree_map(_expand, x) if isinstance(x, (dict, list, tuple)) else _expand(x)
+    return tree_map(_expand, x)
 
 def is_integer_dtype(t):
     if is_tensor(t):
@@ -65,17 +60,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')
@@ -121,20 +117,12 @@ def heuristic_leaf(t, is_vector = False):
     return arr.item() if arr.ndim == 0 else arr
 
 def rebuild_container(x, leaves):
-
-    if isinstance(x, list):
-        return leaves
-
-    if exists(get_attr(type(x), '_fields')):
-        return type(x)(*leaves)
-
-    return type(x)(leaves)
+    return tree_unflatten(leaves, tree_structure(x))
 
 def is_numeric_container(x):
-    # purely numeric nested sequences count as one leaf when no space declares structure
-
     if isinstance(x, (list, tuple)):
-        return len(x) > 0 and all(is_numeric_container(item) for item in x)
+        leaves, _ = tree_flatten(x)
+        return len(leaves) > 0 and all(map(is_scalar, leaves))
 
     return is_scalar(x)
 
@@ -167,6 +155,7 @@ def maybe_squeeze_dim(x, shape_tree = None, is_vector = False, prepend_batch = F
 class AutoBatchedWrapper(EnvWrapper):
 
     is_auto_batched = True
+    priority = 50
 
     def __init__(self, env, is_vector: bool | None = None):
         super().__init__(env)
@@ -186,7 +175,16 @@ class AutoBatchedWrapper(EnvWrapper):
         if not exists(self.action_shape_tree):
             self.refresh_action_space()
 
-        return (maybe_expand_dim(obs), info) if not self.is_vector else (obs, info)
+        if self.is_vector:
+            return obs, info
+
+        obs = maybe_expand_dim(obs)
+        if isinstance(info, dict):
+            for key in FINAL_OBSERVATION_KEYS:
+                if key in info:
+                    info[key] = maybe_expand_dim(info[key])
+
+        return obs, info
 
     def step(self, action):
         action = maybe_squeeze_dim(action, shape_tree = self.action_shape_tree, is_vector = self.is_vector, prepend_batch = self.prepend_batch)

{env_ssl_wrapper-0.5.2 → env_ssl_wrapper-0.5.4}/env_ssl_wrapper/standardize/done_tracker_wrapper.py

@@ -17,6 +17,8 @@ from .helpers import (
 # classes
 
 class DoneTrackerWrapper(EnvWrapper):
+    priority = 70
+
     def __init__(self, env):
         if not is_vectorized(env):
             env = AutoBatchedWrapper(env)

{env_ssl_wrapper-0.5.2 → env_ssl_wrapper-0.5.4}/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,12 +60,14 @@ 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
 
 class EpisodePaddingWrapper(EnvWrapper):
+    priority = 10
+
     def __init__(self, env, pad_autoreset: bool = True):
         super().__init__(env)
         self.is_vector = is_vectorized(env)
@@ -90,7 +93,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 +109,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.4}/env_ssl_wrapper/standardize/flatten_obs_wrapper.py

@@ -37,7 +37,9 @@ def concat_leaves(leaves, is_vector = False):
     axis = -1 if is_vector else 0
     leaves = [flatten_leaf(t, is_vector = is_vector) for t in leaves]
 
-    if is_tensor(leaves[0]):
+    if any(map(is_tensor, leaves)):
+        device = next(t.device for t in leaves if is_tensor(t))
+        leaves = [torch.as_tensor(t, device = device) if not is_tensor(t) else t.to(device) for t in leaves]
         return torch.cat(leaves, dim = axis)
 
     return np.concatenate(leaves, axis = axis)
@@ -45,6 +47,8 @@ def concat_leaves(leaves, is_vector = False):
 # class
 
 class FlattenObsWrapper(TransformObservationWrapper):
+    priority = 90
+
     def __init__(self, env, is_vector: bool | None = None):
         super().__init__(env)
         self.is_vector = default(is_vector, is_vectorized(env))