env-ssl-wrapper 0.4.1__tar.gz → 0.4.2__tar.gz

{env_ssl_wrapper-0.4.1 → env_ssl_wrapper-0.4.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: env-ssl-wrapper
-Version: 0.4.1
+Version: 0.4.2
 Summary: One torch-native interface for any MDP environment
 Project-URL: Homepage, https://pypi.org/project/env-ssl-wrapper/
 Project-URL: Repository, https://codeberg.org/lucidrains/env-ssl-wrapper
@@ -47,14 +47,17 @@ Requires-Dist: fire; extra == 'examples'
 Provides-Extra: test
 Requires-Dist: box2d-py>=2.3.5; extra == 'test'
 Requires-Dist: dm-control; extra == 'test'
+Requires-Dist: gym-pusht; extra == 'test'
 Requires-Dist: gymnasium-robotics; extra == 'test'
 Requires-Dist: gymnasium>=0.29.1; extra == 'test'
 Requires-Dist: imageio; extra == 'test'
+Requires-Dist: mean-conc-beta; extra == 'test'
 Requires-Dist: mujoco; extra == 'test'
 Requires-Dist: pillow; extra == 'test'
 Requires-Dist: populora; extra == 'test'
 Requires-Dist: pybullet; extra == 'test'
 Requires-Dist: pygame; extra == 'test'
+Requires-Dist: pymunk<7; extra == 'test'
 Requires-Dist: pytest; extra == 'test'
 Description-Content-Type: text/markdown
 
@@ -81,7 +84,7 @@ env = StandardizeEnvWrapper(any_env)
 obs, info = env.reset() # torch.float32, batched
 
 while not env.all_done:
-    actions = torch.randint(0, 2, (8,))
+    actions = torch.randint(0, 2, (1,))
     obs, reward, terminated, truncated, info = env.step(actions)
 ```
 

{env_ssl_wrapper-0.4.1 → env_ssl_wrapper-0.4.2}/README.md

@@ -21,7 +21,7 @@ env = StandardizeEnvWrapper(any_env)
 obs, info = env.reset() # torch.float32, batched
 
 while not env.all_done:
-    actions = torch.randint(0, 2, (8,))
+    actions = torch.randint(0, 2, (1,))
     obs, reward, terminated, truncated, info = env.step(actions)
 ```
 

{env_ssl_wrapper-0.4.1 → env_ssl_wrapper-0.4.2}/env_ssl_wrapper/standardize/action_transform_wrapper.py

@@ -97,9 +97,14 @@ class ActionTransformWrapper(EnvWrapper):
         if was_scalar:
             t = rearrange(t, '-> 1')
 
-        dim = t.shape[-1]
-        low = np.broadcast_to(low, (dim,))
-        high = np.broadcast_to(high, (dim,))
+        # bounds may be unbatched (single_action_space) or already batched
+        # (vector envs exposing only a batched action_space)
+
+        if low.shape != t.shape and not (low.ndim > 0 and t.shape[-low.ndim:] == low.shape):
+            dim = t.shape[-1]
+            low = np.broadcast_to(low, (dim,))
+            high = np.broadcast_to(high, (dim,))
+
         valid = np.isfinite(low) & np.isfinite(high)
 
         if is_tensor(t):

{env_ssl_wrapper-0.4.1 → env_ssl_wrapper-0.4.2}/env_ssl_wrapper/standardize/adapters.py

@@ -1,32 +1,20 @@
 from __future__ import annotations
 
-import numpy as np
-import torch
-from torch import is_tensor
-
 from .helpers import (
     EnvWrapper,
     default,
     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
 
-# helpers
-
-def zero_like(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 is_time_step(out):
-    return exists(get_attr(out, 'step_type')) and exists(get_attr(out, 'observation'))
-
 # base adapter
 
 class BaseEnvAdapter:
@@ -38,27 +26,10 @@ class BaseEnvAdapter:
         self.env = env
 
     def step(self, action) -> tuple:
-        out = self.env.step(action)
-        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)
-        if len(out) == 5:
-            obs, reward, term, trunc, info = out
-            return obs, reward, term, trunc, 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
-        raise ValueError(f'cannot standardize step output of length {len(out)}')
+        return normalize_step_out(self.env.step(action))
 
     def reset(self, **kwargs) -> tuple:
-        out = self.env.reset(**kwargs)
-        if is_time_step(out):
-            return out.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 {})
-        return out, {}
+        return normalize_reset_out(self.env.reset(**kwargs))
 
     def seed(self, seed: int):
         if callable(get_attr(self.env, 'seed')):
@@ -381,15 +352,16 @@ class GymnasiumAdapter(BaseEnvAdapter):
 
     @property
     def autoresets(self) -> bool:
-        if truthy_attr(first_existing(self.env, 'autoreset', 'autoresets', 'autoreset_mode')):
-            return True
+        mode = first_existing(self.env, 'autoreset', 'autoresets', 'autoreset_mode')
+
+        if exists(mode):
+            return truthy_attr(mode)
+
         try:
             from gymnasium.vector import VectorEnv
-            if isinstance(self.env, VectorEnv):
-                return True
+            return isinstance(self.env, VectorEnv)
         except ImportError:
-            pass
-        return False
+            return False
 
     def seed(self, seed: int):
         try:

{env_ssl_wrapper-0.4.1 → env_ssl_wrapper-0.4.2}/env_ssl_wrapper/standardize/done_tracker_wrapper.py

@@ -2,8 +2,6 @@ from __future__ import annotations
 
 import numpy as np
 
-from torch.utils._pytree import tree_flatten
-
 from .auto_batched_wrapper import AutoBatchedWrapper
 from .helpers import (
     EnvWrapper,
@@ -11,27 +9,11 @@ from .helpers import (
     env_autoresets,
     env_num_envs,
     exists,
-    get_attr,
-    is_array,
+    get_batch_size,
     is_vectorized,
     to_numpy,
 )
 
-# helper functions
-
-def get_batch_size(tree) -> int | None:
-    leaves, _ = tree_flatten(tree)
-
-    if not leaves:
-        return None
-
-    first = leaves[0]
-
-    if is_array(first):
-        return len(first) if first.ndim > 0 else None
-
-    return len(first) if exists(get_attr(first, '__len__')) else None
-
 # classes
 
 class DoneTrackerWrapper(EnvWrapper):
@@ -113,6 +95,6 @@ class DoneTrackerWrapper(EnvWrapper):
             info['episode_lengths'] = self.episode_lengths.copy()
 
             if self.all_done:
-                info.update(needs_reset = True, all_done = True)
+                info.update(needs_reset = self.needs_reset, all_done = True)
 
         return obs, reward, terminated, truncated, info

{env_ssl_wrapper-0.4.1 → env_ssl_wrapper-0.4.2}/env_ssl_wrapper/standardize/episode_padding_wrapper.py

@@ -21,15 +21,18 @@ from .helpers import (
 
 # helpers
 
-def zero_mask(x, mask, fill_scalar = None):
-    if is_tensor(x):
-        m = torch.as_tensor(mask, device = x.device, dtype = torch.bool)
-        diff = x.ndim - m.ndim
+def broadcast_mask(mask, x):
+    if not is_tensor(x):
+        return mask
+
+    m = torch.as_tensor(mask, device = x.device, dtype = torch.bool)
+    diff = x.ndim - m.ndim
 
-        if diff > 0:
-            m = rearrange(m, f'... -> ... {" ".join(["1"] * diff)}')
+    return rearrange(m, f'... -> ... {" ".join(["1"] * diff)}') if diff > 0 else m
 
-        return torch.where(m, torch.zeros_like(x), x)
+def zero_mask(x, mask, fill_scalar = None):
+    if is_tensor(x):
+        return torch.where(broadcast_mask(mask, x), torch.zeros_like(x), x)
 
     arr = np.asarray(x)
 
@@ -49,13 +52,7 @@ def back_to_mask_type(dones, newly):
 
 def merge_final(current, value, mask):
     if is_tensor(current):
-        m = torch.as_tensor(mask, device = current.device, dtype = torch.bool)
-        diff = current.ndim - m.ndim
-
-        if diff > 0:
-            m = rearrange(m, f'... -> ... {" ".join(["1"] * diff)}')
-
-        return torch.where(m, value, current)
+        return torch.where(broadcast_mask(mask, current), value, current)
 
     curr = np.asarray(current)
 
@@ -97,15 +94,15 @@ class EpisodePaddingWrapper(EnvWrapper):
             if self._is_done is None or len(self._is_done) != len(mask):
                 self._is_done = np.zeros(len(mask), dtype = bool)
 
-            if self.autoreset:
-                self._is_done &= mask
+            # autoreset envs revive every done slot, so each done is a new
+            # terminal transition; non-autoreset slots stay done until reset
+
+            newly = mask if self.autoreset else mask & ~self._is_done
+            self._is_done |= mask
 
             if mask.any():
                 assert exists(self._last_obs), 'environment needs reset before calling step. call env.reset() first'
 
-                newly = mask & ~self._is_done
-                self._is_done |= mask
-
                 if newly.any():
                     final_val = first_existing(info, 'final_observation', 'final_obs')
                     value = final_val if exists(final_val) else self._last_obs

{env_ssl_wrapper-0.4.1 → env_ssl_wrapper-0.4.2}/env_ssl_wrapper/standardize/helpers.py

@@ -22,10 +22,15 @@ def get_attr(obj, name, default = None):
         return default
 
 def truthy_attr(value):
-    # flags arrive as None, methods, numpy scalars — only honest truths count
+    # flags arrive as None, methods, numpy scalars — only honest truths count.
+    # gymnasium's AutoresetMode.DISABLED is a truthy enum that means "off"
+
     if not exists(value) or callable(value):
         return False
 
+    if getattr(value, 'name', None) == 'DISABLED':
+        return False
+
     try:
         return bool(value)
     except Exception:
@@ -46,6 +51,19 @@ def first_existing(obj, *names):
 def is_scalar(v):
     return isinstance(v, (int, float, bool, np.number, np.bool_))
 
+def get_batch_size(tree) -> int | None:
+    leaves, _ = tree_flatten(tree)
+
+    if not leaves:
+        return None
+
+    first = leaves[0]
+
+    if is_array(first):
+        return len(first) if first.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)
 
@@ -71,6 +89,43 @@ def dones_of(terminated, truncated):
         return terminated | truncated
     return tree_map(lambda a, b: a | b, terminated, truncated)
 
+# sim step / reset normalization
+
+def is_time_step(out):
+    return exists(get_attr(out, 'step_type')) and exists(get_attr(out, 'observation'))
+
+def zero_like(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 normalize_reset_out(out):
+    if is_time_step(out):
+        return out.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 {})
+
+    return 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)
+
+    if len(out) == 5:
+        return out
+
+    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
+
+    raise ValueError(f'could not standardize step output of length {len(out)}')
+
 def _zero_leaf(x):
     if is_tensor(x):
         return torch.zeros_like(x)
@@ -272,16 +327,11 @@ class TransformObservationWrapper(EnvWrapper):
         self.takes_done = accepts_done_param(self.transform_obs)
 
     def transform_obs(self, obs, done = None):
-        if hasattr(self, 'observation'):
-            return self.observation(obs)
         return obs
 
     def transform(self, obs, done = None):
         return self.transform_obs(obs, done = done) if self.takes_done else self.transform_obs(obs)
 
-    def observation(self, obs):
-        return self.transform_obs(obs)
-
     def reset(self, **kwargs):
         obs, info = self.env.reset(**kwargs)
         obs = self.transform_obs(obs)

{env_ssl_wrapper-0.4.1 → env_ssl_wrapper-0.4.2}/env_ssl_wrapper/standardize/image_wrapper.py

@@ -6,8 +6,15 @@ import numpy as np
 
 from einops import rearrange
 
-from .helpers import EnvWrapper, env_render, env_render_mode, exists
-from .standardize_wrapper import normalize_reset_out, normalize_step_out
+from .helpers import (
+    FINAL_OBSERVATION_KEYS,
+    TransformObservationWrapper,
+    env_render,
+    env_render_mode,
+    exists,
+    normalize_reset_out,
+    normalize_step_out,
+)
 
 # helper functions
 
@@ -32,12 +39,19 @@ def render_frame(env, image_size = (64, 64), camera = None):
 
         img = env.render()
 
+    if isinstance(img, (list, tuple)):
+        img = np.ascontiguousarray(img)
+
     if isinstance(img, torch.Tensor):
         img = img.detach().cpu().numpy()
 
     img = np.ascontiguousarray(img)
     img = torch.from_numpy(img)
-    return rearrange(img, 'h w c -> 1 c h w')
+
+    if img.ndim == 4:
+        return rearrange(img, 'b h w c -> b c h w'), True
+
+    return rearrange(img, 'h w c -> 1 c h w'), False
 
 def process_image(
     img,
@@ -61,7 +75,7 @@ def process_image(
 
 # class
 
-class ImageObservationWrapper(EnvWrapper):
+class ImageObservationWrapper(TransformObservationWrapper):
     def __init__(
         self,
         env,
@@ -81,18 +95,21 @@ class ImageObservationWrapper(EnvWrapper):
         self.normalize_divisor = normalize_divisor
 
     def render_frame(self):
-        img = render_frame(self.env, image_size = self.image_size, camera = self.camera)
-        return process_image(
+        img, is_batched = render_frame(self.env, image_size = self.image_size, camera = self.camera)
+        processed = process_image(
             img,
             image_size = self.image_size,
             mode = self.mode,
             normalize = self.normalize,
             normalize_divisor = self.normalize_divisor
         )
+        return processed, is_batched
 
     def observation(self, obs):
-        img = self.render_frame()
-        img = rearrange(img, '1 c h w -> c h w')
+        img, is_batched = self.render_frame()
+
+        if not is_batched:
+            img = rearrange(img, '1 c h w -> c h w')
 
         if not isinstance(obs, dict):
             return dict(state = obs, **{self.image_key: img})
@@ -102,10 +119,23 @@ class ImageObservationWrapper(EnvWrapper):
 
         return {**obs, self.image_key: img}
 
+    def transform_obs(self, obs, done = None):
+        return self.observation(obs)
+
     def reset(self, **kwargs):
         obs, info = normalize_reset_out(self.env.reset(**kwargs))
         return self.observation(obs), info
 
     def step(self, action):
         obs, reward, terminated, truncated, info = normalize_step_out(self.env.step(action))
-        return self.observation(obs), reward, terminated, truncated, info
+        out = self.observation(obs)
+
+        # terminal obs are augmented in place, preserving whatever the sim
+        # or the padding wrapper froze as the true terminal observation
+
+        if isinstance(info, dict) and not self.autoresets:
+            for key in FINAL_OBSERVATION_KEYS:
+                if key in info and not (isinstance(info[key], dict) and self.image_key in info[key]):
+                    info[key] = self.observation(info[key])
+
+        return out, reward, terminated, truncated, info

{env_ssl_wrapper-0.4.1 → env_ssl_wrapper-0.4.2}/env_ssl_wrapper/standardize/standardize_env_wrapper.py

@@ -18,11 +18,11 @@ class StandardizeEnvWrapper(EnvWrapper):
         pad_episodes: bool = True,
         done_tracker: bool = True,
         flatten_obs: bool = False,
-        action_transform: bool = False,
+        action_transform: bool | dict = False,
         max_timesteps: int | None = None,
-        image_size: tuple[int, int] | None = None,
+        image_size: int | tuple[int, int] | None = None,
         lambdas: tuple[float, ...] | list[float] | None = None,
-        **kwargs
+        keys: str | tuple[str, ...] | None = None,
     ):
         wrappers = []
 
@@ -30,7 +30,8 @@ class StandardizeEnvWrapper(EnvWrapper):
             wrappers.append(('image', dict(image_size = image_size)))
 
         if action_transform:
-            wrappers.append('action_transform')
+            transform_kwargs = dict(auto = True) if action_transform is True else action_transform
+            wrappers.append(('action_transform', transform_kwargs))
 
         if exists(max_timesteps):
             wrappers.append(('time_limit', dict(max_timesteps = max_timesteps)))
@@ -45,7 +46,7 @@ class StandardizeEnvWrapper(EnvWrapper):
             wrappers.append('done_tracker')
 
         if exists(lambdas):
-            wrappers.append(('memory_trace', dict(lambdas = lambdas, **kwargs)))
+            wrappers.append(('memory_trace', dict(lambdas = lambdas, keys = keys)))
 
         if flatten_obs:
             wrappers.append('flatten_obs')

{env_ssl_wrapper-0.4.1 → env_ssl_wrapper-0.4.2}/env_ssl_wrapper/standardize/standardize_wrapper.py

@@ -1,48 +1,16 @@
 from __future__ import annotations
 
-import numpy as np
-import torch
-from torch import is_tensor
-
-from .adapters import get_adapter, is_time_step, zero_like
+from .adapters import get_adapter
 from .helpers import (
     EnvWrapper,
     default,
     dones_of,
     exists,
-    get_attr,
     instantiate_env,
     mark_terminal_obs,
 )
 from .spaces import infer_observation_space, space_from_action_spec
 
-# helpers
-
-def normalize_reset_out(out):
-    if is_time_step(out):
-        return out.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 {})
-
-    return 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)
-
-    if len(out) == 5:
-        return out
-
-    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
-
-    raise ValueError(f'could not standardize step output of length {len(out)}')
-
 # class
 
 class StandardizeWrapper(EnvWrapper):

{env_ssl_wrapper-0.4.1 → env_ssl_wrapper-0.4.2}/env_ssl_wrapper/standardize/time_limit_wrapper.py

@@ -4,8 +4,18 @@ import numpy as np
 import torch
 from torch import is_tensor
 
-from .helpers import EnvWrapper, dones_of, env_num_envs, is_vectorized, mark_terminal_obs, to_numpy
-from .standardize_wrapper import normalize_reset_out, normalize_step_out
+from .helpers import (
+    EnvWrapper,
+    dones_of,
+    env_num_envs,
+    exists,
+    get_batch_size,
+    is_vectorized,
+    mark_terminal_obs,
+    normalize_reset_out,
+    normalize_step_out,
+    to_numpy,
+)
 
 def back_to_like(t, numpy_arr):
     if is_tensor(t):
@@ -44,8 +54,13 @@ class TimeLimitWrapper(EnvWrapper):
         self.t = np.zeros(self.num_envs, dtype = int)
 
     def reset(self, **kwargs):
+        obs, info = normalize_reset_out(self.env.reset(**kwargs))
+
+        if self.is_vector and exists(batch_size := get_batch_size(obs)):
+            self.num_envs = batch_size
+
         self.t = np.zeros(self.num_envs, dtype = int)
-        return normalize_reset_out(self.env.reset(**kwargs))
+        return obs, info
 
     def step(self, action):
         obs, reward, terminated, truncated, info = normalize_step_out(self.env.step(action))

{env_ssl_wrapper-0.4.1 → env_ssl_wrapper-0.4.2}/env_ssl_wrapper/standardize/utils.py

@@ -25,29 +25,27 @@ WRAPPERS = dict(
     pad_episodes = EpisodePaddingWrapper
 )
 
+def get_wrapper(name):
+    if name in ('standardize_env', 'master'):
+        from .standardize_env_wrapper import StandardizeEnvWrapper
+        return StandardizeEnvWrapper
+
+    if name == 'memory_trace':
+        from ..memory_trace import MemoryTraceWrapper
+        return MemoryTraceWrapper
+
+    if name in WRAPPERS:
+        return WRAPPERS[name]
+
+    raise ValueError(f'unknown wrapper {name!r} — choose from {sorted([*WRAPPERS, "memory_trace", "standardize_env"])}')
+
 def parse_wrapper(wrapper):
     if isinstance(wrapper, str):
-        if wrapper == 'memory_trace':
-            from ..memory_trace import MemoryTraceWrapper
-            wrapper = MemoryTraceWrapper
-        elif wrapper in ('standardize_env', 'master'):
-            from .standardize_env_wrapper import StandardizeEnvWrapper
-            wrapper = StandardizeEnvWrapper
-        elif wrapper not in WRAPPERS:
-            raise ValueError(f'unknown wrapper {wrapper!r} — choose from {sorted([*WRAPPERS, "memory_trace", "standardize_env"])}')
-        else:
-            wrapper = WRAPPERS[wrapper]
+        wrapper = get_wrapper(wrapper)
 
     if isinstance(wrapper, tuple):
         name, kwargs = wrapper
-        if name == 'memory_trace':
-            from ..memory_trace import MemoryTraceWrapper
-            wrapper = partial(MemoryTraceWrapper, **kwargs)
-        elif name in ('standardize_env', 'master'):
-            from .standardize_env_wrapper import StandardizeEnvWrapper
-            wrapper = partial(StandardizeEnvWrapper, **kwargs)
-        else:
-            wrapper = partial(WRAPPERS.get(name, name), **kwargs)
+        wrapper = partial(get_wrapper(name) if isinstance(name, str) else name, **kwargs)
 
     elif isinstance(wrapper, dict):
         raise ValueError("wrapper kwargs must be passed as (name, kwargs), e.g. ('tensor', dict(device = 'cpu'))")

{env_ssl_wrapper-0.4.1 → env_ssl_wrapper-0.4.2}/env_ssl_wrapper/standardize/vector.py

@@ -22,8 +22,6 @@ from .standardize_wrapper import StandardizeWrapper
 
 # leaf helpers
 
-_stack_trees = stack_trees
-
 def _merge_infos(infos):
     infos = [info if isinstance(info, dict) else {} for info in infos]
 
@@ -38,7 +36,7 @@ def _merge_infos(infos):
     for key in keys:
         vals = [info[key] for info in infos]
         try:
-            out[key] = _stack_trees(vals)
+            out[key] = stack_trees(vals)
         except Exception:
             out[key] = np.array(vals, dtype = object)
 
@@ -276,7 +274,7 @@ class MultiprocessingVecEnv:
             _safe_send(conn, ('reset', kwargs))
 
         results = _recv_all(self._conns)
-        return _stack_trees([obs for obs, _ in results]), _merge_infos([info for _, info in results])
+        return stack_trees([obs for obs, _ in results]), _merge_infos([info for _, info in results])
 
     def step(self, actions):
         actions = _split_actions(actions, self.num_envs)
@@ -302,10 +300,10 @@ class MultiprocessingVecEnv:
                 for final, obs_i in zip(final_obs, obs)
             ]
 
-            info['final_observation'] = _stack_trees(final_obs)
+            info['final_observation'] = stack_trees(final_obs)
             info['_final_observation'] = done
 
-        return _stack_trees(obs), reward, terminated, truncated, info
+        return stack_trees(obs), reward, terminated, truncated, info
 
     def close(self):
         _shutdown(self._conns, self._procs)

{env_ssl_wrapper-0.4.1 → env_ssl_wrapper-0.4.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "env-ssl-wrapper"
-version = "0.4.1"
+version = "0.4.2"
 description = "One torch-native interface for any MDP environment"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }
@@ -54,6 +54,9 @@ test = [
     "pillow",
     "imageio",
     "populora",
+    "gym-pusht",
+    "pymunk<7",
+    "mean-conc-beta",
 ]
 
 [tool.pytest.ini_options]