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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: env-ssl-wrapper
-Version: 0.5.3
+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.3 → 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.3 → env_ssl_wrapper-0.5.4}/env_ssl_wrapper/evaluate.py

@@ -12,9 +12,11 @@ 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,
@@ -90,31 +92,34 @@ 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
 
@@ -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.3 → 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.3 → 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,

{env_ssl_wrapper-0.5.3 → 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):
@@ -122,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)
 
@@ -168,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)
@@ -187,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.3 → 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.3 → env_ssl_wrapper-0.5.4}/env_ssl_wrapper/standardize/episode_padding_wrapper.py

@@ -66,6 +66,8 @@ def merge_final(current, value, mask):
 # class
 
 class EpisodePaddingWrapper(EnvWrapper):
+    priority = 10
+
     def __init__(self, env, pad_autoreset: bool = True):
         super().__init__(env)
         self.is_vector = is_vectorized(env)

{env_ssl_wrapper-0.5.3 → 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))

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

@@ -12,7 +12,14 @@ def exists(v):
     return v is not None
 
 def default(v, d):
-    return v if exists(v) else d
+    return v if exists(v) else (d() if callable(d) else d)
+
+def cast_tuple(val, length = 1):
+    if isinstance(val, tuple):
+        return val
+    if isinstance(val, list):
+        return tuple(val)
+    return (val,) * length
 
 def get_attr(obj, name, default = None):
     # properties that raise count as missing
@@ -84,9 +91,8 @@ def to_numpy(t):
     return t.detach().cpu().numpy() if is_tensor(t) else np.asarray(t)
 
 def any_true(x):
-    if is_tensor(x):
-        return bool(x.any())
-    return bool(np.asarray(x).any())
+    leaves, _ = tree_flatten(x)
+    return any(bool(leaf.any()) if is_tensor(leaf) else bool(np.asarray(leaf).any()) for leaf in leaves)
 
 def copy_leaf(x):
     if is_tensor(x):
@@ -98,8 +104,6 @@ def copy_leaf(x):
     return x
 
 def dones_of(terminated, truncated):
-    if not isinstance(terminated, (dict, list, tuple)):
-        return terminated | truncated
     return tree_map(lambda a, b: a | b, terminated, truncated)
 
 # sim step / reset normalization
@@ -107,13 +111,16 @@ def dones_of(terminated, truncated):
 def is_time_step(out):
     return exists(get_attr(out, 'step_type')) and exists(get_attr(out, 'observation'))
 
-def zero_like(x):
+def _zero_bool_leaf(x):
     if is_tensor(x):
         return torch.zeros_like(x, dtype = torch.bool)
 
     arr = np.asarray(x)
     return np.zeros_like(arr, dtype = bool) if arr.ndim > 0 else False
 
+def zero_like(x):
+    return tree_map(_zero_bool_leaf, x)
+
 def normalize_obs(obs):
     # raw arrays / scalars / text / containers pass through — custom observation
     # objects (rlbench, ...) flatten via get_low_dim_data or unpack their
@@ -148,16 +155,8 @@ def normalize_obs(obs):
     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)
+    leaves, _ = tree_flatten(x)
+    return any(map(is_array_like, leaves))
 
 def is_instruction(x):
     if isinstance(x, str):
@@ -235,15 +234,9 @@ def stack_trees(trees):
     if isinstance(first, np.ndarray):
         return np.stack(trees)
 
-    if isinstance(first, dict):
-        return {key: stack_trees([t[key] for t in trees]) for key in first}
-
-    if isinstance(first, tuple):
-        return tuple(stack_trees([t[i] for t in trees]) for i in range(len(first)))
-
     leaves = [tree_flatten(tree)[0] for tree in trees]
     stacked = [_stack_leaves(col) for col in zip(*leaves)]
-    return tree_unflatten(stacked, tree_structure(trees[0]))
+    return tree_unflatten(stacked, tree_structure(first))
 
 def unpack_vector_observations(arr):
     # unpacks a 1D sequence / numpy object array of unbatched single-env observations
@@ -283,6 +276,18 @@ def env_render(env, height, width, camera = None):
 def is_vectorized(env) -> bool:
     return get_adapter(env).is_vectorized
 
+def env_takes_torch(env) -> bool:
+    if get_adapter(env).torch_native:
+        return True
+
+    current = env
+    while isinstance(current, EnvWrapper):
+        if type(current).__name__ == 'TensorWrapper' and getattr(current, 'convert_in', True):
+            return True
+        current = getattr(current, 'env', None)
+
+    return False
+
 # gymnasium 1.x surfaces final observations as 'final_obs' / '_final_obs', everything here uses 'final_observation' / '_final_observation'
 
 FINAL_OBSERVATION_KEYS = ('final_observation', 'final_obs')
@@ -373,6 +378,8 @@ def safe_close(env):
 # base wrapper
 
 class EnvWrapper:
+    priority = 50
+
     def __init__(self, env):
         self.env = env
 

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

