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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: env-ssl-wrapper
-Version: 0.0.1
+Version: 0.0.2
 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.2/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.2/env_ssl_wrapper/__init__.py

@@ -0,0 +1 @@
+from .image_wrapper import ImageObservationWrapper

env_ssl_wrapper-0.0.2/env_ssl_wrapper/env_ssl_wrapper.py

@@ -0,0 +1 @@
+

env_ssl_wrapper-0.0.2/env_ssl_wrapper/image_wrapper.py

@@ -0,0 +1,58 @@
+from __future__ import annotations
+
+import torch
+import numpy as np
+
+import gymnasium as gym
+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(gym.ObservationWrapper):
+    def __init__(
+        self,
+        env,
+        image_size = (64, 64),
+        image_key = 'image',
+        resample_method = Image.BILINEAR,
+        normalize = True,
+        normalize_divisor = 255.0
+    ):
+        super().__init__(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 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

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

@@ -1,6 +1,6 @@
 [project]
 name = "env-ssl-wrapper"
-version = "0.0.1"
+version = "0.0.2"
 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