cusrl 1.0.0__py3-none-any.whl

cusrl/__init__.py

@@ -0,0 +1,107 @@
+from cusrl import environment, hook, logger, module, preset, sampler, template, utils, zoo
+from cusrl.environment import make_gym_env, make_gym_vec, make_isaaclab_env
+from cusrl.module import (
+    CNN,
+    MLP,
+    RNN,
+    Actor,
+    AdaptiveNormalDist,
+    Denormalization,
+    Distribution,
+    DistributionFactoryLike,
+    FeedForward,
+    InferenceModule,
+    LayerFactoryLike,
+    Module,
+    ModuleFactory,
+    ModuleFactoryLike,
+    MultiheadSelfAttention,
+    NormalDist,
+    Normalization,
+    OneHotCategoricalDist,
+    Sequential,
+    Simba,
+    TransformerEncoderLayer,
+    Value,
+)
+from cusrl.sampler import (
+    AutoMiniBatchSampler,
+    MiniBatchSampler,
+    TemporalMiniBatchSampler,
+)
+from cusrl.template import (
+    ActorCritic,
+    Agent,
+    Buffer,
+    Environment,
+    EnvironmentSpec,
+    Hook,
+    Logger,
+    LoggerFactory,
+    LoggerFactoryLike,
+    OptimizerFactory,
+    Player,
+    Sampler,
+    Trainer,
+    Trial,
+)
+from cusrl.utils import (
+    device,
+    set_global_seed,
+)
+
+__all__ = [
+    "hook",
+    "logger",
+    "module",
+    "preset",
+    "sampler",
+    "template",
+    "utils",
+    "zoo",
+    "environment",
+    "Actor",
+    "ActorCritic",
+    "AdaptiveNormalDist",
+    "Agent",
+    "AutoMiniBatchSampler",
+    "Buffer",
+    "CNN",
+    "Denormalization",
+    "Distribution",
+    "DistributionFactoryLike",
+    "Environment",
+    "EnvironmentSpec",
+    "FeedForward",
+    "Hook",
+    "InferenceModule",
+    "LayerFactoryLike",
+    "Logger",
+    "LoggerFactory",
+    "LoggerFactoryLike",
+    "MLP",
+    "MiniBatchSampler",
+    "Module",
+    "ModuleFactory",
+    "ModuleFactoryLike",
+    "MultiheadSelfAttention",
+    "NormalDist",
+    "Normalization",
+    "OneHotCategoricalDist",
+    "OptimizerFactory",
+    "Player",
+    "RNN",
+    "Sampler",
+    "Sequential",
+    "Simba",
+    "TemporalMiniBatchSampler",
+    "Trainer",
+    "TransformerEncoderLayer",
+    "Trial",
+    "Value",
+    "device",
+    "make_gym_env",
+    "make_gym_vec",
+    "make_isaaclab_env",
+    "set_global_seed",
+]

cusrl/environment/__init__.py

@@ -0,0 +1,11 @@
+from .gym import GymEnvAdapter, GymVectorEnvAdapter, make_gym_env, make_gym_vec
+from .isaaclab import IsaacLabEnvAdapter, make_isaaclab_env
+
+__all__ = [
+    "GymEnvAdapter",
+    "GymVectorEnvAdapter",
+    "IsaacLabEnvAdapter",
+    "make_gym_env",
+    "make_gym_vec",
+    "make_isaaclab_env",
+]

cusrl/environment/gym.py

