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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: env-ssl-wrapper
-Version: 0.4.6
+Version: 0.5.0
 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
@@ -37,6 +37,8 @@ Requires-Python: >=3.10
 Requires-Dist: discrete-continuous-embed-readout
 Requires-Dist: einops>=0.8.1
 Requires-Dist: einx>=0.3.0
+Requires-Dist: imageio
+Requires-Dist: imageio-ffmpeg
 Requires-Dist: memmap-replay-buffer
 Requires-Dist: torch-einops-utils>=0.0.29
 Requires-Dist: torch>=2.5
@@ -161,7 +163,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.5.0}/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.5.0}/env_ssl_wrapper/__init__.py

@@ -6,6 +6,7 @@ from .standardize import *
 
 from .memory_trace import MemoryTraceWrapper
 from .action_chunk import ActionChunkWrapper
+from .evaluate import evaluate_actor, EpisodeStats
 
 # Wire backwards-compatibility aliases in sys.modules and module globals
 # so imports like `from env_ssl_wrapper.done_tracker_wrapper import DoneTrackerWrapper`
@@ -40,6 +41,8 @@ __all__ = [
     *standardize.__all__,
     'MemoryTraceWrapper',
     'ActionChunkWrapper',
+    'evaluate_actor',
+    'EpisodeStats',
 ]
 
 def __getattr__(name):

{env_ssl_wrapper-0.4.6 → env_ssl_wrapper-0.5.0}/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.5.0/env_ssl_wrapper/evaluate.py

