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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: env-ssl-wrapper
-Version: 0.4.6
+Version: 0.4.7
 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
@@ -161,7 +161,7 @@ while True:
         obs, info = env.reset()
 ```
 
-Chunks are shaped `(num_envs, chunk_len, *action_shape)` (`(num_envs, chunk_len)` for discrete actions).
+Chunks are shaped `(num_envs, chunk_len, *action_shape)` (`(num_envs, chunk_len)` for discrete actions). The wrapper can also be applied directly to raw gymnasium environments, single or vectorized — torch action chunks are converted to numpy for the underlying simulator.
 
 Execution stops early the moment any env terminates or truncates mid-chunk — the terminal state and done flags of that substep are returned and the rest of the chunk is dropped, so a new episode is never silently advanced.
 

{env_ssl_wrapper-0.4.6 → env_ssl_wrapper-0.4.7}/README.md

@@ -98,7 +98,7 @@ while True:
         obs, info = env.reset()
 ```
 
-Chunks are shaped `(num_envs, chunk_len, *action_shape)` (`(num_envs, chunk_len)` for discrete actions).
+Chunks are shaped `(num_envs, chunk_len, *action_shape)` (`(num_envs, chunk_len)` for discrete actions). The wrapper can also be applied directly to raw gymnasium environments, single or vectorized — torch action chunks are converted to numpy for the underlying simulator.
 
 Execution stops early the moment any env terminates or truncates mid-chunk — the terminal state and done flags of that substep are returned and the rest of the chunk is dropped, so a new episode is never silently advanced.
 

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

@@ -11,8 +11,10 @@ from .standardize.helpers import (
     default,
     any_true,
     dones_of,
+    first_existing,
     get_attr,
     is_vectorized,
+    to_numpy,
 )
 
 # helpers
@@ -62,9 +64,13 @@ class ActionChunkWrapper(EnvWrapper):
         self.gamma = float(gamma)
         self.reward_mode = reward_mode
 
-        action_space = get_attr(env, 'action_space')
+        # vectorized gym envs expose a batched action space - the single env space defines the chunk layout
+
+        action_space = first_existing(env, 'single_action_space', 'action_space')
         self.action_shape = tuple(get_attr(action_space, 'shape', ()) or ())
+
         self.expects_batch = is_vectorized(env)
+        self.raw_env = not isinstance(env, EnvWrapper)
 
     @property
     def chunk_action_shape(self):
@@ -80,6 +86,11 @@ class ActionChunkWrapper(EnvWrapper):
         if not (is_tensor(actions) or isinstance(actions, np.ndarray)):
             actions = np.asarray(actions)
 
+        # raw simulators only understand numpy - wrapper chains from this package handle torch actions
+
+        if self.raw_env and is_tensor(actions):
+            actions = to_numpy(actions)
+
         chunk_axis = actions.ndim - 1 - len(self.action_shape)
 
         assert chunk_axis >= 1, (

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

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