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

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

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

@@ -1,5 +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.3 → env_ssl_wrapper-0.0.4}/env_ssl_wrapper/auto_batched_wrapper.py

@@ -51,8 +51,8 @@ class AutoBatchedWrapper:
     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-0.0.4}/env_ssl_wrapper/tensor_wrapper.py

@@ -67,7 +67,7 @@ class TensorWrapper:
     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
 

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.3 → env_ssl_wrapper-0.0.4}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "env-ssl-wrapper"
-version = "0.0.3"
+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.3/env_ssl_wrapper/utils.py

@@ -1,30 +0,0 @@
-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