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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: env-ssl-wrapper
-Version: 0.0.2
+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

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.2 → env_ssl_wrapper-0.0.3}/env_ssl_wrapper/image_wrapper.py

@@ -3,7 +3,6 @@ from __future__ import annotations
 import torch
 import numpy as np
 
-import gymnasium as gym
 from PIL import Image
 from einops import rearrange
 
@@ -14,7 +13,7 @@ def cast_tuple(t, length = 1):
 
 # class
 
-class ImageObservationWrapper(gym.ObservationWrapper):
+class ImageObservationWrapper:
     def __init__(
         self,
         env,
@@ -24,13 +23,18 @@ class ImageObservationWrapper(gym.ObservationWrapper):
         normalize = True,
         normalize_divisor = 255.0
     ):
-        super().__init__(env)
+        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)
@@ -56,3 +60,11 @@ class ImageObservationWrapper(gym.ObservationWrapper):
         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.2 → env_ssl_wrapper-0.0.3}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "env-ssl-wrapper"
-version = "0.0.2"
+version = "0.0.3"
 description = "An RL environment wrapper for learning SSL in the background"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

env_ssl_wrapper-0.0.2/env_ssl_wrapper/__init__.py

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

env_ssl_wrapper-0.0.2/env_ssl_wrapper/env_ssl_wrapper.py

@@ -1 +0,0 @@
-