env-ssl-wrapper 0.3.0__tar.gz → 0.4.0__tar.gz

{env_ssl_wrapper-0.3.0 → env_ssl_wrapper-0.4.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: env-ssl-wrapper
-Version: 0.3.0
+Version: 0.4.0
 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
@@ -43,6 +43,7 @@ Requires-Dist: torch>=2.5
 Requires-Dist: x-mlps-pytorch
 Requires-Dist: x-transformers
 Provides-Extra: examples
+Requires-Dist: fire; extra == 'examples'
 Provides-Extra: test
 Requires-Dist: box2d-py>=2.3.5; extra == 'test'
 Requires-Dist: dm-control; extra == 'test'
@@ -67,19 +68,18 @@ One line turns any simulator's environment — mujoco warp, isaac sim, pybullet,
 pip install env-ssl-wrapper
 ```
 
-## Usage
+## Standardize
 
 ```python
 import torch
-from env_ssl_wrapper import compose_env
+from env_ssl_wrapper import StandardizeEnvWrapper
 
-env = compose_env(
-    any_env,                                 # any env from any sim
-    ('tensor', dict(device='cpu')),          # wrap with whatever you need
-    'done_tracker',
-)
+# one master wrapper to standardize any simulator environment
+
+env = StandardizeEnvWrapper(any_env)
+
+obs, info = env.reset() # torch.float32, batched
 
-obs, info = env.reset()                      # torch.float32, batched
 while not env.all_done:
     actions = torch.randint(0, 2, (8,))
     obs, reward, terminated, truncated, info = env.step(actions)
@@ -87,6 +87,50 @@ while not env.all_done:
 
 Works identically for every simulator.
 
+Or compose individual wrappers piecemeal:
+
+```python
+from env_ssl_wrapper import compose_env
+
+env = compose_env(
+    any_env,
+    ('tensor', dict(device = 'cpu')),
+    'done_tracker',
+)
+```
+
+## Research Wrappers
+
+### Memory Trace
+
+A cheap memory that keeps track of an exponential moving average (EMA) of a state, or subset ([Eberhard et al., 2025](https://arxiv.org/abs/2503.15200)).
+
+```python
+import gymnasium as gym
+from env_ssl_wrapper import StandardizeEnvWrapper, MemoryTraceWrapper
+
+env = StandardizeEnvWrapper(gym.make('LunarLander-v3'))
+env = MemoryTraceWrapper(env, lambdas = (0.9, 0.99))
+
+obs, info = env.reset()
+# obs['obs']        -> (1, 8)
+# obs['trace_0.9']  -> (1, 8)
+# obs['trace_0.99'] -> (1, 8)
+```
+
+Or on a subset of dictionary observations:
+
+```python
+env = MemoryTraceWrapper(env, lambdas = (0.9, 0.99), keys = 'proprio')
+# obs['proprio_trace_0.9'] -> (1, 2)
+```
+
+Run the PPO benchmark on POMDP LunarLander:
+
+```bash
+uv run test_memory_trace.py
+```
+
 ## Wrappers
 
 Pass wrappers as strings (default config) or `(name, dict)` tuples (custom config), in any order.
@@ -166,15 +210,6 @@ obs, info = env.reset()
 obs, reward, terminated, truncated, info = env.step(torch.randn(1, 1))
 ```
 
-## Mock sims
-
-`env_ssl_wrapper.mocks` ships dependency-free stand-ins emulating each simulator's quirks (`GymnasiumMockEnv`, `IsaacMockEnv`, `DMControlMockEnv`, ...) for testing your code without installing the real sims.
-
-```python
-from env_ssl_wrapper.mocks import IsaacMockEnv
-env = compose_env(IsaacMockEnv(), 'tensor', 'done_tracker')
-```
-
 ## Multiprocessing
 
 Parallelize any single environment or factory into an autoresetting vector env:
@@ -191,5 +226,16 @@ with MultiprocessingVecEnv('CartPole-v1', num_envs = 8) as env:
 
 ```bash
 uv sync --extra test
-uv run pytest tests/test_real_envs.py
+uv run pytest
+```
+
+## Citations
+
+```bibtex
+@article{eberhard2025partially,
+    title   = {Partially Observable Reinforcement Learning with Memory Traces},
+    author  = {Onno Eberhard and Michael Muehlebach and Claire Vernade},
+    journal = {arXiv preprint arXiv:2503.15200},
+    year    = {2025}
+}
 ```

{env_ssl_wrapper-0.3.0 → env_ssl_wrapper-0.4.0}/README.md

@@ -8,19 +8,18 @@ One line turns any simulator's environment — mujoco warp, isaac sim, pybullet,
 pip install env-ssl-wrapper
 ```
 
-## Usage
+## Standardize
 
 ```python
 import torch
-from env_ssl_wrapper import compose_env
+from env_ssl_wrapper import StandardizeEnvWrapper
 
-env = compose_env(
-    any_env,                                 # any env from any sim
-    ('tensor', dict(device='cpu')),          # wrap with whatever you need
-    'done_tracker',
-)
+# one master wrapper to standardize any simulator environment
+
+env = StandardizeEnvWrapper(any_env)
+
+obs, info = env.reset() # torch.float32, batched
 
-obs, info = env.reset()                      # torch.float32, batched
 while not env.all_done:
     actions = torch.randint(0, 2, (8,))
     obs, reward, terminated, truncated, info = env.step(actions)
@@ -28,6 +27,50 @@ while not env.all_done:
 
 Works identically for every simulator.
 
+Or compose individual wrappers piecemeal:
+
+```python
+from env_ssl_wrapper import compose_env
+
+env = compose_env(
+    any_env,
+    ('tensor', dict(device = 'cpu')),
+    'done_tracker',
+)
+```
+
+## Research Wrappers
+
+### Memory Trace
+
+A cheap memory that keeps track of an exponential moving average (EMA) of a state, or subset ([Eberhard et al., 2025](https://arxiv.org/abs/2503.15200)).
+
+```python
+import gymnasium as gym
+from env_ssl_wrapper import StandardizeEnvWrapper, MemoryTraceWrapper
+
+env = StandardizeEnvWrapper(gym.make('LunarLander-v3'))
+env = MemoryTraceWrapper(env, lambdas = (0.9, 0.99))
+
+obs, info = env.reset()
+# obs['obs']        -> (1, 8)
+# obs['trace_0.9']  -> (1, 8)
+# obs['trace_0.99'] -> (1, 8)
+```
+
+Or on a subset of dictionary observations:
+
+```python
+env = MemoryTraceWrapper(env, lambdas = (0.9, 0.99), keys = 'proprio')
+# obs['proprio_trace_0.9'] -> (1, 2)
+```
+
+Run the PPO benchmark on POMDP LunarLander:
+
+```bash
+uv run test_memory_trace.py
+```
+
 ## Wrappers
 
 Pass wrappers as strings (default config) or `(name, dict)` tuples (custom config), in any order.
@@ -107,15 +150,6 @@ obs, info = env.reset()
 obs, reward, terminated, truncated, info = env.step(torch.randn(1, 1))
 ```
 
-## Mock sims
-
-`env_ssl_wrapper.mocks` ships dependency-free stand-ins emulating each simulator's quirks (`GymnasiumMockEnv`, `IsaacMockEnv`, `DMControlMockEnv`, ...) for testing your code without installing the real sims.
-
-```python
-from env_ssl_wrapper.mocks import IsaacMockEnv
-env = compose_env(IsaacMockEnv(), 'tensor', 'done_tracker')
-```
-
 ## Multiprocessing
 
 Parallelize any single environment or factory into an autoresetting vector env:
@@ -132,5 +166,16 @@ with MultiprocessingVecEnv('CartPole-v1', num_envs = 8) as env:
 
 ```bash
 uv sync --extra test
-uv run pytest tests/test_real_envs.py
+uv run pytest
+```
+
+## Citations
+
+```bibtex
+@article{eberhard2025partially,
+    title   = {Partially Observable Reinforcement Learning with Memory Traces},
+    author  = {Onno Eberhard and Michael Muehlebach and Claire Vernade},
+    journal = {arXiv preprint arXiv:2503.15200},
+    year    = {2025}
+}
 ```

env_ssl_wrapper-0.4.0/env_ssl_wrapper/__init__.py

@@ -0,0 +1,47 @@
+from __future__ import annotations
+
+import sys
+from . import standardize
+from .standardize import *
+
+from .memory_trace import MemoryTraceWrapper
+
+# Wire backwards-compatibility aliases in sys.modules and module globals
+# so imports like `from env_ssl_wrapper.done_tracker_wrapper import DoneTrackerWrapper`
+# or `import env_ssl_wrapper.mocks` continue to work without breaking.
+
+_STANDARDIZE_SUBMODULES = (
+    'adapters',
+    'auto_batched_wrapper',
+    'action_transform_wrapper',
+    'done_tracker_wrapper',
+    'episode_padding_wrapper',
+    'flatten_obs_wrapper',
+    'helpers',
+    'image_wrapper',
+    'mocks',
+    'spaces',
+    'standardize_wrapper',
+    'tensor_wrapper',
+    'time_limit_wrapper',
+    'standardize_env_wrapper',
+    'utils',
+    'vector',
+)
+
+for _name in _STANDARDIZE_SUBMODULES:
+    _mod = getattr(standardize, _name, None)
+    if _mod is not None:
+        sys.modules[f'{__name__}.{_name}'] = _mod
+        globals()[_name] = _mod
+
+__all__ = [
+    *standardize.__all__,
+    'MemoryTraceWrapper',
+]
+
+def __getattr__(name):
+    if hasattr(standardize, name):
+        return getattr(standardize, name)
+    raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
+

env_ssl_wrapper-0.4.0/env_ssl_wrapper/memory_trace.py

@@ -0,0 +1,93 @@
+from __future__ import annotations
+from typing import Sequence
+
+import torch
+from torch import is_tensor
+from torch_einops_utils import pad_right_ndim_to
+
+from .standardize.helpers import (
+    TransformObservationWrapper,
+    exists,
+    default,
+    any_true,
+)
+
+# 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
+
+    # reset trace on done, else decay
+
+    if not is_tensor(done):
+        done = torch.as_tensor(done, device = x.device)
+
+    weight = torch.where(done, 1., 1. - lam).to(x)
+    return pad_right_ndim_to(weight, x.ndim)
+
+# Partially Observable Reinforcement Learning with Memory Traces - Eberhard et al.
+# https://arxiv.org/abs/2503.15200
+
+class MemoryTraceWrapper(TransformObservationWrapper):
+
+    def __init__(
+        self,
+        env,
+        lambdas: float | Sequence[float] = (0.9, 0.99),
+        trace_key: str = 'trace',
+        obs_key: str = 'obs',
+        keys: str | Sequence[str] | None = None,
+    ):
+        super().__init__(env)
+
+        self.lambdas = tuple(float(l) for l in cast_tuple(lambdas))
+        assert all(0. <= l <= 1. for l in self.lambdas), f'lambdas must be within [0, 1], got {self.lambdas}'
+
+        self.trace_key = trace_key
+        self.obs_key = obs_key
+        self.keys = set(cast_tuple(keys)) if exists(keys) else None
+
+        self.traces = dict()
+
+    def trace_key_for(self, key, lam):
+        prefix = self.trace_key if key == self.obs_key else f'{key}_{self.trace_key}'
+        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}
+        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):
+                continue
+
+            if not is_tensor(val):
+                val = torch.as_tensor(val)
+                out[key] = val
+
+            val_float = val.float() if not val.is_floating_point() else val
+
+            # init or lerp update traces
+            # z_t = λ * z_{t-1} + (1 - λ) * y_t
+
+            if key not in self.traces:
+                traces = [val_float.clone() for _ in self.lambdas]
+            else:
+                traces = [prev.lerp(val_float, calc_lerp_weight(lam, done, val_float)) for lam, prev in zip(self.lambdas, self.traces[key])]
+
+            self.traces[key] = traces
+
+            for lam, trace in zip(self.lambdas, traces):
+                out[self.trace_key_for(key, lam)] = trace
+
+        return out
+
+    def reset(self, **kwargs):
+        self.traces = dict()
+        return super().reset(**kwargs)

env_ssl_wrapper-0.4.0/env_ssl_wrapper/standardize/__init__.py

@@ -0,0 +1,166 @@
+from __future__ import annotations
+
+from . import (
+    adapters,
+    auto_batched_wrapper,
+    action_transform_wrapper,
+    done_tracker_wrapper,
+    episode_padding_wrapper,
+    flatten_obs_wrapper,
+    helpers,
+    image_wrapper,
+    mocks,
+    spaces,
+    standardize_wrapper,
+    tensor_wrapper,
+    time_limit_wrapper,
+    utils,
+    vector,
+)
+
+from .standardize_wrapper import StandardizeWrapper
+from .image_wrapper import ImageObservationWrapper
+from .auto_batched_wrapper import AutoBatchedWrapper
+from .tensor_wrapper import TensorWrapper
+from .action_transform_wrapper import ActionTransformWrapper
+from .done_tracker_wrapper import DoneTrackerWrapper
+from .flatten_obs_wrapper import FlattenObsWrapper
+from .episode_padding_wrapper import EpisodePaddingWrapper
+from .time_limit_wrapper import TimeLimitWrapper
+from .vector import MultiprocessingVecEnv
+from .standardize_env_wrapper import StandardizeEnvWrapper, StandardizeEnv, StandardizedEnv
+
+from .adapters import (
+    BaseEnvAdapter,
+    get_adapter,
+    register_adapter,
+    MujocoWarpAdapter,
+    IsaacAdapter,
+    PyBulletAdapter,
+    DMControlAdapter,
+    PufferLibAdapter,
+    RoboticsAdapter,
+    GymnasiumAdapter,
+    LegacyGymAdapter,
+    DefaultAdapter,
+)
+
+from .spaces import (
+    InferredSpace,
+    infer_observation_space,
+    space_from_action_spec,
+    action_space_dim,
+    action_space_is_discrete,
+    action_space_is_box,
+    action_space_bounds,
+    action_dim_of,
+    obs_dim_of,
+)
+
+from .mocks import (
+    MockEnv,
+    GymnasiumMockEnv,
+    GymnasiumDiscreteMockEnv,
+    LegacyGymMockEnv,
+    PyBulletMockEnv,
+    DMControlMockEnv,
+    IsaacMockEnv,
+    AutoresetVectorMockEnv,
+    PufferVectorMockEnv,
+    PufferTensorMockEnv,
+    ManiSkillMockEnv,
+    BraxMockEnv,
+    MjxMockEnv,
+    MetaWorldMockEnv,
+    TrifingerMockEnv,
+    HabitatMockEnv,
+    TupleObsMockEnv,
+    JaxArray,
+)
+
+from .utils import wrap_env, compose_env
+from .helpers import (
+    TransformObservationWrapper,
+    ObservationWrapper,
+    has_final_observation,
+    maybe_get_final_observation,
+    get_final_observation,
+    maybe_transform_final_observation,
+)
+
+__all__ = [
+    'TransformObservationWrapper',
+    'ObservationWrapper',
+    'adapters',
+    'auto_batched_wrapper',
+    'action_transform_wrapper',
+    'done_tracker_wrapper',
+    'episode_padding_wrapper',
+    'flatten_obs_wrapper',
+    'helpers',
+    'image_wrapper',
+    'mocks',
+    'spaces',
+    'standardize_wrapper',
+    'tensor_wrapper',
+    'time_limit_wrapper',
+    'utils',
+    'vector',
+    'StandardizeWrapper',
+    'ImageObservationWrapper',
+    'AutoBatchedWrapper',
+    'TensorWrapper',
+    'ActionTransformWrapper',
+    'DoneTrackerWrapper',
+    'FlattenObsWrapper',
+    'EpisodePaddingWrapper',
+    'TimeLimitWrapper',
+    'MultiprocessingVecEnv',
+    'StandardizeEnvWrapper',
+    'StandardizeEnv',
+    'BaseEnvAdapter',
+    'get_adapter',
+    'register_adapter',
+    'MujocoWarpAdapter',
+    'IsaacAdapter',
+    'PyBulletAdapter',
+    'DMControlAdapter',
+    'PufferLibAdapter',
+    'RoboticsAdapter',
+    'GymnasiumAdapter',
+    'LegacyGymAdapter',
+    'DefaultAdapter',
+    'InferredSpace',
+    'infer_observation_space',
+    'space_from_action_spec',
+    'action_space_dim',
+    'action_space_is_discrete',
+    'action_space_is_box',
+    'action_space_bounds',
+    'action_dim_of',
+    'obs_dim_of',
+    'MockEnv',
+    'GymnasiumMockEnv',
+    'GymnasiumDiscreteMockEnv',
+    'LegacyGymMockEnv',
+    'PyBulletMockEnv',
+    'DMControlMockEnv',
+    'IsaacMockEnv',
+    'AutoresetVectorMockEnv',
+    'PufferVectorMockEnv',
+    'PufferTensorMockEnv',
+    'ManiSkillMockEnv',
+    'BraxMockEnv',
+    'MjxMockEnv',
+    'MetaWorldMockEnv',
+    'TrifingerMockEnv',
+    'HabitatMockEnv',
+    'TupleObsMockEnv',
+    'JaxArray',
+    'wrap_env',
+    'compose_env',
+    'has_final_observation',
+    'maybe_get_final_observation',
+    'get_final_observation',
+    'maybe_transform_final_observation',
+]

{env_ssl_wrapper-0.3.0/env_ssl_wrapper → env_ssl_wrapper-0.4.0/env_ssl_wrapper/standardize}/adapters.py

@@ -175,7 +175,7 @@ class WrapperAdapter(BaseEnvAdapter):
     @property
     def autoresets(self) -> bool:
         val = first_existing(self.env, 'autoreset', 'autoresets', 'autoreset_mode')
-        return bool(val) if exists(val) else self.inner_adapter.autoresets
+        return truthy_attr(val) if exists(val) else self.inner_adapter.autoresets
 
     @property
     def action_space(self):
@@ -202,19 +202,6 @@ class DMControlAdapter(BaseEnvAdapter):
         has_physics = exists(get_attr(env, 'physics')) and callable(get_attr(get_attr(env, 'physics'), 'render'))
         return is_dm_type or has_physics
 
-    def step(self, action):
-        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)
-        return super().step(action)
-
-    def reset(self, **kwargs):
-        out = self.env.reset(**kwargs)
-        if is_time_step(out):
-            return out.observation, {}
-        return super().reset(**kwargs)
-
     def seed(self, seed: int):
         random_state = get_attr(get_attr(self.env, 'task'), '_random')
         if exists(random_state) and callable(get_attr(random_state, 'seed')):