@@ -9,6 +9,7 @@ from einops import rearrange
 from .helpers import (
     FINAL_OBSERVATION_KEYS,
     TransformObservationWrapper,
+    cast_tuple,
     env_render,
     env_render_mode,
     exists,
@@ -18,9 +19,6 @@ from .helpers import (
 
 # helper functions
 
-def cast_tuple(t, length = 1):
-    return t if isinstance(t, tuple) else ((t,) * length)
-
 def render_frame(env, image_size = (64, 64), camera = None):
     # render from any sim — the shared env_render probe knows each image
     # surface (dm_control physics, pybullet client, robosuite sim); envs
@@ -72,6 +70,8 @@ def process_image(
 # class
 
 class ImageObservationWrapper(TransformObservationWrapper):
+    priority = 20
+
     def __init__(
         self,
         env,

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

@@ -71,7 +71,7 @@ def space_dim(space) -> int | None:
     if exists(shape):
         return shape_dim(shape)
 
-    subspaces = get_attr(space, 'spaces')
+    subspaces = space if isinstance(space, (dict, list, tuple)) else get_attr(space, 'spaces')
     if not exists(subspaces):
         return None
 

{env_ssl_wrapper-0.5.3 → env_ssl_wrapper-0.5.4}/env_ssl_wrapper/standardize/standardize_wrapper.py

@@ -14,6 +14,8 @@ from .spaces import infer_observation_space, space_from_action_spec
 # class
 
 class StandardizeWrapper(EnvWrapper):
+    priority = 0
+
     # normalizes any sim into (obs, reward, terminated, truncated, info) via its adapter
 
     def __init__(self, env, adapter = None):

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

@@ -43,9 +43,6 @@ def numpy_to_torch(x, device, cast_obs_to_float = True):
     if isinstance(x, np.ndarray) and x.dtype == object:
         x = unpack_vector_observations(x)
 
-    if not isinstance(x, (dict, list, tuple)):
-        return to_torch_leaf(x, device, cast_obs_to_float)
-
     return tree_map(partial(to_torch_leaf, device = device, cast_obs_to_float = cast_obs_to_float), x)
 
 def to_numpy_leaf(t):
@@ -66,9 +63,6 @@ def to_numpy_leaf(t):
 
 def torch_to_numpy(x):
     # torch to numpy; 0-dim collapses to scalar, float64 → float32
-    if not isinstance(x, (dict, list, tuple)):
-        return to_numpy_leaf(x)
-
     return tree_map(to_numpy_leaf, x)
 
 # rewards float32, dones bool
@@ -85,6 +79,8 @@ def contract(t, to_float = False):
 # class
 
 class TensorWrapper(EnvWrapper):
+    priority = 60
+
     def __init__(
         self,
         env,
@@ -138,17 +134,15 @@ class TensorWrapper(EnvWrapper):
             leaf = to_torch_leaf(leaf, self.device, cast_obs_to_float = False)
             return contract(leaf, to_float = to_float)
 
-        if not isinstance(t, (dict, list, tuple)):
-            return convert(t)
-
         return tree_map(convert, t)
 
     def step(self, action):
         if self.convert_in:
             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
+            elif is_tensor(action):
+                action = action.to(self.device)
+            else:
                 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.3 → env_ssl_wrapper-0.5.4}/env_ssl_wrapper/standardize/time_limit_wrapper.py

@@ -33,6 +33,8 @@ def back_to_like(t, numpy_arr):
     return numpy_arr
 
 class TimeLimitWrapper(EnvWrapper):
+    priority = 40
+
     """
     caps episode length at max_timesteps, setting truncated = True for
     capped envs (vectorized and single alike); timers reset per episode.

{env_ssl_wrapper-0.5.3 → env_ssl_wrapper-0.5.4}/env_ssl_wrapper/standardize/utils.py

@@ -60,48 +60,27 @@ def parse_wrapper(wrapper):
 def compose_env(env, *wrappers, pad_episodes: bool = True):
     env = instantiate_env(env)
 
-    funcs = []
-    classes = []
-
-    for wrapper in wrappers:
-        func, cls = parse_wrapper(wrapper)
-        funcs.append(func)
-        classes.append(cls)
+    items = [parse_wrapper(wrapper) for wrapper in wrappers]
+    classes = {cls for _, cls in items}
 
     if StandardizeWrapper not in classes:
-        funcs.insert(0, StandardizeWrapper)
-        classes.insert(0, StandardizeWrapper)
+        items.append((StandardizeWrapper, StandardizeWrapper))
 
     # vectorized envs get standardized episode padding + a persistent final_observation
     # (autoresetting envs re-emit the true terminal obs, so no padding for those)
 
     if pad_episodes and EpisodePaddingWrapper not in classes and is_vectorized(env):
         pad_wrapper = partial(EpisodePaddingWrapper, pad_autoreset = not env_autoresets(env))
-        funcs.insert(1, pad_wrapper)
-        classes.insert(1, EpisodePaddingWrapper)
-
-    from ..memory_trace import MemoryTraceWrapper
-    if MemoryTraceWrapper in classes and TensorWrapper in classes:
-        idx_mem = classes.index(MemoryTraceWrapper)
-        idx_ten = classes.index(TensorWrapper)
-        if idx_mem < idx_ten:
-            f = funcs.pop(idx_mem)
-            c = classes.pop(idx_mem)
-            idx_ten = classes.index(TensorWrapper)
-            funcs.insert(idx_ten + 1, f)
-            classes.insert(idx_ten + 1, c)
-
-    # action chunking changes the step signature, so it always goes outermost
-
-    from ..action_chunk import ActionChunkWrapper
-    if ActionChunkWrapper in classes:
-        idx_chunk = classes.index(ActionChunkWrapper)
-        funcs.append(funcs.pop(idx_chunk))
-        classes.append(classes.pop(idx_chunk))
+        items.append((pad_wrapper, EpisodePaddingWrapper))
+
+    # sort by canonical pipeline priority (stable sort preserves user order for equal priorities)
+
+    items.sort(key = lambda item: getattr(item[1], 'priority', 50))
 
+    classes = [cls for _, cls in items]
     assert len(set(classes)) == len(classes), 'duplicate wrappers found'
 
-    for func in funcs:
+    for func, _ in items:
         env = func(env)
 
     return env

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

@@ -3,20 +3,22 @@ from __future__ import annotations
 from multiprocessing import get_context
 import numpy as np
 from torch import is_tensor
-from torch.utils._pytree import tree_map
+from torch.utils._pytree import tree_flatten, tree_map, tree_unflatten
 
 from .helpers import (
     any_true,
     dones_of,
     exists,
+    get_adapter,
     get_attr,
     instantiate_env,
     safe_close,
     stack_trees,
+    to_numpy,
     truthy_attr,
     _zero_leaf,
 )
-from .spaces import action_dim_of, obs_dim_of
+from .spaces import action_dim_of, action_space_is_discrete, obs_dim_of
 from .standardize_wrapper import StandardizeWrapper
 
 # leaf helpers
@@ -48,6 +50,15 @@ class WorkerError(RuntimeError):
 
 def _exec(env, cmd, payload):
     if cmd == 'step':
+        if is_tensor(payload) and not get_adapter(env).torch_native:
+            payload = to_numpy(payload)
+
+        if action_space_is_discrete(env.action_space):
+            if hasattr(payload, 'item') and (not hasattr(payload, 'ndim') or payload.ndim <= 1):
+                payload = payload.item()
+            elif isinstance(payload, (list, tuple)) and len(payload) == 1:
+                payload = payload[0]
+
         obs, reward, terminated, truncated, info = env.step(payload)
 
         if any_true(dones_of(terminated, truncated)):
@@ -182,16 +193,14 @@ def _shutdown(conns, procs):
             pass
 
 def _split_actions(actions, num_envs):
-    if isinstance(actions, dict):
-        assert all(len(v) == num_envs for v in actions.values()), f'expected {num_envs} actions per key'
-        return [{k: v[i] for k, v in actions.items()} for i in range(num_envs)]
-
-    if isinstance(actions, tuple) and not is_tensor(actions) and not isinstance(actions, np.ndarray):
-        assert all(len(v) == num_envs for v in actions), f'expected {num_envs} actions per element'
-        return [tuple(elem[i] for elem in actions) for i in range(num_envs)]
+    if isinstance(actions, list) and len(actions) == num_envs:
+        return actions
 
-    assert len(actions) == num_envs, f'expected {num_envs} actions, but got {len(actions)}'
-    return actions
+    leaves, spec = tree_flatten(actions)
+    assert all(len(leaf) == num_envs for leaf in leaves), (
+        f'expected {num_envs} actions, but got {len(leaves[0]) if leaves else 0}'
+    )
+    return [tree_unflatten([leaf[i] for leaf in leaves], spec) for i in range(num_envs)]
 
 # class
 
@@ -270,16 +279,19 @@ class MultiprocessingVecEnv:
         _recv_all(self._conns)
 
     def reset(self, seed = None, **kwargs):
-        if exists(seed):
-            self.seed(seed)
-
-        for conn in self._conns:
-            _safe_send(conn, ('reset', kwargs))
+        for i, conn in enumerate(self._conns):
+            worker_kwargs = dict(kwargs)
+            if exists(seed):
+                worker_kwargs['seed'] = seed + i
+            _safe_send(conn, ('reset', worker_kwargs))
 
         results = _recv_all(self._conns)
         return stack_trees([obs for obs, _ in results]), _merge_infos([info for _, info in results])
 
     def step(self, actions):
+        if is_tensor(actions):
+            actions = to_numpy(actions)
+
         actions = _split_actions(actions, self.num_envs)
 
         for conn, action in zip(self._conns, actions):

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

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