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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: env-ssl-wrapper
-Version: 0.4.4
+Version: 0.4.5
 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.5}/env_ssl_wrapper/action_chunk.py

@@ -7,7 +7,7 @@ from torch import is_tensor
 
 from .standardize.helpers import (
     EnvWrapper,
-    exists,
+    copy_leaf,
     default,
     any_true,
     dones_of,
@@ -83,7 +83,7 @@ class ActionChunkWrapper(EnvWrapper):
 
         for action in chunk_steps(actions, chunk_axis):
             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.5}/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.5}/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.5}/pyproject.toml

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