{env_ssl_wrapper-0.3.0/env_ssl_wrapper → env_ssl_wrapper-0.4.0/env_ssl_wrapper/standardize}/auto_batched_wrapper.py

@@ -39,6 +39,12 @@ def maybe_expand_dim(x):
 
         return rearrange(arr, '... -> 1 ...')
 
+    if is_tensor(x):
+        return rearrange(x, '-> 1') if x.ndim == 0 else rearrange(x, '... -> 1 ...')
+
+    if isinstance(x, np.ndarray) and x.dtype.kind in 'biufc':
+        return rearrange(x, '-> 1') if x.ndim == 0 else rearrange(x, '... -> 1 ...')
+
     return tree_map(_expand, x)
 
 def is_integer_dtype(t):

{env_ssl_wrapper-0.3.0/env_ssl_wrapper → env_ssl_wrapper-0.4.0/env_ssl_wrapper/standardize}/done_tracker_wrapper.py

@@ -12,6 +12,7 @@ from .helpers import (
     env_num_envs,
     exists,
     get_attr,
+    is_array,
     is_vectorized,
     to_numpy,
 )
@@ -25,6 +26,10 @@ def get_batch_size(tree) -> int | None:
         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

{env_ssl_wrapper-0.3.0/env_ssl_wrapper → env_ssl_wrapper-0.4.0/env_ssl_wrapper/standardize}/episode_padding_wrapper.py