@@ -0,0 +1,337 @@
+from __future__ import annotations
+
+from pathlib import Path
+from typing import NamedTuple
+
+import numpy as np
+import torch
+from torch import nn, is_tensor
+from torch.distributions import Categorical, Distribution
+from torch.utils._pytree import tree_flatten, tree_map
+
+from .standardize.helpers import (
+    EnvWrapper,
+    env_autoresets,
+    env_num_envs,
+    env_render,
+    exists,
+    first_existing,
+    get_adapter,
+    get_attr,
+    is_array,
+    is_vectorized,
+    normalize_reset_out,
+    normalize_step_out,
+    to_numpy,
+)
+
+# results
+
+class EpisodeStats(NamedTuple):
+    returns: torch.Tensor
+    lengths: torch.Tensor
+    videos: tuple[str, ...] = ()
+
+    @property
+    def num_episodes(self):
+        return self.returns.numel()
+
+    @property
+    def mean(self):
+        return self.returns.mean().item()
+
+    @property
+    def std(self):
+        return self.returns.std(unbiased = False).item()
+
+    @property
+    def min(self):
+        return self.returns.min().item()
+
+    @property
+    def max(self):
+        return self.returns.max().item()
+
+    @property
+    def video_path(self):
+        return self.videos[0] if self.videos else None
+
+    def success_rate(self, threshold):
+        return (self.returns >= threshold).float().mean().item()
+
+    def summary(self):
+        return dict(episodes = self.num_episodes, mean = self.mean, std = self.std, min = self.min, max = self.max)
+
+# actor output -> action
+
+def resolve_actor(actor):
+    # a callable `.dist` takes precedence, sb3 style
+
+    module = actor if isinstance(actor, nn.Module) else None
+    policy = get_attr(actor, 'dist')
+
+    if not callable(policy):
+        policy = actor
+
+    assert callable(policy), 'actor must be callable, or expose a callable `.dist`'
+
+    return policy, module
+
+def is_distribution(out):
+    if isinstance(out, Distribution):
+        return True
+
+    if is_tensor(out):
+        return False
+
+    return callable(get_attr(out, 'sample')) and (exists(get_attr(out, 'mean')) or exists(get_attr(out, 'probs')))
+
+def to_action(env, out, deterministic = True):
+    # distributions collapse to their mode, raw floating logits over a discrete space to their argmax
+
+    if is_distribution(out):
+        if not deterministic:
+            return out.sample()
+
+        return out.probs.argmax(dim = -1) if exists(get_attr(out, 'probs')) else out.mean
+
+    if is_tensor(out) and out.is_floating_point() and out.ndim > 0:
+        num_actions = get_attr(get_adapter(env).action_space, 'n')
+
+        if exists(num_actions) and out.shape[-1] == num_actions:
+            return out.argmax(dim = -1) if deterministic else Categorical(logits = out).sample()
+
+    return out
+
+# states and actions
+
+def to_actor_obs(obs, device, batched):
+    def convert(x):
+        if not (is_array(x) or callable(get_attr(x, '__array__'))):
+            return x
+
+        tensor = x if is_tensor(x) else torch.as_tensor(np.asarray(x))
+
+        if tensor.is_floating_point():
+            tensor = tensor.float()
+
+        tensor = tensor.to(device)
+
+        return tensor if batched or tensor.ndim == 0 else tensor[None]
+
+    return tree_map(convert, obs)
+
+def to_env_action(env, action):
+    if not is_tensor(action):
+        action = torch.as_tensor(action)
+
+    device = get_attr(env, 'device')
+    action = action.to(device) if exists(device) else action
+
+    # raw simulators take unbatched numpy
+
+    if not isinstance(env, EnvWrapper):
+        action = to_numpy(action)
+
+        if action.ndim > 1 and action.shape[0] == 1:
+            action = action[0]
+
+        return action.item() if action.ndim == 0 else action
+
+    # wrappers take batched torch actions - pad up to the declared rank
+
+    chunk_shape = get_attr(env, 'chunk_action_shape')
+
+    if exists(chunk_shape):
+        rank = len(chunk_shape) + 1
+    else:
+        space = first_existing(env, 'single_action_space', 'action_space')
+        rank = len(get_attr(space, 'shape', ())) + 1
+
+    while action.ndim < rank:
+        action = action[None]
+
+    return action
+
+
+# video
+
+def capture_frame(env):
+    render = get_attr(env, 'render')
+    frame = render() if callable(render) else None
+
+    if not exists(frame):
+        try:
+            frame = env_render(env, 256, 256)
+        except Exception:
+            return None
+
+    if is_tensor(frame):
+        frame = frame.detach().cpu().numpy()
+
+    frame = np.asarray(frame)
+
+    if frame.ndim == 4:
+        frame = frame[0]
+
+    if frame.ndim == 3 and frame.shape[0] in (3, 4) and frame.shape[-1] not in (3, 4):
+        frame = np.moveaxis(frame, 0, -1)
+
+    if frame.ndim == 3 and frame.shape[-1] == 4:
+        frame = frame[..., :3]
+
+    if frame.dtype != np.uint8:
+        scale = 255. if frame.max() <= 1. else 1.
+        frame = np.clip(frame * scale, 0., 255.).astype(np.uint8)
+
+    return frame
+
+def save_video(frames, path, fps = 30):
+    import imageio
+
+    path = Path(path)
+    path.parent.mkdir(parents = True, exist_ok = True)
+    imageio.mimsave(path, frames, fps = fps, macro_block_size = None)
+
+def video_path_at(path, index):
+    path = Path(path)
+    return str(path if index == 0 else path.with_name(f'{path.stem}_ep{index}{path.suffix}'))
+
+# evaluate
+
+@torch.no_grad()
+def evaluate_actor(
+    actor,
+    env,
+    episodes = 10,
+    max_steps = 1000,
+    seed = None,
+    goal = None,
+    device = None,
+    deterministic = True,
+    video_path = None,
+    video_episodes = 1,
+    fps = 30,
+    verbose = False
+):
+    # seed fixes the episode seeds, so checkpoints are compared on identical episodes
+    # video_path records the first video_episodes episodes
+
+    policy, module = resolve_actor(actor)
+
+    if not exists(device):
+        params = next(module.parameters(), None) if exists(module) else None
+        device = get_attr(params, 'device', torch.device('cpu'))
+
+    was_training = get_attr(module, 'training')
+
+    if exists(module):
+        module.eval()
+
+    if exists(goal):
+        goal = torch.as_tensor(goal) if not is_tensor(goal) else goal
+        goal = goal.to(device)
+        goal = goal[None] if goal.ndim == 1 else goal
+
+    num_envs = env_num_envs(env)
+    record = exists(video_path)
+
+    assert not (record and num_envs > 1), 'video recording requires a single env'
+
+    obs, info = normalize_reset_out(env.reset(seed = seed) if exists(seed) else env.reset())
+    batched = is_vectorized(env) or all(map(is_tensor, tree_flatten(obs)[0]))
+
+    goal = goal.expand(num_envs, -1) if exists(goal) and goal.shape[0] == 1 and num_envs > 1 else goal
+
+    returns = torch.zeros(num_envs)
+    lengths = torch.zeros(num_envs)
+    steps = torch.zeros(num_envs, dtype = torch.long)
+
+    done_returns, done_lengths, videos, frames = [], [], [], []
+
+    try:
+        while len(done_returns) < episodes:
+            if record and len(done_returns) < video_episodes:
+                frame = capture_frame(env)
+                if exists(frame):
+                    frames.append(frame)
+
+            state = to_actor_obs(obs, device, batched)
+            out = policy(state, goal) if exists(goal) else policy(state)
+            action = to_env_action(env, to_action(env, out, deterministic))
+
+            obs, reward, terminated, truncated, info = normalize_step_out(env.step(action))
+
+            # reward chunks collapse to one scalar per env
+
+            reward = torch.as_tensor(to_numpy(reward), dtype = torch.float32)
+
+            while reward.ndim > 1:
+                reward = reward.sum(dim = -1)
+
+            if reward.numel() > num_envs:
+                reward = reward.reshape(num_envs, -1).sum(dim = -1)
+
+            returns += reward if reward.numel() == num_envs else reward.expand(num_envs)
+
+            lengths += info.get('chunk_length', 1) if isinstance(info, dict) else 1
+            steps += 1
+
+            done = to_numpy(terminated | truncated).astype(bool).reshape(-1)
+
+            if exists(max_steps):
+                done |= to_numpy(steps >= max_steps).reshape(-1)
+
+            for i in np.where(done)[0]:
+                if len(done_returns) >= episodes:
+                    break
+
+                episode = len(done_returns)
+                done_returns.append(returns[i].item())
+                done_lengths.append(lengths[i].item())
+
+                if verbose:
+                    print(f'episode {episode:3d} | return {done_returns[-1]:8.1f} | length {done_lengths[-1]:7.0f}', flush = True)
+
+                if record and episode < video_episodes:
+                    frame = capture_frame(env)
+                    if exists(frame):
+                        frames.append(frame)
+
+                    if frames:
+                        path = video_path_at(video_path, episode)
+                        save_video(frames, path, fps)
+                        videos.append(path)
+
+                returns[i] = 0.
+                lengths[i] = 0.
+                steps[i] = 0
+
+            if done.any() and len(done_returns) < episodes:
+                if record:
+                    frames = []
+
+                # autoresetting envs reset themselves, everything else resets here
+
+                if not env_autoresets(env) and (num_envs == 1 or done.all() or bool(get_attr(env, 'needs_reset', False))):
+                    reset_seed = seed + len(done_returns) if exists(seed) and num_envs == 1 else None
+                    obs, info = normalize_reset_out(env.reset(seed = reset_seed) if exists(reset_seed) else env.reset())
+    finally:
+        if exists(module) and exists(was_training):
+            module.train(was_training)
+
+    return EpisodeStats(
+        returns = torch.tensor(done_returns, dtype = torch.float32),
+        lengths = torch.tensor(done_lengths, dtype = torch.float32),
+        videos = tuple(videos)
+    )
+
+if __name__ == '__main__':
+    import gymnasium as gym
+    from env_ssl_wrapper import StandardizeEnvWrapper
+
+    env = StandardizeEnvWrapper(gym.make('CartPole-v1'))
+
+    stats = evaluate_actor(lambda obs: torch.zeros(obs.shape[0], 2), env, episodes = 10, seed = 0, verbose = True)
+
+    print(stats.summary())

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

@@ -1,6 +1,6 @@
 [project]
 name = "env-ssl-wrapper"
-version = "0.4.6"
+version = "0.5.0"
 description = "One torch-native interface for any MDP environment"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }
@@ -27,6 +27,8 @@ dependencies = [
     "discrete-continuous-embed-readout",
     "einx>=0.3.0",
     "einops>=0.8.1",
+    "imageio",
+    "imageio-ffmpeg",
     "memmap-replay-buffer",
     "torch>=2.5",
     "torch-einops-utils>=0.0.29",