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

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

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

@@ -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)
@@ -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.4}/pyproject.toml

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