@@ -8,7 +8,7 @@ from torch import is_tensor
 from torch.utils._pytree import tree_map
 from einops import rearrange
 
-from .helpers import EnvWrapper, copy_leaf, dones_of, exists, is_vectorized, to_numpy
+from .helpers import EnvWrapper, copy_leaf, dones_of, env_autoresets, exists, is_vectorized, to_numpy
 
 # helpers
 
@@ -63,6 +63,7 @@ class EpisodePaddingWrapper(EnvWrapper):
     def __init__(self, env):
         super().__init__(env)
         self.is_vector = is_vectorized(env)
+        self.autoreset = env_autoresets(env)
         self._last_obs = None
         self._is_done = None
         self._final_obs = None
@@ -83,12 +84,15 @@ 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):
+                self._is_done = np.zeros(len(mask), dtype = bool)
+
+            if self.autoreset:
+                self._is_done &= mask
+
             if mask.any():
                 assert exists(self._last_obs), 'environment needs reset before calling step. call env.reset() first'
 
-                if self._is_done is None or len(self._is_done) != len(mask):
-                    self._is_done = np.zeros(len(mask), dtype = bool)
-
                 newly = mask & ~self._is_done
                 self._is_done |= mask
 

env_ssl_wrapper-0.4.0/env_ssl_wrapper/standardize/flatten_obs_wrapper.py

@@ -0,0 +1,75 @@
+from __future__ import annotations
+
+import numpy as np
+import torch
+from torch import is_tensor
+from torch.utils._pytree import tree_flatten
+from einops import rearrange
+
+from .helpers import (
+    TransformObservationWrapper,
+    default,
+    is_array,
+    is_scalar,
+    is_vectorized,
+)
+
+# helpers
+
+def flattenable(t):
+    if is_array(t) or is_scalar(t):
+        return True
+
+    try:
+        return np.asarray(t).dtype.kind not in 'USO'
+    except Exception:
+        return False
+
+def flatten_leaf(t, is_vector = False):
+    if not is_array(t):
+        t = np.asarray(t)
+
+    if not is_vector:
+        if t.ndim == 0:
+            return rearrange(t, '-> 1')
+        return rearrange(t, '... -> (...)')
+
+    if t.ndim == 0:
+        return rearrange(t, '-> 1 1')
+
+    if t.ndim == 1:
+        return rearrange(t, 'b -> b 1')
+
+    return rearrange(t, 'b ... -> b (...)')
+
+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]):
+        return torch.cat(leaves, dim = axis)
+
+    return np.concatenate(leaves, axis = axis)
+
+# class
+
+class FlattenObsWrapper(TransformObservationWrapper):
+    def __init__(self, env, is_vector: bool | None = None):
+        super().__init__(env)
+        self.is_vector = default(is_vector, is_vectorized(env))
+
+    def transform_obs(self, obs):
+        if is_array(obs):
+            return obs
+
+        leaves, _ = tree_flatten(obs)
+        leaves = [t for t in leaves if flattenable(t)]
+
+        if len(leaves) == 0:
+            return obs
+
+        if len(leaves) == 1:
+            return flatten_leaf(leaves[0], is_vector = self.is_vector)
+
+        return concat_leaves(leaves, is_vector = self.is_vector)
+

