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

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

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

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

env_ssl_wrapper-0.0.4/env_ssl_wrapper/action_transform_wrapper.py

@@ -0,0 +1,93 @@
+from __future__ import annotations
+
+import numpy as np
+import torch
+from torch import is_tensor
+from torch.utils._pytree import tree_map
+
+# helpers
+
+def exists(v):
+    return v is not None
+
+def default(v, d):
+    return v if exists(v) else d
+
+def rescale(
+    t,
+    from_range: tuple[float, float],
+    to_range: tuple[float, float]
+):
+    from_min, from_max = from_range
+    to_min, to_max = to_range
+    return (t - from_min) / (from_max - from_min) * (to_max - to_min) + to_min
+
+# classes
+
+class ActionTransformWrapper:
+    def __init__(
+        self,
+        env,
+        transforms = None,
+        clip = None
+    ):
+        self.env = env
+
+        if isinstance(transforms, dict):
+            transforms = [transforms]
+
+        self.transforms = default(transforms, [])
+        self.clip = clip
+
+    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):
+        return self.env.reset(**kwargs)
+
+    def step(self, action):
+        def transform_action(t):
+            is_torch_float = is_tensor(t) and t.is_floating_point()
+            is_np_float = isinstance(t, np.ndarray) and np.issubdtype(t.dtype, np.floating)
+            is_scalar_float = isinstance(t, float)
+
+            if not (is_torch_float or is_np_float or is_scalar_float):
+                return t
+
+            if is_tensor(t):
+                t = t.clone()
+            elif isinstance(t, np.ndarray):
+                t = np.copy(t)
+
+            for ind, transform in enumerate(self.transforms):
+                indices = transform.get('indices')
+                rescale_from_to = transform.get('rescale_from_to')
+
+                if not exists(indices) and len(self.transforms) > 1:
+                    indices = ind
+
+                if exists(rescale_from_to):
+                    from_range, to_range = rescale_from_to
+
+                    if not exists(indices):
+                        t = rescale(t, from_range, to_range)
+                    else:
+                        part = t[..., indices]
+                        t[..., indices] = rescale(part, from_range, to_range)
+
+            if exists(self.clip):
+                min_clip, max_clip = self.clip
+
+                if is_tensor(t):
+                    t = torch.clamp(t, min_clip, max_clip)
+                elif isinstance(t, np.ndarray):
+                    t = np.clip(t, min_clip, max_clip)
+                else:
+                    t = max(min_clip, min(t, max_clip))
+
+            return t
+
+        action = tree_map(transform_action, action)
+        return self.env.step(action)

env_ssl_wrapper-0.0.4/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.4}/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.4/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.4/env_ssl_wrapper/utils.py

@@ -0,0 +1,45 @@
+from __future__ import annotations
+from functools import partial
+
+from .image_wrapper import ImageObservationWrapper
+from .auto_batched_wrapper import AutoBatchedWrapper
+from .tensor_wrapper import TensorWrapper
+from .action_transform_wrapper import ActionTransformWrapper
+
+WRAPPERS = dict(
+    image = ImageObservationWrapper,
+    auto_batch = AutoBatchedWrapper,
+    tensor = TensorWrapper,
+    action_transform = ActionTransformWrapper
+)
+
+def is_unique(arr):
+    return len(set(arr)) == len(arr)
+
+def compose_env(env, *wrappers):
+    funcs = []
+    classes = []
+
+    for wrapper in wrappers:
+        if isinstance(wrapper, str):
+            wrapper = WRAPPERS[wrapper]
+
+        if isinstance(wrapper, tuple):
+            name, kwargs = wrapper
+            wrapper = partial(WRAPPERS.get(name, name), **kwargs)
+
+        cls = wrapper.func if isinstance(wrapper, partial) else wrapper
+
+        funcs.append(wrapper)
+        classes.append(cls)
+
+    assert is_unique(classes), 'duplicate wrappers found'
+
+    for func in funcs:
+        env = func(env)
+
+    return env
+
+# alias
+
+wrap_env = compose_env

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

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