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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: env-ssl-wrapper
-Version: 0.3.0
+Version: 0.4.1
 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.1}/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.1/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.1/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.1/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.1/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.1/env_ssl_wrapper/standardize}/auto_batched_wrapper.py

@@ -7,7 +7,18 @@ import torch
 from torch.utils._pytree import tree_map
 from einops import rearrange
 
-from .helpers import EnvWrapper, default, exists, first_existing, get_attr, is_array, is_scalar, is_tensor, is_vectorized
+from .helpers import (
+    FINAL_OBSERVATION_KEYS,
+    EnvWrapper,
+    default,
+    exists,
+    first_existing,
+    get_attr,
+    is_array,
+    is_scalar,
+    is_tensor,
+    is_vectorized,
+)
 from .spaces import space_from_action_spec
 
 # helper functions
@@ -39,6 +50,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):
@@ -196,7 +213,9 @@ class AutoBatchedWrapper(EnvWrapper):
 
         obs, reward, terminated, truncated, info = *maybe_expand_dim(out[:4]), out[4]
 
-        if isinstance(info, dict) and 'final_observation' in info:
-            info['final_observation'] = maybe_expand_dim(info['final_observation'])
+        if isinstance(info, dict):
+            for key in FINAL_OBSERVATION_KEYS:
+                if key in info:
+                    info[key] = maybe_expand_dim(info[key])
 
         return obs, reward, terminated, truncated, info

{env_ssl_wrapper-0.3.0/env_ssl_wrapper → env_ssl_wrapper-0.4.1/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