{env_ssl_wrapper-0.3.0/env_ssl_wrapper → env_ssl_wrapper-0.4.0/env_ssl_wrapper/standardize}/helpers.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import inspect
 import numpy as np
 from torch import is_tensor
 from torch.utils._pytree import tree_map
@@ -62,6 +63,8 @@ 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)
 
 # environment probes
@@ -85,6 +88,33 @@ def env_render(env, height, width, camera = None):
 def is_vectorized(env) -> bool:
     return get_adapter(env).is_vectorized
 
+def has_final_observation(info):
+    return isinstance(info, dict) and 'final_observation' in info
+
+def maybe_get_final_observation(info):
+    if isinstance(info, dict):
+        return info.get('final_observation')
+
+    return None
+
+def get_final_observation(info, default = None):
+    final_obs = maybe_get_final_observation(info)
+
+    if exists(final_obs):
+        return final_obs
+
+    if exists(default):
+        return default
+
+    raise KeyError("no 'final_observation' found in info")
+
+def maybe_transform_final_observation(info, fn):
+    if not has_final_observation(info):
+        return info
+
+    info['final_observation'] = fn(info['final_observation'])
+    return info
+
 def mark_terminal_obs(info, obs, dones, is_vector):
     # single-env terminal contract — vector envs handled by EpisodePaddingWrapper
     if not is_vector and isinstance(info, dict) and 'final_observation' not in info and any_true(dones):
@@ -132,3 +162,60 @@ class EnvWrapper:
         if name.startswith('_'):
             raise AttributeError(f"attempted to get missing private attribute '{name}'")
         return getattr(self.env, name)
+
+def accepts_done_param(fn):
+    try:
+        params = inspect.signature(fn).parameters
+        return 'done' in params or any(p.kind == inspect.Parameter.VAR_KEYWORD for p in params.values())
+    except (ValueError, TypeError):
+        return False
+
+class TransformObservationWrapper(EnvWrapper):
+    """
+    Base observation wrapper that automatically handles:
+    - Calling transform_obs on observations in reset() and step()
+    - Detecting environment autoreset and passing `done` to stateful transforms
+    - Propagating transformed observations to info['final_observation']
+    """
+
+    def __init__(self, env):
+        super().__init__(env)
+        self.autoresets = env_autoresets(env)
+        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)
+
+        if isinstance(info, dict) and 'final_observation' in info:
+            info['final_observation'] = self.transform_obs(info['final_observation'])
+
+        return obs, info
+
+    def step(self, action):
+        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)
+
+        if isinstance(info, dict) and 'final_observation' in info:
+            if not self.takes_done:
+                info['final_observation'] = self.transform_obs(info['final_observation'])
+            elif not self.autoresets:
+                info['final_observation'] = out
+
+        return out, reward, terminated, truncated, info
+
+ObservationWrapper = TransformObservationWrapper
+

env_ssl_wrapper-0.4.0/env_ssl_wrapper/standardize/standardize_env_wrapper.py

@@ -0,0 +1,66 @@
+from __future__ import annotations
+
+import torch
+from .helpers import EnvWrapper, exists, default
+from .utils import compose_env
+
+class StandardizeEnvWrapper(EnvWrapper):
+    """
+    Master environment wrapper that turns any simulator environment
+    into a torch-native, batched, standardized RL environment.
+    """
+
+    def __init__(
+        self,
+        env,
+        device: str | torch.device = 'cpu',
+        auto_batch: bool = True,
+        pad_episodes: bool = True,
+        done_tracker: bool = True,
+        flatten_obs: bool = False,
+        action_transform: bool = False,
+        max_timesteps: int | None = None,
+        image_size: tuple[int, int] | None = None,
+        lambdas: tuple[float, ...] | list[float] | None = None,
+        **kwargs
+    ):
+        wrappers = []
+
+        if exists(image_size):
+            wrappers.append(('image', dict(image_size = image_size)))
+
+        if action_transform:
+            wrappers.append('action_transform')
+
+        if exists(max_timesteps):
+            wrappers.append(('time_limit', dict(max_timesteps = max_timesteps)))
+
+        if auto_batch:
+            wrappers.append('auto_batch')
+
+        if exists(device):
+            wrappers.append(('tensor', dict(device = device)))
+
+        if done_tracker:
+            wrappers.append('done_tracker')
+
+        if exists(lambdas):
+            wrappers.append(('memory_trace', dict(lambdas = lambdas, **kwargs)))
+
+        if flatten_obs:
+            wrappers.append('flatten_obs')
+
+        wrapped = compose_env(env, *wrappers, pad_episodes = pad_episodes)
+
+        super().__init__(wrapped)
+
+    def reset(self, **kwargs):
+        return self.env.reset(**kwargs)
+
+    def step(self, action):
+        return self.env.step(action)
+
+# aliases
+
+StandardizeEnv = StandardizeEnvWrapper
+StandardizedEnv = StandardizeEnvWrapper

{env_ssl_wrapper-0.3.0/env_ssl_wrapper → env_ssl_wrapper-0.4.0/env_ssl_wrapper/standardize}/tensor_wrapper.py

@@ -11,43 +11,49 @@ from .helpers import EnvWrapper, exists, get_attr, is_scalar
 
 # helpers
 
+def to_torch_leaf(t, device, cast_obs_to_float = True):
+    if not is_tensor(t):
+        if isinstance(t, np.ndarray):
+            t = from_numpy(t.copy())
+        elif is_scalar(t):
+            t = tensor(t)
+        elif exists(get_attr(t, '__array__')):
+            t = from_numpy(np.asarray(t))
+        else:
+            return t
+
+    dtype = torch.float32 if cast_obs_to_float and t.dtype != torch.bool else t.dtype
+    return t.to(device = device, dtype = dtype)
+
 def numpy_to_torch(x, device, cast_obs_to_float = True):
     # numpy / scalars / foreign array-likes to torch; float32 unless disabled
+    if not isinstance(x, (dict, list, tuple)):
+        return to_torch_leaf(x, device, cast_obs_to_float)
 
-    def _to_torch(t):
-        if not is_tensor(t):
-            if isinstance(t, np.ndarray):
-                t = from_numpy(np.array(t))
-            elif is_scalar(t):
-                t = tensor(t)
-            elif exists(get_attr(t, '__array__')):
-                t = from_numpy(np.array(t))
-            else:
-                return t
-
-        dtype = torch.float32 if cast_obs_to_float and t.dtype != torch.bool else t.dtype
-        return t.to(device = device, dtype = dtype)
-    return tree_map(_to_torch, x)
+    return tree_map(partial(to_torch_leaf, device = device, cast_obs_to_float = cast_obs_to_float), x)
 
-def torch_to_numpy(x):
-    # torch to numpy; 0-dim collapses to scalar, float64 → float32
+def to_numpy_leaf(t):
+    if is_tensor(t):
+        t = t.detach().cpu().numpy()
+    elif is_scalar(t):
+        t = np.asarray(t)
+    else:
+        return t
 
-    def _to_numpy(t):
-        if is_tensor(t):
-            t = t.detach().cpu().numpy()
-        elif is_scalar(t):
-            t = np.asarray(t)
-        else:
-            return t
+    if t.ndim == 0:
+        return t.item()
 
-        if t.ndim == 0:
-            return t.item()
+    if t.dtype == np.float64:
+        t = t.astype(np.float32)
 
-        if t.dtype == np.float64:
-            t = t.astype(np.float32)
+    return t
 
-        return t
-    return tree_map(_to_numpy, x)
+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
 
@@ -96,6 +102,12 @@ class TensorWrapper(EnvWrapper):
 
         return obs, info
 
+    def to_contract(self, t, to_float = False):
+        if not isinstance(t, (dict, list, tuple)):
+            leaf = to_torch_leaf(t, self.device, cast_obs_to_float = False)
+            return contract(leaf, to_float = to_float)
+        return contract(self.cast(t), to_float = to_float)
+
     def step(self, action):
         action = torch_to_numpy(action) if self.convert_in else action
         obs, reward, terminated, truncated, info = self.env.step(action)