@@ -0,0 +1,157 @@
+import random
+import warnings
+from collections.abc import Callable, Sequence
+from typing import Any
+
+import gymnasium as gym
+import numpy as np
+from gymnasium.envs.registration import EnvSpec
+
+import cusrl.utils
+from cusrl.template import Environment
+from cusrl.utils.typing import Array, Slice
+
+__all__ = ["GymEnvAdapter", "GymVectorEnvAdapter", "make_gym_env", "make_gym_vec"]
+
+
+class GymEnvAdapter(Environment):
+    def __init__(self, wrapped: gym.Env):
+        if not isinstance(wrapped.observation_space, gym.spaces.Box):
+            raise ValueError("Only Box observation space is supported.")
+        if not len(wrapped.observation_space.shape) == 1:
+            raise ValueError("Only 1D observation space is supported.")
+        if isinstance(wrapped.action_space, gym.spaces.Box):
+            if not len(wrapped.action_space.shape) == 1:
+                raise ValueError("For Box action space, only 1D action space is supported.")
+            action_dim = wrapped.action_space.shape[0]
+        elif isinstance(wrapped.action_space, gym.spaces.Discrete):
+            action_dim = int(wrapped.action_space.n)
+        else:
+            raise ValueError(f"Unsupported action space type: {wrapped.action_space}.")
+
+        super().__init__(
+            num_instances=1,
+            observation_dim=wrapped.observation_space.shape[0],
+            action_dim=action_dim,
+            spec={"gym_spec": wrapped.spec, "gym_metadata": wrapped.metadata},
+        )
+        wrapped.reset(seed=random.getrandbits(4))
+        self.wrapped = wrapped
+
+    def reset(self, *, indices: Array | Slice | None = None):
+        observation, info = self.wrapped.reset()
+        observation = observation.reshape(1, -1)
+        if self.wrapped.render_mode is not None:
+            self.wrapped.render()
+        # TODO: process arrays in info
+        return observation, None, info
+
+    def step(self, action: Array):
+        if isinstance(self.wrapped.action_space, gym.spaces.Discrete):
+            action = np.argmax(action, axis=-1)
+        action = action.squeeze(0)
+        observation, reward, terminated, truncated, info = self.wrapped.step(action)
+        observation = observation.reshape(1, -1)
+        reward = np.array([[reward]], dtype=np.float32)
+        terminated = np.array([[terminated]])
+        truncated = np.array([[truncated]])
+        if self.wrapped.render_mode is not None:
+            self.wrapped.render()
+        # TODO: process arrays in info
+        return observation, None, reward, terminated, truncated, info
+
+
+class GymVectorEnvAdapter(Environment):
+    def __init__(self, wrapped: gym.vector.VectorEnv):
+        if not isinstance(wrapped.single_observation_space, gym.spaces.Box):
+            raise ValueError("Only Box observation space is supported.")
+        if not len(wrapped.single_observation_space.shape) == 1:
+            raise ValueError("Only 1D observation space is supported.")
+        if isinstance(wrapped.single_action_space, gym.spaces.Box):
+            if not len(wrapped.single_action_space.shape) == 1:
+                raise ValueError("For Box action space, only 1D action space is supported.")
+            action_dim = wrapped.single_action_space.shape[0]
+        elif isinstance(wrapped.single_action_space, gym.spaces.Discrete):
+            action_dim = int(wrapped.single_action_space.n)
+        else:
+            raise ValueError(f"Unsupported action space type: {wrapped.single_action_space}.")
+
+        if (autoreset_mode := wrapped.metadata.get("autoreset_mode")) is None:
+            if cusrl.utils.is_main_process():
+                warnings.warn("GymVectorEnvAdapter: make sure 'autoreset_mode' is 'DISABLED'.")
+        elif autoreset_mode != gym.vector.AutoresetMode.DISABLED:
+            raise ValueError("'autoreset_mode' of vector environments must be 'DISABLED'.")
+
+        super().__init__(
+            num_instances=wrapped.num_envs,
+            observation_dim=wrapped.single_observation_space.shape[0],
+            action_dim=action_dim,
+            spec={"gym_spec": wrapped.spec, "gym_metadata": wrapped.metadata},
+        )
+        wrapped.reset(seed=random.getrandbits(4))
+        self.wrapped = wrapped
+
+    def reset(self, *, indices: Array | Slice | None = None):
+        if indices is None:
+            observation, info = self.wrapped.reset()
+            return observation, None, info
+        if not isinstance(indices, np.ndarray):
+            mask = np.zeros(self.num_instances, dtype=bool)
+            mask[indices] = True
+            indices = mask
+
+        observation, info = self.wrapped.reset(options={"reset_mask": indices})
+        if self.wrapped.render_mode is not None:
+            self.wrapped.render()
+        # TODO: process arrays in info
+        return observation, None, info
+
+    def step(self, action: Array):
+        if isinstance(self.wrapped.single_action_space, gym.spaces.Discrete):
+            action = np.argmax(action, axis=-1)
+        observation, reward, terminated, truncated, info = self.wrapped.step(action)
+        if isinstance(reward, np.ndarray):
+            reward = reward.astype(np.float32)
+        reward = reward.reshape(-1, 1)
+        terminated = terminated.reshape(-1, 1)
+        truncated = truncated.reshape(-1, 1)
+        if self.wrapped.render_mode is not None:
+            self.wrapped.render()
+        # TODO: process arrays in info
+        return observation, None, reward, terminated, truncated, info
+
+
+def make_gym_env(
+    id: str | EnvSpec,
+    max_episode_steps: int | None = None,
+    disable_env_checker: bool | None = None,
+    **kwargs: Any,
+) -> Environment:
+    return GymEnvAdapter(
+        gym.make(
+            id=id,
+            max_episode_steps=max_episode_steps,
+            disable_env_checker=disable_env_checker,
+            **kwargs,
+        )
+    )
+
+
+def make_gym_vec(
+    id: str | EnvSpec,
+    num_envs: int = 1,
+    vectorization_mode: gym.VectorizeMode | str | None = None,
+    vector_kwargs: dict[str, Any] | None = None,
+    wrappers: Sequence[Callable[[gym.Env], gym.Wrapper]] | None = None,
+    **kwargs,
+) -> Environment:
+    return GymVectorEnvAdapter(
+        gym.make_vec(
+            id=id,
+            num_envs=num_envs,
+            vectorization_mode=vectorization_mode,
+            vector_kwargs=(vector_kwargs or {}) | {"autoreset_mode": gym.vector.AutoresetMode.DISABLED},
+            wrappers=wrappers,
+            **kwargs,
+        )
+    )

