env-ssl-wrapper 0.0.1__tar.gz → 0.0.3__tar.gz

{env_ssl_wrapper-0.0.1 → env_ssl_wrapper-0.0.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: env-ssl-wrapper
-Version: 0.0.1
+Version: 0.0.3
 Summary: An RL environment wrapper for learning SSL in the background
 Project-URL: Homepage, https://pypi.org/project/env-ssl-wrapper/
 Project-URL: Repository, https://codeberg.org/lucidrains/env-ssl-wrapper
@@ -34,10 +34,14 @@ Classifier: License :: OSI Approved :: MIT License
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Requires-Python: >=3.10
+Requires-Dist: discrete-continuous-embed-readout
 Requires-Dist: einops>=0.8.1
 Requires-Dist: einx>=0.3.0
+Requires-Dist: memmap-replay-buffer
 Requires-Dist: torch-einops-utils>=0.0.29
 Requires-Dist: torch>=2.5
+Requires-Dist: x-mlps-pytorch
+Requires-Dist: x-transformers
 Provides-Extra: examples
 Provides-Extra: test
 Requires-Dist: pytest; extra == 'test'
@@ -46,3 +50,63 @@ Description-Content-Type: text/markdown
 ## env-ssl-wrapper (wip)
 
 Wrappers around environments that will take care of providing representations from self supervised learning automagically
+
+## Citations
+
+```bibtex
+@misc{schwarzer2021dataefficientreinforcementlearningselfpredictive,
+    title   = {Data-Efficient Reinforcement Learning with Self-Predictive Representations},
+    author  = {Max Schwarzer and Ankesh Anand and Rishab Goel and R Devon Hjelm and Aaron Courville and Philip Bachman},
+    year    = {2021},
+    eprint  = {2007.05929},
+    archivePrefix = {arXiv},
+    primaryClass = {cs.LG},
+    url     = {https://arxiv.org/abs/2007.05929},
+}
+```
+
+```bibtex
+@misc{schmidt2024learningactactions,
+    title   = {Learning to Act without Actions},
+    author  = {Dominik Schmidt and Minqi Jiang},
+    year    = {2024},
+    eprint  = {2312.10812},
+    archivePrefix = {arXiv},
+    primaryClass = {cs.LG},
+    url     = {https://arxiv.org/abs/2312.10812},
+}
+```
+
+```bibtex
+@misc{eysenbach2023contrastivelearninggoalconditionedreinforcement,
+    title   = {Contrastive Learning as Goal-Conditioned Reinforcement Learning},
+    author  = {Benjamin Eysenbach and Tianjun Zhang and Ruslan Salakhutdinov and Sergey Levine},
+    year    = {2023},
+    eprint  = {2206.07568},
+    archivePrefix = {arXiv},
+    primaryClass = {cs.LG},
+    url     = {https://arxiv.org/abs/2206.07568},
+}
+```
+
+```bibtex
+@misc{ashlag2025stateentropyregularizationrobust,
+    title   = {State Entropy Regularization for Robust Reinforcement Learning},
+    author  = {Yonatan Ashlag and Uri Koren and Mirco Mutti and Esther Derman and Pierre-Luc Bacon and Shie Mannor},
+    year    = {2025},
+    eprint  = {2506.07085},
+    archivePrefix = {arXiv},
+    primaryClass = {cs.LG},
+    url     = {https://arxiv.org/abs/2506.07085},
+}
+```
+
+```bibtex
+@inproceedings{park2026dual,
+    title   = {Dual Goal Representations},
+    author  = {Seohong Park and Deepinder Mann and Sergey Levine},
+    booktitle = {The Fourteenth International Conference on Learning Representations},
+    year    = {2026},
+    url     = {https://openreview.net/forum?id=aMKFTidLSM}
+}
+```

env_ssl_wrapper-0.0.3/README.md

@@ -0,0 +1,63 @@
+## env-ssl-wrapper (wip)
+
+Wrappers around environments that will take care of providing representations from self supervised learning automagically
+
+## Citations
+
+```bibtex
+@misc{schwarzer2021dataefficientreinforcementlearningselfpredictive,
+    title   = {Data-Efficient Reinforcement Learning with Self-Predictive Representations},
+    author  = {Max Schwarzer and Ankesh Anand and Rishab Goel and R Devon Hjelm and Aaron Courville and Philip Bachman},
+    year    = {2021},
+    eprint  = {2007.05929},
+    archivePrefix = {arXiv},
+    primaryClass = {cs.LG},
+    url     = {https://arxiv.org/abs/2007.05929},
+}
+```
+
+```bibtex
+@misc{schmidt2024learningactactions,
+    title   = {Learning to Act without Actions},
+    author  = {Dominik Schmidt and Minqi Jiang},
+    year    = {2024},
+    eprint  = {2312.10812},
+    archivePrefix = {arXiv},
+    primaryClass = {cs.LG},
+    url     = {https://arxiv.org/abs/2312.10812},
+}
+```
+
+```bibtex
+@misc{eysenbach2023contrastivelearninggoalconditionedreinforcement,
+    title   = {Contrastive Learning as Goal-Conditioned Reinforcement Learning},
+    author  = {Benjamin Eysenbach and Tianjun Zhang and Ruslan Salakhutdinov and Sergey Levine},
+    year    = {2023},
+    eprint  = {2206.07568},
+    archivePrefix = {arXiv},
+    primaryClass = {cs.LG},
+    url     = {https://arxiv.org/abs/2206.07568},
+}
+```
+
+```bibtex
+@misc{ashlag2025stateentropyregularizationrobust,
+    title   = {State Entropy Regularization for Robust Reinforcement Learning},
+    author  = {Yonatan Ashlag and Uri Koren and Mirco Mutti and Esther Derman and Pierre-Luc Bacon and Shie Mannor},
+    year    = {2025},
+    eprint  = {2506.07085},
+    archivePrefix = {arXiv},
+    primaryClass = {cs.LG},
+    url     = {https://arxiv.org/abs/2506.07085},
+}
+```
+
+```bibtex
+@inproceedings{park2026dual,
+    title   = {Dual Goal Representations},
+    author  = {Seohong Park and Deepinder Mann and Sergey Levine},
+    booktitle = {The Fourteenth International Conference on Learning Representations},
+    year    = {2026},
+    url     = {https://openreview.net/forum?id=aMKFTidLSM}
+}
+```

env_ssl_wrapper-0.0.3/env_ssl_wrapper/__init__.py

@@ -0,0 +1,5 @@
+from .image_wrapper import ImageObservationWrapper
+from .auto_batched_wrapper import AutoBatchedWrapper
+from .tensor_wrapper import TensorWrapper
+
+from .utils import wrap_env, compose_env

env_ssl_wrapper-0.0.3/env_ssl_wrapper/auto_batched_wrapper.py

@@ -0,0 +1,58 @@
+from __future__ import annotations
+
+import numpy as np
+from torch import is_tensor
+from torch.utils._pytree import tree_map
+from einops import rearrange
+
+# helper functions
+
+def is_vectorized(env) -> bool:
+    if hasattr(env, 'num_envs') and getattr(env, 'num_envs', 0) > 0:
+        return True
+    try:
+        from gymnasium.vector import VectorEnv
+        return isinstance(getattr(env, 'unwrapped', env), VectorEnv)
+    except ImportError:
+        return False
+
+def maybe_expand_dim(x):
+    def _expand(t):
+        if isinstance(t, np.ndarray) or is_tensor(t):
+            return rearrange(t, '... -> 1 ...')
+        if isinstance(t, (int, float, bool, np.number, np.bool_)):
+            return np.array([t])
+        return t
+    return tree_map(_expand, x)
+
+def maybe_squeeze_dim(x):
+    def _squeeze(t):
+        if isinstance(t, np.ndarray) or is_tensor(t):
+            return rearrange(t, '1 ... -> ...')
+        return t
+    return tree_map(_squeeze, x)
+
+# classes
+
+class AutoBatchedWrapper:
+    def __init__(self, env, is_vector: bool | None = None):
+        self.env = env
+        self.is_vector = is_vector if is_vector is not None else is_vectorized(env)
+
+    def __getattr__(self, name):
+        if name.startswith('_'):
+            raise AttributeError(f"attempted to get missing private attribute '{name}'")
+        return getattr(self.env, name)
+
+    def reset(self, **kwargs):
+        obs, info = self.env.reset(**kwargs)
+        return (maybe_expand_dim(obs), info) if not self.is_vector else (obs, info)
+
+    def step(self, action):
+        action = maybe_squeeze_dim(action) if not self.is_vector else action
+        out = self.env.step(action)
+        
+        if self.is_vector:
+            return out
+            
+        return *maybe_expand_dim(out[:4]), out[4]

env_ssl_wrapper-0.0.3/env_ssl_wrapper/image_wrapper.py

@@ -0,0 +1,70 @@
+from __future__ import annotations
+
+import torch
+import numpy as np
+
+from PIL import Image
+from einops import rearrange
+
+# functions
+
+def cast_tuple(t, length = 1):
+    return t if isinstance(t, tuple) else ((t,) * length)
+
+# class
+
+class ImageObservationWrapper:
+    def __init__(
+        self,
+        env,
+        image_size = (64, 64),
+        image_key = 'image',
+        resample_method = Image.BILINEAR,
+        normalize = True,
+        normalize_divisor = 255.0
+    ):
+        self.env = env
+        self.image_size = cast_tuple(image_size, 2)
+        self.image_key = image_key
+        self.resample_method = resample_method
+        self.normalize = normalize
+        self.normalize_divisor = normalize_divisor
+
+    def __getattr__(self, name):
+        if name.startswith('_'):
+            raise AttributeError(f"attempted to get missing private attribute '{name}'")
+        return getattr(self.env, name)
+
+    def render_frame(self):
+        img = self.env.render()
+        img = Image.fromarray(img).resize(self.image_size, resample = self.resample_method)
+        img_tensor = torch.from_numpy(np.array(img))
+        img = rearrange(img_tensor, 'h w c -> 1 c h w')
+
+        if self.normalize:
+            img = img.float() / self.normalize_divisor
+
+        return img
+
+    def observation(self, obs):
+        img_tensor = self.render_frame()
+        img_tensor = rearrange(img_tensor, '1 c h w -> c h w')
+
+        if not isinstance(obs, dict):
+            return dict(state = obs, **{self.image_key: img_tensor})
+
+        if self.image_key in obs:
+            raise ValueError(f"Key '{self.image_key}' is already present in the observation dictionary.")
+
+        obs = dict(obs)
+        obs.update({self.image_key: img_tensor})
+
+        return obs
+
+    def reset(self, **kwargs):
+        obs, info = self.env.reset(**kwargs)
+        return self.observation(obs), info
+
+    def step(self, action):
+        obs, reward, terminated, truncated, info = self.env.step(action)
+        return self.observation(obs), reward, terminated, truncated, info

env_ssl_wrapper-0.0.3/env_ssl_wrapper/tensor_wrapper.py

@@ -0,0 +1,74 @@
+from __future__ import annotations
+
+import numpy as np
+from torch import tensor, is_tensor, from_numpy, float64, device as torch_device
+from torch.utils._pytree import tree_map
+
+# helper functions
+
+def numpy_to_torch(x, device, cast_float64_to_float32 = False):
+    def _to_torch(t):
+        if isinstance(t, np.ndarray):
+            t = from_numpy(t)
+        elif isinstance(t, (int, float, bool, np.number, np.bool_)):
+            t = tensor(t)
+
+        if not is_tensor(t):
+            return t
+
+        if cast_float64_to_float32 and t.dtype == float64:
+            t = t.float()
+
+        return t.to(device)
+    return tree_map(_to_torch, x)
+
+def torch_to_numpy(x, cast_float64_to_float32 = False):
+    def _to_numpy(t):
+        if is_tensor(t):
+            t = t.detach().cpu().numpy()
+        elif isinstance(t, (int, float, bool, np.number, np.bool_)):
+            t = np.array(t)
+
+        if not isinstance(t, np.ndarray):
+            return t
+
+        if cast_float64_to_float32 and t.dtype == np.float64:
+            t = t.astype(np.float32)
+
+        return t
+    return tree_map(_to_numpy, x)
+
+# classes
+
+class TensorWrapper:
+    def __init__(
+        self,
+        env,
+        device: str | torch_device = 'cpu',
+        convert_in: bool = True,
+        convert_out: bool = True,
+        cast_float64_to_float32: bool = False
+    ):
+        self.env = env
+        self.device = torch_device(device)
+        self.convert_in = convert_in
+        self.convert_out = convert_out
+        self.cast_float64_to_float32 = cast_float64_to_float32
+
+    def __getattr__(self, name):
+        if name.startswith('_'):
+            raise AttributeError(f"attempted to get missing private attribute '{name}'")
+        return getattr(self.env, name)
+
+    def reset(self, **kwargs):
+        obs, info = self.env.reset(**kwargs)
+        return (numpy_to_torch(obs, self.device, self.cast_float64_to_float32), info) if self.convert_out else (obs, info)
+
+    def step(self, action):
+        action = torch_to_numpy(action, self.cast_float64_to_float32) if self.convert_in else action
+        out = self.env.step(action)
+        
+        if not self.convert_out:
+            return out
+
+        return *numpy_to_torch(out[:4], self.device, self.cast_float64_to_float32), out[4]

env_ssl_wrapper-0.0.3/env_ssl_wrapper/utils.py

@@ -0,0 +1,30 @@
+from __future__ import annotations
+from functools import partial
+
+from .image_wrapper import ImageObservationWrapper
+from .auto_batched_wrapper import AutoBatchedWrapper
+from .tensor_wrapper import TensorWrapper
+
+WRAPPERS = dict(
+    image = ImageObservationWrapper,
+    auto_batch = AutoBatchedWrapper,
+    tensor = TensorWrapper
+)
+
+def compose_env(env, *wrappers):
+    for wrapper in wrappers:
+        if isinstance(wrapper, str):
+            wrapper = WRAPPERS[wrapper]
+
+        if isinstance(wrapper, tuple):
+            name_or_fn, kwargs = wrapper
+            fn = WRAPPERS.get(name_or_fn, name_or_fn)
+            wrapper = partial(fn, **kwargs)
+            
+        env = wrapper(env)
+        
+    return env
+
+# alias
+
+wrap_env = compose_env

{env_ssl_wrapper-0.0.1 → env_ssl_wrapper-0.0.3}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "env-ssl-wrapper"
-version = "0.0.1"
+version = "0.0.3"
 description = "An RL environment wrapper for learning SSL in the background"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }
@@ -24,10 +24,14 @@ classifiers=[
 ]
 
 dependencies = [
+    "discrete-continuous-embed-readout",
     "einx>=0.3.0",
     "einops>=0.8.1",
+    "memmap-replay-buffer",
     "torch>=2.5",
     "torch-einops-utils>=0.0.29",
+    "x-transformers",
+    "x-mlps-pytorch",
 ]
 
 [project.urls]

env_ssl_wrapper-0.0.1/README.md

@@ -1,3 +0,0 @@
-## env-ssl-wrapper (wip)
-
-Wrappers around environments that will take care of providing representations from self supervised learning automagically