@@ -104,9 +116,9 @@ class TensorWrapper(EnvWrapper):
             return obs, reward, terminated, truncated, info
 
         obs = self.cast(obs)
-        reward = contract(self.cast(reward), to_float = True)
-        terminated = contract(self.cast(terminated))
-        truncated = contract(self.cast(truncated))
+        reward = self.to_contract(reward, to_float = True)
+        terminated = self.to_contract(terminated)
+        truncated = self.to_contract(truncated)
         self.cast_info(info)
 
         return obs, reward, terminated, truncated, info

{env_ssl_wrapper-0.3.0/env_ssl_wrapper → env_ssl_wrapper-0.4.0/env_ssl_wrapper/standardize}/time_limit_wrapper.py

@@ -26,8 +26,15 @@ def back_to_like(t, numpy_arr):
     return numpy_arr
 
 class TimeLimitWrapper(EnvWrapper):
-    # caps episode length at max_timesteps, setting truncated = True for
-    # capped envs (vectorized and single alike); timers reset per episode
+    """
+    caps episode length at max_timesteps, setting truncated = True for
+    capped envs (vectorized and single alike); timers reset per episode.
+
+    Note on vector environments:
+    Marking a slot as truncated signals the time limit to downstream consumers,
+    but does not stop the underlying simulator slot if it runs continuously or
+    autoresets. Downstream loops should check `truncated` to reset consumer state.
+    """
 
     def __init__(self, env, max_timesteps):
         super().__init__(env)

{env_ssl_wrapper-0.3.0/env_ssl_wrapper → env_ssl_wrapper-0.4.0/env_ssl_wrapper/standardize}/utils.py

@@ -27,14 +27,27 @@ WRAPPERS = dict(
 
 def parse_wrapper(wrapper):
     if isinstance(wrapper, str):
-        if wrapper not in WRAPPERS:
-            raise ValueError(f'unknown wrapper {wrapper!r} — choose from {sorted(WRAPPERS)}')
-
-        wrapper = WRAPPERS[wrapper]
+        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]
 
     if isinstance(wrapper, tuple):
         name, kwargs = wrapper
-        wrapper = partial(WRAPPERS.get(name, name), **kwargs)
+        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)
 
     elif isinstance(wrapper, dict):
         raise ValueError("wrapper kwargs must be passed as (name, kwargs), e.g. ('tensor', dict(device = 'cpu'))")
@@ -63,6 +76,17 @@ def compose_env(env, *wrappers, pad_episodes: bool = True):
         funcs.insert(1, EpisodePaddingWrapper)
         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)
+
     assert len(set(classes)) == len(classes), 'duplicate wrappers found'
 
     for func in funcs:

{env_ssl_wrapper-0.3.0/env_ssl_wrapper → env_ssl_wrapper-0.4.0/env_ssl_wrapper/standardize}/vector.py

@@ -6,7 +6,7 @@ import torch
 from torch import is_tensor
 from torch.utils._pytree import tree_flatten, tree_map, tree_structure, tree_unflatten
 
-from .helpers import any_true, dones_of, exists, get_attr, instantiate_env, safe_close
+from .helpers import any_true, dones_of, exists, get_attr, instantiate_env, safe_close, truthy_attr
 from .spaces import action_dim_of
 from .standardize_wrapper import StandardizeWrapper
 
@@ -23,6 +23,20 @@ def _stack_leaves(leaves):
     return np.stack(leaves)
 
 def _stack_trees(trees):
+    first = trees[0]
+
+    if is_tensor(first):
+        return torch.stack(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]))
@@ -57,7 +71,10 @@ def _exec(env, cmd, payload):
         obs, reward, terminated, truncated, info = env.step(payload)
 
         if any_true(dones_of(terminated, truncated)):
-            obs, final_obs = env.reset()[0], obs
+            if truthy_attr(get_attr(env.adapter, 'autoresets')):
+                final_obs = info.get('final_observation', obs)
+            else:
+                obs, final_obs = env.reset()[0], obs
         else:
             final_obs = None
 
@@ -165,15 +182,34 @@ def _shutdown(conns, procs):
     for conn in conns:
         try:
             conn.send(('close', None))
-        except (EOFError, BrokenPipeError):
+        except (EOFError, BrokenPipeError, OSError):
             pass
 
     for proc in procs:
-        proc.join(timeout = 2)
+        if proc.is_alive():
+            proc.join(timeout = 0.2)
 
         if proc.is_alive():
             proc.terminate()
 
+    for conn in conns:
+        try:
+            conn.close()
+        except Exception:
+            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)]
+
+    assert len(actions) == num_envs, f'expected {num_envs} actions, but got {len(actions)}'
+    return actions
+
 # class
 
 class MultiprocessingVecEnv:
@@ -260,7 +296,7 @@ class MultiprocessingVecEnv:
         return _stack_trees([obs for obs, _ in results]), _merge_infos([info for _, info in results])
 
     def step(self, actions):
-        assert len(actions) == self.num_envs, f'expected {self.num_envs} actions, but got {len(actions)}'
+        actions = _split_actions(actions, self.num_envs)
 
         for conn, action in zip(self._conns, actions):
             _safe_send(conn, ('step', action))

