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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: env-ssl-wrapper
-Version: 0.4.3
+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
@@ -167,6 +167,8 @@ Execution stops early the moment any env terminates or truncates mid-chunk — t
 
 Pass `gamma` (default `1.`) to discount intra-chunk rewards $r = \sum_{i=0}^{L-1} \gamma^i r_i$. The macro-transition discount factor to the next state is provided as `info['discount'] = gamma ** chunk_length`. `reward_mode` can also be `'mean'` or `'last'`.
 
+Set `reward_mode = 'chunk'` to get the raw per-step rewards back as a reward chunk of shape `(num_envs, chunk_length)` (`(chunk_length,)` for a single env) instead of a reduced scalar — useful when a learner wants to do per-step credit assignment within the chunk. The chunk is truncated to the number of substeps actually executed (`info['chunk_length']`) and is undiscounted, so `gamma` only affects the aggregated modes.
+
 Can also be passed directly to `StandardizeEnvWrapper` or `compose_env`:
 
 ```python

{env_ssl_wrapper-0.4.3 → env_ssl_wrapper-0.4.5}/README.md

@@ -104,6 +104,8 @@ Execution stops early the moment any env terminates or truncates mid-chunk — t
 
 Pass `gamma` (default `1.`) to discount intra-chunk rewards $r = \sum_{i=0}^{L-1} \gamma^i r_i$. The macro-transition discount factor to the next state is provided as `info['discount'] = gamma ** chunk_length`. `reward_mode` can also be `'mean'` or `'last'`.
 
+Set `reward_mode = 'chunk'` to get the raw per-step rewards back as a reward chunk of shape `(num_envs, chunk_length)` (`(chunk_length,)` for a single env) instead of a reduced scalar — useful when a learner wants to do per-step credit assignment within the chunk. The chunk is truncated to the number of substeps actually executed (`info['chunk_length']`) and is undiscounted, so `gamma` only affects the aggregated modes.
+
 Can also be passed directly to `StandardizeEnvWrapper` or `compose_env`:
 
 ```python

{env_ssl_wrapper-0.4.3 → 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,
@@ -45,7 +45,7 @@ class ActionChunkWrapper(EnvWrapper):
         gamma = default(discount, gamma)
         assert chunk_len >= 1, f'chunk_len must be at least 1, got {chunk_len}'
         assert 0. <= gamma <= 1., f'gamma must be between 0 and 1, got {gamma}'
-        assert reward_mode in ('sum', 'mean', 'last'), f'unknown reward_mode {reward_mode!r}'
+        assert reward_mode in ('sum', 'mean', 'last', 'chunk'), f'unknown reward_mode {reward_mode!r}'
 
         self.chunk_len = chunk_len
         self.gamma = float(gamma)
@@ -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
@@ -93,7 +93,9 @@ class ActionChunkWrapper(EnvWrapper):
         rewards = stack_steps(rewards)
         executed_len = rewards.shape[-1]
 
-        if self.reward_mode == 'last':
+        if self.reward_mode == 'chunk':
+            reward = rewards
+        elif self.reward_mode == 'last':
             reward = rewards[..., -1]
         elif self.reward_mode == 'mean':
             reward = reduce(rewards, '... k -> ...', 'mean')

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

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