cusrl/environment/isaaclab.py

@@ -0,0 +1,133 @@
+import argparse
+from collections.abc import Sequence
+from typing import Any
+
+import gymnasium as gym
+import torch
+
+import cusrl.utils
+from cusrl.template import Environment, EnvironmentSpec
+from cusrl.utils.typing import Array, Slice
+
+__all__ = ["IsaacLabEnvAdapter", "make_isaaclab_env"]
+
+
+class IsaacLabEnvAdapter(Environment):
+    def __init__(self, id: str, argv: Sequence[str] | None = None, **kwargs):
+        from isaaclab.app import AppLauncher
+
+        parser = argparse.ArgumentParser(prog="--environment-args", description="IsaacLab environment")
+        parser.add_argument("--num_envs", type=int, metavar="N", help="Number of environments to simulate.")
+        AppLauncher.add_app_launcher_args(parser)
+        args = parser.parse_args(argv or [])
+        args.device = str(cusrl.device())
+        self.app_launcher = AppLauncher(args)
+        self.simulation_app = self.app_launcher.app
+
+        from isaaclab.envs import DirectMARLEnv, DirectRLEnv, ManagerBasedRLEnv, multi_agent_to_single_agent
+        from isaaclab_tasks.utils.parse_cfg import load_cfg_from_registry
+
+        env_cfg = load_cfg_from_registry(id, "env_cfg_entry_point")
+        env_cfg.sim.device = args.device
+        if args.num_envs is not None:
+            env_cfg.scene.num_envs = args.num_envs
+        env_cfg.scene.num_envs = max(env_cfg.scene.num_envs // cusrl.utils.distributed.world_size(), 1)
+        isaaclab_env = gym.make(id, cfg=env_cfg, disable_env_checker=True, **kwargs)
+        if isinstance(isaaclab_env.unwrapped, DirectMARLEnv):
+            isaaclab_env = multi_agent_to_single_agent(isaaclab_env)
+        self.wrapped: ManagerBasedRLEnv | DirectRLEnv = isaaclab_env.unwrapped
+        self.device = self.wrapped.device
+        self.metrics = cusrl.utils.Metrics()
+        super().__init__(
+            self.wrapped.num_envs,
+            self._get_observation_dim(),
+            self._get_action_dim(),
+            self._get_state_dim(),
+            EnvironmentSpec(autoreset=True, final_state_is_missing=True),
+        )
+
+        # Avoid terminal color issues
+        print("\033[0m", end="")
+
+    def __del__(self):
+        if hasattr(self, "wrapped"):
+            self.wrapped.close()
+        if hasattr(self, "simulation_app"):
+            self.simulation_app.close()
+
+    def _get_observation_dim(self) -> int:
+        if hasattr(self.wrapped, "observation_manager"):
+            shape = self.wrapped.observation_manager.group_obs_dim["policy"]
+        else:
+            shape = self.wrapped.single_observation_space["policy"].shape
+
+        if not len(shape) == 1:
+            raise ValueError("Only 1D observation space is supported. ")
+        return shape[0]
+
+    def _get_action_dim(self) -> int:
+        if hasattr(self.wrapped, "action_manager"):
+            return self.wrapped.action_manager.total_action_dim
+        space = self.wrapped.single_action_space
+        if not len(space.shape) == 1:
+            raise ValueError("Only 1D action space is supported. ")
+        return space.shape[0]
+
+    def _get_state_dim(self) -> int | None:
+        shape = None
+        if hasattr(self.wrapped, "observation_manager"):
+            shape = self.wrapped.observation_manager.group_obs_dim.get("critic")
+        else:
+            space = self.wrapped.single_observation_space.get("critic")
+            if space is not None:
+                shape = space.shape
+
+        if shape is None:
+            return None
+        if not len(shape) == 1:
+            raise ValueError("Only 1D state space is supported. ")
+        return shape[0]
+
+    def reset(self, *, indices: Array | Slice | None = None):
+        if indices is None:
+            observation_dict, _ = self.wrapped.reset()
+            self.wrapped.episode_length_buf.random_(int(self.wrapped.max_episode_length))
+            observation = observation_dict.pop("policy")
+            state = observation_dict.pop("critic", None)
+            extras = observation_dict
+        else:
+            if isinstance(indices, slice):
+                indices = torch.arange(self.num_instances, device=self.device)[indices]
+            observation_dict, _ = self.wrapped.reset(env_ids=torch.as_tensor(indices, device=self.device))
+
+            observation = observation_dict.pop("policy", None)
+            state = observation_dict.pop("critic", None)
+            extras = {key: value[indices] for key, value in observation_dict.items()}
+            if observation is not None:
+                observation = observation[indices]
+            if state is not None:
+                state = state[indices]
+
+        return observation, state, extras
+
+    def step(self, action):
+        observation_dict, reward, terminated, truncated, extras = self.wrapped.step(action)
+        observation = observation_dict.pop("policy")
+        state = observation_dict.pop("critic", None)
+        reward = reward.unsqueeze(-1)
+        terminated = terminated.unsqueeze(-1)
+        truncated = truncated.unsqueeze(-1)
+        self.metrics.record(
+            **extras.get("log", {}),
+            **extras.get("episode", {}),
+        )
+        return observation, state, reward, terminated, truncated, observation_dict
+
+    def get_metrics(self):
+        metrics = self.metrics.summary()
+        self.metrics.clear()
+        return metrics
+
+
+def make_isaaclab_env(id: str, argv: Sequence[str] | None = None, **kwargs: Any) -> Environment:
+    return IsaacLabEnvAdapter(id, argv, **kwargs)

cusrl/hook/__init__.py

@@ -0,0 +1,58 @@
+from .advantage import AdvantageNormalization, AdvantageReduction
+from .condition import ConditionalObjectiveActivation
+from .gae import GAE
+from .gradient import GradientClipping
+from .initialization import ModuleInitialization
+from .lr_schedule import AdaptiveLRSchedule, MiniBatchWiseLRSchedule, ThresholdLRSchedule
+from .normalization import ObservationNormalization
+from .on_policy import OnPolicyPreparation
+from .ppo import EntropyLoss, PPOSurrogateLoss
+from .representation import (
+    NextStatePrediction,
+    ReturnPrediction,
+    StatePrediction,
+)
+from .rnd import RandomNetworkDistillation
+from .schedule import (
+    HookActivationSchedule,
+    OnPolicyBufferCapacitySchedule,
+    ParameterSchedule,
+)
+from .smoothness import ActionSmoothnessLoss
+from .statistics import OnPolicyStatistics
+from .symmetry import (
+    SymmetricArchitecture,
+    SymmetricDataAugmentation,
+    SymmetryLoss,
+)
+from .value import ValueComputation, ValueLoss
+
+__all__ = [
+    "ActionSmoothnessLoss",
+    "AdaptiveLRSchedule",
+    "AdvantageNormalization",
+    "AdvantageReduction",
+    "ConditionalObjectiveActivation",
+    "EntropyLoss",
+    "GAE",
+    "GradientClipping",
+    "HookActivationSchedule",
+    "MiniBatchWiseLRSchedule",
+    "ModuleInitialization",
+    "NextStatePrediction",
+    "ObservationNormalization",
+    "OnPolicyBufferCapacitySchedule",
+    "OnPolicyPreparation",
+    "OnPolicyStatistics",
+    "ParameterSchedule",
+    "PPOSurrogateLoss",
+    "RandomNetworkDistillation",
+    "ReturnPrediction",
+    "StatePrediction",
+    "SymmetricArchitecture",
+    "SymmetricDataAugmentation",
+    "SymmetryLoss",
+    "ThresholdLRSchedule",
+    "ValueComputation",
+    "ValueLoss",
+]

cusrl/hook/advantage.py

@@ -0,0 +1,100 @@
+from collections.abc import Sequence
+from typing import Literal
+
+import torch
+from torch import Tensor
+
+from cusrl.template import ActorCritic, Hook
+from cusrl.utils import distributed
+
+__all__ = ["AdvantageReduction", "AdvantageNormalization"]
+
+
+class AdvantageReduction(Hook):
+    """A hook to reduce a multidimensional advantage tensor into a scalar.
+
+    This hook reduces the advantage tensor along its last dimension. This is useful
+    in multi-goal settings where the advantage is a vector.
+
+    Args:
+        reduction (Literal["sum", "mean"], optional):
+            The reduction method to apply, either "sum" or "mean".
+            Defaults to "sum".
+        weight (Sequence[float] | None, optional):
+            An optional sequence of weights to apply element-wise to the
+            advantage tensor before reduction. Defaults to None.
+
+    Raises:
+        ValueError: If an invalid reduction method is provided.
+    """
+
+    def __init__(
+        self,
+        reduction: Literal["sum", "mean"] = "sum",
+        weight: Sequence[float] | None = None,
+    ):
+        if reduction not in ("sum", "mean"):
+            raise ValueError(f"Unknown reduction: '{reduction}'.")
+
+        self.reduction = reduction
+        self.weight: Tensor | None = weight
+
+    def init(self):
+        if self.weight is not None:
+            self.weight = self.agent.to_tensor(self.weight)
+
+    def objective(self, batch: dict[str, Tensor]):
+        if (advantage := batch["advantage"]).size(-1) == 1:
+            return
+        if self.weight is not None:
+            advantage = advantage * self.weight
+        if self.reduction == "sum":
+            advantage = advantage.sum(-1, keepdim=True)
+        elif self.reduction == "mean":
+            advantage = advantage.mean(-1, keepdim=True)
+        else:
+            raise ValueError(f"Unknown reduction: '{self.reduction}'.")
+        batch["advantage"] = advantage
+
+
+class AdvantageNormalization(Hook[ActorCritic]):
+    """A hook to normalize advantages in actor-critic algorithms.
+
+    This hook standardizes the advantages to have a mean of 0 and a standard
+    deviation of 1. This can help stabilize training by preventing the scale of
+    advantages from fluctuating wildly. Normalization can be configured to occur
+    either once on the entire buffer before updates begin, or on each mini-batch
+    during the objective calculation.
+
+    The normalization correctly handles distributed training by averaging the
+    mean and variance across all processes.
+
+    Args:
+        mini_batch_wise (bool, optional):
+            If `True`, normalization is applied to each mini-batch. Defaults to
+            `False`,
+        synchronize (bool, optional):
+            If `True`, the mean and variance are synchronized across all processes
+            in distributed training. Defaults to `True`.
+    """
+
+    def __init__(self, mini_batch_wise: bool = False, synchronize: bool = True):
+        self.mini_batch_wise = mini_batch_wise
+        self.synchronize = synchronize
+
+    def pre_update(self, buffer):
+        if not self.mini_batch_wise:
+            self.normalize_(buffer["advantage"])
+
+    def objective(self, batch):
+        if self.mini_batch_wise:
+            self.normalize_(batch["advantage"])
+
+    @torch.no_grad()
+    def normalize_(self, advantage: Tensor):
+        dims = tuple(range(advantage.ndim - 1))
+        var, mean = torch.var_mean(advantage, dim=dims, correction=0)
+        if self.synchronize:
+            distributed.reduce_mean_var_(mean, var)
+        std = (var + 1e-8).sqrt()
+        advantage.sub_(mean).div_(std)

cusrl/hook/condition.py

@@ -0,0 +1,57 @@
+from collections.abc import Callable, Iterable
+from typing import Any
+
+import torch
+
+from cusrl.template import ActorCritic, Buffer, Hook
+
+__all__ = ["ConditionalObjectiveActivation", "EpochIndexCondition"]
+
+
+class EpochIndexCondition:
+    """Checks if the current epoch index is in a specified set of epoch indices.
+
+    Args:
+        epoch_index (int | Iterable[int]):
+            A single epoch index or an iterable of epoch indices.
+    """
+
+    def __init__(self, epoch_index: int | Iterable[int]):
+        if isinstance(epoch_index, int):
+            epoch_index = [epoch_index]
+        self.epoch_index = set(epoch_index)
+
+    def __call__(self, agent: ActorCritic, batch: dict[str, Any]) -> bool:
+        return batch["epoch_index"] in self.epoch_index
+
+
+class ConditionalObjectiveActivation(Hook[ActorCritic]):
+    """A hook to activate other objective hooks based on specified conditions.
+
+    This hook must be placed before any objective hooks it controls.
+
+    Args:
+        named_conditions (Callable[[ActorCritic, dict[str, Any]], bool]):
+            Keyword arguments mapping the name of an objective hook to a callable
+            condition. The condition determines whether the corresponding hook
+            should be active. It receives the agent and the current batch and
+            returns `True` if the hook should be active, `False` otherwise.
+    """
+
+    def __init__(self, **named_conditions: Callable[[ActorCritic, dict[str, Any]], bool]):
+        self.named_conditions = named_conditions
+        self.named_activation = {}
+
+    def pre_update(self, buffer: "Buffer"):
+        # Store the current activation state of the hooks
+        for name in self.named_conditions:
+            self.named_activation[name] = self.agent.hook[name].active
+
+    def objective(self, batch: dict[str, torch.Tensor]) -> torch.Tensor | None:
+        for name, condition in self.named_conditions.items():
+            self.agent.hook[name].active = self.named_activation[name] and condition(self.agent, batch)
+
+    def post_update(self):
+        # Restore the activation state of the hooks
+        for name in self.named_conditions:
+            self.agent.hook[name].active = self.named_activation[name]