{env_ssl_wrapper-0.3.0 → env_ssl_wrapper-0.4.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "env-ssl-wrapper"
-version = "0.3.0"
+version = "0.4.0"
 description = "One torch-native interface for any MDP environment"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }
@@ -39,7 +39,9 @@ Homepage = "https://pypi.org/project/env-ssl-wrapper/"
 Repository = "https://codeberg.org/lucidrains/env-ssl-wrapper"
 
 [project.optional-dependencies]
-examples = []
+examples = [
+    "fire",
+]
 test = [
     "pytest",
     "gymnasium>=0.29.1",
@@ -66,10 +68,6 @@ testpaths = [
 requires = ["hatchling"]
 build-backend = "hatchling.build"
 
-[tool.rye]
-managed = true
-dev-dependencies = []
-
 [tool.hatch.metadata]
 allow-direct-references = true
 

env_ssl_wrapper-0.3.0/env_ssl_wrapper/__init__.py

@@ -1,59 +0,0 @@
-from .standardize_wrapper import StandardizeWrapper
-from .image_wrapper import ImageObservationWrapper
-from .auto_batched_wrapper import AutoBatchedWrapper
-from .tensor_wrapper import TensorWrapper
-from .action_transform_wrapper import ActionTransformWrapper
-from .done_tracker_wrapper import DoneTrackerWrapper
-from .flatten_obs_wrapper import FlattenObsWrapper
-from .episode_padding_wrapper import EpisodePaddingWrapper
-from .time_limit_wrapper import TimeLimitWrapper
-from .vector import MultiprocessingVecEnv
-
-from .adapters import (
-    BaseEnvAdapter,
-    get_adapter,
-    register_adapter,
-    MujocoWarpAdapter,
-    IsaacAdapter,
-    PyBulletAdapter,
-    DMControlAdapter,
-    PufferLibAdapter,
-    RoboticsAdapter,
-    GymnasiumAdapter,
-    LegacyGymAdapter,
-    DefaultAdapter,
-)
-from .spaces import (
-    InferredSpace,
-    infer_observation_space,
-    space_from_action_spec,
-    action_space_dim,
-    action_space_is_discrete,
-    action_space_is_box,
-    action_space_bounds,
-    action_dim_of,
-    obs_dim_of,
-)
-
-from .mocks import (
-    MockEnv,
-    GymnasiumMockEnv,
-    GymnasiumDiscreteMockEnv,
-    LegacyGymMockEnv,
-    PyBulletMockEnv,
-    DMControlMockEnv,
-    IsaacMockEnv,
-    AutoresetVectorMockEnv,
-    PufferVectorMockEnv,
-    PufferTensorMockEnv,
-    ManiSkillMockEnv,
-    BraxMockEnv,
-    MjxMockEnv,
-    MetaWorldMockEnv,
-    TrifingerMockEnv,
-    HabitatMockEnv,
-    TupleObsMockEnv,
-    JaxArray
-)
-
-from .utils import wrap_env, compose_env

env_ssl_wrapper-0.3.0/env_ssl_wrapper/flatten_obs_wrapper.py

@@ -1,79 +0,0 @@
-from __future__ import annotations
-
-import numpy as np
-import torch
-from torch import is_tensor
-from torch.utils._pytree import tree_flatten
-from einops import rearrange
-
-from .helpers import EnvWrapper, is_array, is_scalar
-
-# helpers
-
-def flattenable(t):
-    if is_array(t) or is_scalar(t):
-        return True
-
-    try:
-        return np.asarray(t).dtype.kind not in 'USO'
-    except Exception:
-        return False
-
-def flatten_leaf(t):
-    if not is_array(t):
-        t = np.asarray(t)
-
-    if t.ndim == 0:
-        return rearrange(t, '-> 1 1')
-
-    if t.ndim == 1:
-        return rearrange(t, 'b -> b 1')
-
-    return rearrange(t, 'b ... -> b (...)')
-
-def concat_leaves(leaves):
-    axis, flatten = -1, True
-
-    if leaves[0].ndim == 1 and not all(len(t) == len(leaves[0]) for t in leaves):
-        axis, flatten = 0, False
-
-    leaves = [flatten_leaf(t) for t in leaves] if flatten else leaves
-
-    if is_tensor(leaves[0]):
-        return torch.cat(leaves, dim = axis)
-
-    return np.concatenate(leaves, axis = axis)
-
-# class
-
-class FlattenObsWrapper(EnvWrapper):
-    def __init__(self, env):
-        super().__init__(env)
-
-    def observation(self, obs):
-        if is_array(obs):
-            return obs
-
-        leaves, _ = tree_flatten(obs)
-        leaves = [t for t in leaves if flattenable(t)]
-
-        if len(leaves) == 0:
-            return obs
-
-        if len(leaves) == 1:
-            return flatten_leaf(leaves[0])
-
-        return concat_leaves(leaves)
-
-    def final_observation(self, info):
-        if isinstance(info, dict) and 'final_observation' in info:
-            info['final_observation'] = self.observation(info['final_observation'])
-        return info
-
-    def reset(self, **kwargs):
-        obs, info = self.env.reset(**kwargs)
-        return self.observation(obs), self.final_observation(info)
-
-    def step(self, action):
-        obs, reward, terminated, truncated, info = self.env.step(action)
-        return self.observation(obs), reward, terminated, truncated, self.final_observation(info)