env-ssl-wrapper 0.4.4__tar.gz → 0.4.6__tar.gz

{env_ssl_wrapper-0.4.4 → env_ssl_wrapper-0.4.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: env-ssl-wrapper
-Version: 0.4.4
+Version: 0.4.6
 Summary: One torch-native interface for any MDP environment
 Project-URL: Homepage, https://pypi.org/project/env-ssl-wrapper/
 Project-URL: Repository, https://codeberg.org/lucidrains/env-ssl-wrapper

{env_ssl_wrapper-0.4.4 → env_ssl_wrapper-0.4.6}/env_ssl_wrapper/action_chunk.py

@@ -7,11 +7,12 @@ from torch import is_tensor
 
 from .standardize.helpers import (
     EnvWrapper,
-    exists,
+    copy_leaf,
     default,
     any_true,
     dones_of,
     get_attr,
+    is_vectorized,
 )
 
 # helpers
@@ -22,6 +23,16 @@ def stack_steps(steps):
 def chunk_steps(actions, axis):
     return actions.unbind(dim = axis) if is_tensor(actions) else np.moveaxis(actions, axis, 0)
 
+def unbatch_step(action):
+    # chunks always carry a num_envs axis - drop it for plain single envs
+    if is_tensor(action):
+        return action.squeeze(0)
+
+    if isinstance(action, np.ndarray):
+        return np.squeeze(action, axis = 0) if action.ndim > 0 else action
+
+    return action
+
 # wrapper
 
 class ActionChunkWrapper(EnvWrapper):
@@ -53,6 +64,7 @@ class ActionChunkWrapper(EnvWrapper):
 
         action_space = get_attr(env, 'action_space')
         self.action_shape = tuple(get_attr(action_space, 'shape', ()) or ())
+        self.expects_batch = is_vectorized(env)
 
     @property
     def chunk_action_shape(self):
@@ -82,8 +94,11 @@ class ActionChunkWrapper(EnvWrapper):
         out = None
 
         for action in chunk_steps(actions, chunk_axis):
+            if not self.expects_batch:
+                action = unbatch_step(action)
+
             out = self.env.step(action)
-            rewards.append(out[1])
+            rewards.append(copy_leaf(out[1]))
 
             if any_true(dones_of(out[2], out[3])):
                 break

{env_ssl_wrapper-0.4.4 → env_ssl_wrapper-0.4.6}/env_ssl_wrapper/standardize/auto_batched_wrapper.py

@@ -153,18 +153,14 @@ def is_numeric_container(x):
 def maybe_squeeze_dim(x, shape_tree = None, is_vector = False, prepend_batch = False):
     # reshape actions to match the env's space, falling back to heuristics
 
-    if isinstance(shape_tree, (list, dict)):
-        keyed = isinstance(shape_tree, dict)
-
-        ok = x.keys() == shape_tree.keys() if keyed else isinstance(x, (list, tuple)) and len(x) == len(shape_tree)
-        assert ok, f'action structure does not match its {"dict" if keyed else "tuple"} action space'
-
-        children = list(x.values()) if keyed else list(x)
-        subtrees = list(shape_tree.values()) if keyed else shape_tree
-
-        leaves = [maybe_squeeze_dim(child, subtree, is_vector, prepend_batch) for child, subtree in zip(children, subtrees)]
-
-        return dict(zip(x.keys(), leaves)) if keyed else rebuild_container(x, leaves)
+    if isinstance(shape_tree, dict):
+        assert isinstance(x, dict) and x.keys() == shape_tree.keys(), 'action structure does not match its dict action space'
+        return {key: maybe_squeeze_dim(child, shape_tree[key], is_vector, prepend_batch) for key, child in x.items()}
+
+    if isinstance(shape_tree, list):
+        assert isinstance(x, (list, tuple)) and len(x) == len(shape_tree), 'action structure does not match its tuple action space'
+        leaves = [maybe_squeeze_dim(child, subtree, is_vector, prepend_batch) for child, subtree in zip(x, shape_tree)]
+        return rebuild_container(x, leaves)
 
     # leaf-shaped tree — claims the whole input
 

{env_ssl_wrapper-0.4.4 → env_ssl_wrapper-0.4.6}/env_ssl_wrapper/standardize/helpers.py

@@ -84,6 +84,9 @@ def copy_leaf(x):
 
     return x
 
+def copy_tree(tree):
+    return tree_map(copy_leaf, tree)
+
 def dones_of(terminated, truncated):
     if not isinstance(terminated, (dict, list, tuple)):
         return terminated | truncated

{env_ssl_wrapper-0.4.4 → env_ssl_wrapper-0.4.6}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "env-ssl-wrapper"
-version = "0.4.4"
+version = "0.4.6"
 description = "One torch-native interface for any MDP environment"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }