diff --git a/.DS_Store b/.DS_Store new file mode 100644 index 0000000..6a88466 Binary files /dev/null and b/.DS_Store differ diff --git a/README.md b/README.md index 7bd2d30..212efed 100644 --- a/README.md +++ b/README.md @@ -15,13 +15,13 @@ cd humanoid_skateboarding uv sync && uv pip install -e . ``` -**(可选)LeRobot v3 导出 / 边播边录** 需要额外安装 `lerobot`(不在默认 `pyproject` 依赖里): +**(可选)LeRobot v3 导出 / 边播边录** 需要额外安装 `lerobot`(不在默认依赖里,以免与现有 PyTorch/CUDA 栈冲突): ```bash uv pip install lerobot ``` -若安装后出现 `import torch` 报 NCCL 符号错误,可尝试: +若安装 `lerobot` 等包后出现 `import torch` 报 **NCCL 符号错误**(例如 `undefined symbol: ncclDevCommDestroy`):多为 **`nvidia-nccl-cu12` 与 `torch`(cu13)并存**,二者都往 `site-packages/nvidia/nccl/lib/` 装 `libnccl.so.2`,旧库覆盖了新库。可卸载 cu12 并重装 cu13 的 NCCL: ```bash uv pip uninstall nvidia-nccl-cu12 @@ -58,7 +58,46 @@ uv run play Mjlab-Skater-Flat-Unitree-G1 --checkpoint_file ckpts/test.pt - **`--viewer auto`**(默认):有 `DISPLAY` / `WAYLAND_DISPLAY` 时用 **native**,否则 **rerun**。 - **`--viewer native`**:本机有图形界面时使用 MuJoCo 原生窗口。 - **`--viewer rerun`**:Rerun Web Viewer(无头服务器常用)。 -- **`--viewer viser`**:Viser。 +- **`--viewer rerun_native`**:同一套 mjlab 仿真与策略步进,**同时**打开仓库自带的 **MuJoCo 原生 viewer**(`NativeMujocoViewer` / GLFW)并把离屏相机 / qpos **推到 Rerun**。与「两个进程各跑一套仿真」无关:仍是 **单一 `env`/单一仿真循环**。**不能**与 `--lerobot-record` 共用。 +- **`--viewer viser`**:Viser(浏览器三维面板,另一种自带前端)。 +- **`--viewer rerun_viser`**:**同一进程、单一仿真**,同时在浏览器里打开 **Rerun**(`--rerun-web-port` / `--rerun-grpc-port`)与 **Viser mjlab 面板**(`--viser-port`)。三者各占不同端口;适合 RoboHub 左 Rerun、右 Mujoco 双 iframe。 + - **注意**:`rerun_native` 里的 **native 是 GLFW 桌面窗口,不占用 HTTP 端口**;若你要「两个端口都是网页服务」,用 **`rerun_viser`**,不要用 `rerun_native` 来凑端口。 + +`rerun_native` 示例(端口与 `rerun` 相同,见下节 SSH 转发): + +```bash +uv run play Mjlab-Skater-Flat-Unitree-G1 --checkpoint_file ckpts/test.pt \ + --viewer rerun_native --rerun-web-port 18080 --rerun-grpc-port 19876 +``` + +`rerun_viser` 示例(**三个端口**:Rerun Web、Rerun gRPC、Viser,互不重复): + +```bash +uv run play Mjlab-Skater-Flat-Unitree-G1 --checkpoint_file ckpts/test.pt \ + --viewer rerun_viser \ + --rerun-web-port 18080 \ + --rerun-grpc-port 19876 \ + --viser-port 19090 +``` + +远程浏览器需 **三个** 本地转发(把示例端口换成你实际用的): + +```bash +ssh -N \ + -L 18080:127.0.0.1:18080 \ + -L 19876:127.0.0.1:19876 \ + -L 19090:127.0.0.1:19090 \ + user@云主机 +``` + +**SSH / 无桌面 / RoboHub 技能里 `RuntimeError: … DISPLAY`:** +`rerun_native` 里的「自带 viewer」是 **本机 X11/Wayland 上的 GLFW 窗口**,不是 Rerun 网页。若 shell 里 **没有** `DISPLAY` 或 `WAYLAND_DISPLAY`(很多容器/编排默认不传),会报错。处理方式: + +- **只想要浏览器里看 Rerun**(单后端、无 MuJoCo 小窗):用 `--viewer rerun`。 +- **仍要 `rerun_native` 但机器无物理桌面**:可装 `xvfb` 用虚拟显示,例如: + `xvfb-run -a uv run play Mjlab-Skater-Flat-Unitree-G1 ... --viewer rerun_native` + (具体以你镜像是否已含 `xvfb` 为准。) +- **RoboHub 侧**:需在技能/容器环境注入 `DISPLAY` 或把启动命令包在 `xvfb-run` 里,否则与本地终端直跑表现一致。 完整参数: @@ -68,14 +107,22 @@ uv run play Mjlab-Skater-Flat-Unitree-G1 --help ### 无头 OpenGL(MuJoCo 离屏相机) -在无 `DISPLAY` 的 Linux 上,`play` 会在导入 MuJoCo 前尽量设置 **`MUJOCO_GL=egl`**(见 `mjlab_husky/mujoco_gl.py`)。若仍失败可手动指定: +在无 `DISPLAY` / `WAYLAND_DISPLAY` 的 Linux 上,`play` 等在 **`import mujoco` 之前** 调用 `mujlab_husky/mujoco_gl.py`:未设置 `MUJOCO_GL` 时默认 **`osmesa`**。仅设 `MUJOCO_GL` 不够:无头时 PyOpenGL 仍可能按 `linux` 选 **GLX**,导致 `glGetError` / `eglQueryString`;因此脚本会同步设置 **`PYOPENGL_PLATFORM=osmesa`**(或在你使用 `MUJOCO_GL=egl` 时为 **`egl`**)。可按需手动指定: ```bash -export MUJOCO_GL=egl # GPU 无头(常见) -# 或 -export MUJOCO_GL=osmesa # 纯 CPU 软件光栅(更慢) +export MUJOCO_GL=egl # GPU + 可用 NVIDIA EGL 时(更快) +# 未设置时由 mujoco_gl 默认 osmesa;或显式 CPU 光栅: +export MUJOCO_GL=osmesa # 需系统已装 libosmesa6(见下) ``` +**Ubuntu(OSMesa)**:若仍报 OpenGL / `glGetError`,请先安装运行时: + +```bash +sudo apt-get update && sudo apt-get install -y libosmesa6 +``` + +若报错 **`mjENBL_MULTICCD`**:来自 **MuJoCo Python 枚举与 `mujoco-warp` Git 修订不一致**。本项目用 PyPI `mujoco==3.8.x` 时,`uv.lock` 已将 **`mujoco-warp` 固定为上游标签 `v3.8.0`**;若在别处自行 `uv lock --upgrade-package mjlab`,需再次确认锁里两处一致。 + ### Rerun:端口与远程浏览器 Rerun 需要 **两个端口**:**Web**(默认 `8080`)+ **gRPC**(默认多为 `9876`,以终端打印为准)。 diff --git a/dataset/.DS_Store b/dataset/.DS_Store new file mode 100644 index 0000000..fcfb195 Binary files /dev/null and b/dataset/.DS_Store differ diff --git a/pyproject.toml b/pyproject.toml index 558084c..5fc73a1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -31,6 +31,12 @@ play-lerobot-rerun = "mjlab_husky.scripts.play_lerobot_rerun:main" src = ["src"] # Helpful for recognizing first-party imports. indent-width = 4 +[tool.uv] +# py.mujoco.org 上的 3.7.0.dev* 预发布包会被撤下,wheel 404;强制使用 PyPI 稳定版。 +override-dependencies = ["mujoco>=3.8.0,<3.9"] +# `mjlab` 声明的 mujoco-warp git rev(1dc288c)依赖 mjENBL_MULTICCD,PyPI 的 mujoco 3.8.x 尚无该枚举。 +# `uv.lock` 将 mujoco-warp 固定为上游标签 v3.8.0(与 mujoco 3.8.0 对齐);升级 mjlab / 跑 `uv lock` 后若冲突请复查该包。 + [tool.uv.sources] -mjlab = { path = "/opt/vendor/mjlab" } -rsl-rl-lib = { path = "rsl_rl" } \ No newline at end of file +mjlab = { git = "https://github.com/mujocolab/mjlab.git", rev = "13212ad" } +rsl-rl-lib = { path = "rsl_rl" } diff --git a/rsl_rl/.DS_Store b/rsl_rl/.DS_Store new file mode 100644 index 0000000..ba25014 Binary files /dev/null and b/rsl_rl/.DS_Store differ diff --git a/rsl_rl/__pycache__/__init__.cpython-312.pyc b/rsl_rl/__pycache__/__init__.cpython-312.pyc deleted file mode 100644 index 9679b24..0000000 Binary files a/rsl_rl/__pycache__/__init__.cpython-312.pyc and /dev/null differ diff --git a/rsl_rl/algorithms/__pycache__/__init__.cpython-312.pyc b/rsl_rl/algorithms/__pycache__/__init__.cpython-312.pyc deleted file mode 100644 index 6d2ed60..0000000 Binary files a/rsl_rl/algorithms/__pycache__/__init__.cpython-312.pyc and /dev/null differ diff --git a/rsl_rl/algorithms/__pycache__/amp_ppo.cpython-312.pyc b/rsl_rl/algorithms/__pycache__/amp_ppo.cpython-312.pyc deleted file mode 100644 index a36b38a..0000000 Binary files a/rsl_rl/algorithms/__pycache__/amp_ppo.cpython-312.pyc and /dev/null differ diff --git a/rsl_rl/algorithms/__pycache__/distillation.cpython-312.pyc b/rsl_rl/algorithms/__pycache__/distillation.cpython-312.pyc deleted file mode 100644 index 145edc8..0000000 Binary files a/rsl_rl/algorithms/__pycache__/distillation.cpython-312.pyc and /dev/null differ diff --git a/rsl_rl/algorithms/__pycache__/ppo.cpython-312.pyc b/rsl_rl/algorithms/__pycache__/ppo.cpython-312.pyc deleted file mode 100644 index d1ce2e5..0000000 Binary files a/rsl_rl/algorithms/__pycache__/ppo.cpython-312.pyc and /dev/null differ diff --git a/rsl_rl/build/lib/rsl_rl/__init__.py b/rsl_rl/build/lib/rsl_rl/__init__.py deleted file mode 100644 index ebc2e20..0000000 --- a/rsl_rl/build/lib/rsl_rl/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -"""Main module for the rsl_rl package.""" diff --git a/rsl_rl/build/lib/rsl_rl/algorithms/__init__.py b/rsl_rl/build/lib/rsl_rl/algorithms/__init__.py deleted file mode 100644 index 0ef9163..0000000 --- a/rsl_rl/build/lib/rsl_rl/algorithms/__init__.py +++ /dev/null @@ -1,11 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -"""Implementation of different RL agents.""" - -from .distillation import Distillation -from .ppo import PPO -from .amp_ppo import AMP_PPO -__all__ = ["PPO", "Distillation", "AMP_PPO"] diff --git a/rsl_rl/build/lib/rsl_rl/algorithms/amp_ppo.py b/rsl_rl/build/lib/rsl_rl/algorithms/amp_ppo.py deleted file mode 100644 index a2164eb..0000000 --- a/rsl_rl/build/lib/rsl_rl/algorithms/amp_ppo.py +++ /dev/null @@ -1,571 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -from __future__ import annotations -from torch._tensor import Tensor -from torch._tensor import Tensor -from typing import Any -from copy import deepcopy - -import torch -import torch.nn as nn -import torch.optim as optim -from itertools import chain - -from rsl_rl.modules import ActorCritic -from rsl_rl.modules.rnd import RandomNetworkDistillation -from rsl_rl.storage import RolloutStorage, ReplayBufferMulti -from rsl_rl.utils import string_to_callable - - -class AMP_PPO: - """Proximal Policy Optimization algorithm (https://arxiv.org/abs/1707.06347).""" - - policy: ActorCritic - """The actor critic module.""" - - def __init__( - self, - policy, - discriminator, - amp_data, - amp_normalizer, - amp_num_frames=1, - amp_replay_buffer_size=100000, - num_learning_epochs=5, - num_mini_batches=4, - clip_param=0.2, - gamma=0.99, - lam=0.95, - value_loss_coef=1.0, - entropy_coef=0.01, - learning_rate=0.001, - max_grad_norm=1.0, - use_clipped_value_loss=True, - schedule="adaptive", - desired_kl=0.01, - device="cpu", - normalize_advantage_per_mini_batch=False, - # RND parameters - rnd_cfg: dict | None = None, - # Symmetry parameters - symmetry_cfg: dict | None = None, - # Distributed training parameters - multi_gpu_cfg: dict | None = None, - ): - # device-related parameters - self.device = device - self.is_multi_gpu = multi_gpu_cfg is not None - # Multi-GPU parameters - if multi_gpu_cfg is not None: - self.gpu_global_rank = multi_gpu_cfg["global_rank"] - self.gpu_world_size = multi_gpu_cfg["world_size"] - else: - self.gpu_global_rank = 0 - self.gpu_world_size = 1 - - # RND components - if rnd_cfg is not None: - # Extract parameters used in ppo - rnd_lr = rnd_cfg.pop("learning_rate", 1e-3) - # Create RND module - self.rnd = RandomNetworkDistillation(device=self.device, **rnd_cfg) - # Create RND optimizer - params = self.rnd.predictor.parameters() - self.rnd_optimizer = optim.Adam(params, lr=rnd_lr) - else: - self.rnd = None - self.rnd_optimizer = None - - # Symmetry components - if symmetry_cfg is not None: - # Check if symmetry is enabled - use_symmetry = symmetry_cfg["use_data_augmentation"] or symmetry_cfg["use_mirror_loss"] - # Print that we are not using symmetry - if not use_symmetry: - print("Symmetry not used for learning. We will use it for logging instead.") - # If function is a string then resolve it to a function - if isinstance(symmetry_cfg["data_augmentation_func"], str): - symmetry_cfg["data_augmentation_func"] = string_to_callable(symmetry_cfg["data_augmentation_func"]) - # Check valid configuration - if symmetry_cfg["use_data_augmentation"] and not callable(symmetry_cfg["data_augmentation_func"]): - raise ValueError( - "Data augmentation enabled but the function is not callable:" - f" {symmetry_cfg['data_augmentation_func']}" - ) - # Store symmetry configuration - self.symmetry = symmetry_cfg - else: - self.symmetry = None - - ## AMP components - self.discriminator = discriminator - self.discriminator.to(self.device) - - self.amp_storage = ReplayBufferMulti(discriminator.state_dim, amp_replay_buffer_size, amp_num_frames, device) - self.amp_data = amp_data - self.amp_normalizer = amp_normalizer - - # PPO components - self.policy = policy - self.policy.to(self.device) - - # Create rollout storage - self.storage: RolloutStorage = None # type: ignore - self.transition = RolloutStorage.Transition() - self.amp_transition = RolloutStorage.Transition() - params = [ - {'params': self.policy.parameters(), 'name': 'policy'}, - ] - - params.append({ - 'params': self.discriminator.trunk.parameters(), - 'weight_decay': 10e-4, - 'name': f'amp_trunk' - }) - params.append({ - 'params': self.discriminator.amp_linear.parameters(), - 'weight_decay': 10e-2, - 'name': f'amp_head' - }) - - # Create optimizer - self.optimizer = optim.Adam(params, lr=learning_rate) - - # PPO parameters - self.clip_param = clip_param - self.num_learning_epochs = num_learning_epochs - self.num_mini_batches = num_mini_batches - self.value_loss_coef = value_loss_coef - self.entropy_coef = entropy_coef - self.gamma = gamma - self.lam = lam - self.max_grad_norm = max_grad_norm - self.use_clipped_value_loss = use_clipped_value_loss - self.desired_kl = desired_kl - self.schedule = schedule - self.learning_rate = learning_rate - self.normalize_advantage_per_mini_batch = normalize_advantage_per_mini_batch - - - def init_storage(self, training_type, num_envs, num_transitions_per_env, obs, actions_shape): - # create rollout storage - self.storage = RolloutStorage( - training_type, - num_envs, - num_transitions_per_env, - obs, - actions_shape, - self.device, - ) - - def act(self, obs, amp_obs): - if self.policy.is_recurrent: - self.transition.hidden_states = self.policy.get_hidden_states() - # compute the actions and values - self.transition.actions = self.policy.act(obs).detach() - self.transition.values = self.policy.evaluate(obs).detach() - self.transition.actions_log_prob = self.policy.get_actions_log_prob(self.transition.actions).detach() - self.transition.action_mean = self.policy.action_mean.detach() - self.transition.action_sigma = self.policy.action_std.detach() - # need to record obs before env.step() - self.transition.observations = obs - self.amp_transition.observations = amp_obs - return self.transition.actions - - def process_env_step(self, obs, rewards, dones, extras,amp_obs, amp_obs_frames=None): - # update the normalizers - self.policy.update_normalization(obs) - if self.rnd: - self.rnd.update_normalization(obs) - - # Record the rewards and dones - # Note: we clone here because later on we bootstrap the rewards based on timeouts - self.transition.rewards = rewards.clone() - self.transition.dones = dones - - # Compute the intrinsic rewards and add to extrinsic rewards - if self.rnd: - # Compute the intrinsic rewards - self.intrinsic_rewards = self.rnd.get_intrinsic_reward(obs) - # Add intrinsic rewards to extrinsic rewards - self.transition.rewards += self.intrinsic_rewards - - # Bootstrapping on time outs - if "time_outs" in extras: - self.transition.rewards += self.gamma * torch.squeeze( - self.transition.values * extras["time_outs"].unsqueeze(1).to(self.device), 1 - ) - - if amp_obs_frames is not None: - self.amp_storage.insert(amp_obs_frames) - else: - self.amp_storage.insert(self.amp_transition.observations, amp_obs) - - # record the transition - self.storage.add_transitions(self.transition) - self.transition.clear() - self.amp_transition.clear() - self.policy.reset(dones) - - def compute_returns(self, obs): - # compute value for the last step - last_values = self.policy.evaluate(obs).detach() - self.storage.compute_returns( - last_values, self.gamma, self.lam, normalize_advantage=not self.normalize_advantage_per_mini_batch - ) - - def update(self): # noqa: C901 - mean_value_loss = 0 - mean_surrogate_loss = 0 - mean_entropy = 0 - mean_amp_loss = 0 - mean_grad_pen_loss = 0 - mean_policy_pred = 0 - mean_expert_pred = 0 - # -- RND loss - if self.rnd: - mean_rnd_loss = 0 - else: - mean_rnd_loss = None - # -- Symmetry loss - if self.symmetry: - mean_symmetry_loss = 0 - else: - mean_symmetry_loss = None - - # generator for mini batches - if self.policy.is_recurrent: - generator = self.storage.recurrent_mini_batch_generator(self.num_mini_batches, self.num_learning_epochs) - else: - generator = self.storage.mini_batch_generator(self.num_mini_batches, self.num_learning_epochs) - - - amp_policy_generator = self.amp_storage.feed_forward_generator( - self.num_learning_epochs * self.num_mini_batches, - self.storage.num_envs * self.storage.num_transitions_per_env // self.num_mini_batches, - ) - - amp_expert_generator = self.amp_data.feed_forward_generator_23dof_multi( - self.num_learning_epochs * self.num_mini_batches, - self.storage.num_envs * self.storage.num_transitions_per_env // self.num_mini_batches, - ) - - # iterate over batches - for sample, sample_amp_policy, sample_amp_expert in zip(generator, amp_policy_generator, amp_expert_generator): - ( - obs_batch, - actions_batch, - target_values_batch, - advantages_batch, - returns_batch, - old_actions_log_prob_batch, - old_mu_batch, - old_sigma_batch, - hid_states_batch, - masks_batch, - ) = sample - - # number of augmentations per sample - # we start with 1 and increase it if we use symmetry augmentation - num_aug = 1 - # original batch size - # we assume policy group is always there and needs augmentation - original_batch_size = obs_batch.batch_size[0] - - # check if we should normalize advantages per mini batch - if self.normalize_advantage_per_mini_batch: - with torch.no_grad(): - advantages_batch = (advantages_batch - advantages_batch.mean()) / (advantages_batch.std() + 1e-8) - - # Perform symmetric augmentation - if self.symmetry and self.symmetry["use_data_augmentation"]: - # augmentation using symmetry - data_augmentation_func = self.symmetry["data_augmentation_func"] - # returned shape: [batch_size * num_aug, ...] - obs_batch, actions_batch = data_augmentation_func( - obs=obs_batch, - actions=actions_batch, - env=self.symmetry["_env"], - ) - # compute number of augmentations per sample - # we assume policy group is always there and needs augmentation - num_aug = int(obs_batch.batch_size[0] / original_batch_size) - # repeat the rest of the batch - # -- actor - old_actions_log_prob_batch = old_actions_log_prob_batch.repeat(num_aug, 1) - # -- critic - target_values_batch = target_values_batch.repeat(num_aug, 1) - advantages_batch = advantages_batch.repeat(num_aug, 1) - returns_batch = returns_batch.repeat(num_aug, 1) - - # Recompute actions log prob and entropy for current batch of transitions - # Note: we need to do this because we updated the policy with the new parameters - # -- actor - self.policy.act(obs_batch, masks=masks_batch, hidden_states=hid_states_batch[0]) - actions_log_prob_batch = self.policy.get_actions_log_prob(actions_batch) - # -- critic - value_batch = self.policy.evaluate(obs_batch, masks=masks_batch, hidden_states=hid_states_batch[1]) - # -- entropy - # we only keep the entropy of the first augmentation (the original one) - mu_batch = self.policy.action_mean[:original_batch_size] - sigma_batch = self.policy.action_std[:original_batch_size] - entropy_batch = self.policy.entropy[:original_batch_size] - - # KL - if self.desired_kl is not None and self.schedule == "adaptive": - with torch.inference_mode(): - kl = torch.sum( - torch.log(sigma_batch / old_sigma_batch + 1.0e-5) - + (torch.square(old_sigma_batch) + torch.square(old_mu_batch - mu_batch)) - / (2.0 * torch.square(sigma_batch)) - - 0.5, - axis=-1, - ) - kl_mean = torch.mean(kl) - - # Reduce the KL divergence across all GPUs - if self.is_multi_gpu: - torch.distributed.all_reduce(kl_mean, op=torch.distributed.ReduceOp.SUM) - kl_mean /= self.gpu_world_size - - # Update the learning rate - # Perform this adaptation only on the main process - # TODO: Is this needed? If KL-divergence is the "same" across all GPUs, - # then the learning rate should be the same across all GPUs. - if self.gpu_global_rank == 0: - if kl_mean > self.desired_kl * 2.0: - self.learning_rate = max(1e-5, self.learning_rate / 1.5) - elif kl_mean < self.desired_kl / 2.0 and kl_mean > 0.0: - self.learning_rate = min(1e-2, self.learning_rate * 1.5) - - # Update the learning rate for all GPUs - if self.is_multi_gpu: - lr_tensor = torch.tensor(self.learning_rate, device=self.device) - torch.distributed.broadcast(lr_tensor, src=0) - self.learning_rate = lr_tensor.item() - - # Update the learning rate for all parameter groups - for param_group in self.optimizer.param_groups: - param_group["lr"] = self.learning_rate - - # Surrogate loss - ratio = torch.exp(actions_log_prob_batch - torch.squeeze(old_actions_log_prob_batch)) - surrogate = -torch.squeeze(advantages_batch) * ratio - surrogate_clipped = -torch.squeeze(advantages_batch) * torch.clamp( - ratio, 1.0 - self.clip_param, 1.0 + self.clip_param - ) - surrogate_loss = torch.max(surrogate, surrogate_clipped).mean() - - # Value function loss - if self.use_clipped_value_loss: - value_clipped = target_values_batch + (value_batch - target_values_batch).clamp( - -self.clip_param, self.clip_param - ) - value_losses = (value_batch - returns_batch).pow(2) - value_losses_clipped = (value_clipped - returns_batch).pow(2) - value_loss = torch.max(value_losses, value_losses_clipped).mean() - else: - value_loss = (returns_batch - value_batch).pow(2).mean() - - loss = surrogate_loss + self.value_loss_coef * value_loss - self.entropy_coef * entropy_batch.mean() - - # Symmetry loss - if self.symmetry: - # obtain the symmetric actions - # if we did augmentation before then we don't need to augment again - if not self.symmetry["use_data_augmentation"]: - data_augmentation_func = self.symmetry["data_augmentation_func"] - obs_batch, _ = data_augmentation_func(obs=obs_batch, actions=None, env=self.symmetry["_env"]) - # compute number of augmentations per sample - num_aug = int(obs_batch.shape[0] / original_batch_size) - - # actions predicted by the actor for symmetrically-augmented observations - mean_actions_batch = self.policy.act_inference(obs_batch.detach().clone()) - - # compute the symmetrically augmented actions - # note: we are assuming the first augmentation is the original one. - # We do not use the action_batch from earlier since that action was sampled from the distribution. - # However, the symmetry loss is computed using the mean of the distribution. - action_mean_orig = mean_actions_batch[:original_batch_size] - _, actions_mean_symm_batch = data_augmentation_func( - obs=None, actions=action_mean_orig, env=self.symmetry["_env"] - ) - - # compute the loss (we skip the first augmentation as it is the original one) - mse_loss = torch.nn.MSELoss() - symmetry_loss = mse_loss( - mean_actions_batch[original_batch_size:], actions_mean_symm_batch.detach()[original_batch_size:] - ) - # add the loss to the total loss - if self.symmetry["use_mirror_loss"]: - loss += self.symmetry["mirror_loss_coeff"] * symmetry_loss - else: - symmetry_loss = symmetry_loss.detach() - - # Random Network Distillation loss - # TODO: Move this processing to inside RND module. - if self.rnd: - # extract the rnd_state - # TODO: Check if we still need torch no grad. It is just an affine transformation. - with torch.no_grad(): - rnd_state_batch = self.rnd.get_rnd_state(obs_batch[:original_batch_size]) - rnd_state_batch = self.rnd.state_normalizer(rnd_state_batch) - # predict the embedding and the target - predicted_embedding = self.rnd.predictor(rnd_state_batch) - target_embedding = self.rnd.target(rnd_state_batch).detach() - # compute the loss as the mean squared error - mseloss = torch.nn.MSELoss() - rnd_loss = mseloss(predicted_embedding, target_embedding) - - expert_states = sample_amp_expert - policy_states = sample_amp_policy - - with torch.no_grad(): - expert_states = self.amp_normalizer.normalize_torch(expert_states.to(self.device), self.device) - policy_states = self.amp_normalizer.normalize_torch(policy_states, self.device) - - contact_phase_push = obs_batch['critic'][:, -4] - mask_push = contact_phase_push == 1. - - if mask_push.any(): - policy_d = self.discriminator(policy_states.flatten(1)) - expert_states = expert_states.to(self.device) - expert_d = self.discriminator(expert_states.flatten(1)) - - expert_loss = torch.nn.MSELoss()(expert_d, torch.ones(expert_d.size(), device=self.device)) - policy_loss = torch.nn.MSELoss()(policy_d, -1 * torch.ones(policy_d.size(), device=self.device)) - amp_loss = 0.5 * (expert_loss + policy_loss) - - # grad penalty - grad_pen_loss = self.discriminator.compute_grad_pen(expert_states, lambda_=5) - else: - amp_loss = torch.tensor(0.0, device=self.device) - grad_pen_loss = torch.tensor(0.0, device=self.device) - expert_loss = torch.tensor(0.0, device=self.device) - policy_loss = torch.tensor(0.0, device=self.device) - - loss += (amp_loss + grad_pen_loss) - self.amp_normalizer.update(policy_states.cpu().numpy()) - self.amp_normalizer.update(expert_states.cpu().numpy()) - - # Compute the gradients - # -- For PPO - self.optimizer.zero_grad() - loss.backward() - # -- For RND - if self.rnd: - self.rnd_optimizer.zero_grad() # type: ignore - rnd_loss.backward() - - # Collect gradients from all GPUs - if self.is_multi_gpu: - self.reduce_parameters() - - # Apply the gradients - # -- For PPO - nn.utils.clip_grad_norm_(self.policy.parameters(), self.max_grad_norm) - self.optimizer.step() - # -- For RND - if self.rnd_optimizer: - self.rnd_optimizer.step() - - # Store the losses - mean_value_loss += value_loss.item() - mean_surrogate_loss += surrogate_loss.item() - mean_entropy += entropy_batch.mean().item() - mean_amp_loss += amp_loss.item() - mean_grad_pen_loss += grad_pen_loss.item() - mean_policy_pred += policy_loss.mean().item() - mean_expert_pred += expert_loss.mean().item() - # -- RND loss - if mean_rnd_loss is not None: - mean_rnd_loss += rnd_loss.item() - # -- Symmetry loss - if mean_symmetry_loss is not None: - mean_symmetry_loss += symmetry_loss.item() - - # -- For PPO - num_updates = self.num_learning_epochs * self.num_mini_batches - mean_value_loss /= num_updates - mean_surrogate_loss /= num_updates - mean_entropy /= num_updates - mean_amp_loss /= num_updates - mean_grad_pen_loss /= num_updates - mean_policy_pred /= num_updates - mean_expert_pred /= num_updates - - # -- For RND - if mean_rnd_loss is not None: - mean_rnd_loss /= num_updates - # -- For Symmetry - if mean_symmetry_loss is not None: - mean_symmetry_loss /= num_updates - # -- Clear the storage - self.storage.clear() - - # construct the loss dictionary - loss_dict = { - "value_function": mean_value_loss, - "surrogate": mean_surrogate_loss, - "entropy": mean_entropy, - "amp": mean_amp_loss, - "amp_grad_pen": mean_grad_pen_loss, - "amp_policy_pred": mean_policy_pred, - "amp_expert_pred": mean_expert_pred, - } - if self.rnd: - loss_dict["rnd"] = mean_rnd_loss - if self.symmetry: - loss_dict["symmetry"] = mean_symmetry_loss - - return loss_dict - - """ - Helper functions - """ - - def broadcast_parameters(self): - """Broadcast model parameters to all GPUs.""" - # obtain the model parameters on current GPU - model_params = [self.policy.state_dict()] - if self.rnd: - model_params.append(self.rnd.predictor.state_dict()) - # broadcast the model parameters - torch.distributed.broadcast_object_list(model_params, src=0) - # load the model parameters on all GPUs from source GPU - self.policy.load_state_dict(model_params[0]) - if self.rnd: - self.rnd.predictor.load_state_dict(model_params[1]) - - def reduce_parameters(self): - """Collect gradients from all GPUs and average them. - - This function is called after the backward pass to synchronize the gradients across all GPUs. - """ - # Create a tensor to store the gradients - grads = [param.grad.view(-1) for param in self.policy.parameters() if param.grad is not None] - if self.rnd: - grads += [param.grad.view(-1) for param in self.rnd.parameters() if param.grad is not None] - all_grads = torch.cat(grads) - - # Average the gradients across all GPUs - torch.distributed.all_reduce(all_grads, op=torch.distributed.ReduceOp.SUM) - all_grads /= self.gpu_world_size - - # Get all parameters - all_params = self.policy.parameters() - if self.rnd: - all_params = chain(all_params, self.rnd.parameters()) - - # Update the gradients for all parameters with the reduced gradients - offset = 0 - for param in all_params: - if param.grad is not None: - numel = param.numel() - # copy data back from shared buffer - param.grad.data.copy_(all_grads[offset : offset + numel].view_as(param.grad.data)) - # update the offset for the next parameter - offset += numel diff --git a/rsl_rl/build/lib/rsl_rl/algorithms/distillation.py b/rsl_rl/build/lib/rsl_rl/algorithms/distillation.py deleted file mode 100644 index 3a86e00..0000000 --- a/rsl_rl/build/lib/rsl_rl/algorithms/distillation.py +++ /dev/null @@ -1,185 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -import torch -import torch.nn as nn - -from rsl_rl.modules import StudentTeacher, StudentTeacherRecurrent -from rsl_rl.storage import RolloutStorage -from rsl_rl.utils import resolve_optimizer - - -class Distillation: - """Distillation algorithm for training a student model to mimic a teacher model.""" - - policy: StudentTeacher | StudentTeacherRecurrent - """The student teacher model.""" - - def __init__( - self, - policy, - num_learning_epochs=1, - gradient_length=15, - learning_rate=1e-3, - max_grad_norm=None, - loss_type="mse", - optimizer="adam", - device="cpu", - # Distributed training parameters - multi_gpu_cfg: dict | None = None, - ): - # device-related parameters - self.device = device - self.is_multi_gpu = multi_gpu_cfg is not None - # Multi-GPU parameters - if multi_gpu_cfg is not None: - self.gpu_global_rank = multi_gpu_cfg["global_rank"] - self.gpu_world_size = multi_gpu_cfg["world_size"] - else: - self.gpu_global_rank = 0 - self.gpu_world_size = 1 - - # distillation components - self.policy = policy - self.policy.to(self.device) - self.storage = None # initialized later - - # initialize the optimizer - self.optimizer = resolve_optimizer(optimizer)(self.policy.parameters(), lr=learning_rate) - - # initialize the transition - self.transition = RolloutStorage.Transition() - self.last_hidden_states = None - - # distillation parameters - self.num_learning_epochs = num_learning_epochs - self.gradient_length = gradient_length - self.learning_rate = learning_rate - self.max_grad_norm = max_grad_norm - - # initialize the loss function - loss_fn_dict = { - "mse": nn.functional.mse_loss, - "huber": nn.functional.huber_loss, - } - if loss_type in loss_fn_dict: - self.loss_fn = loss_fn_dict[loss_type] - else: - raise ValueError(f"Unknown loss type: {loss_type}. Supported types are: {list(loss_fn_dict.keys())}") - - self.num_updates = 0 - - def init_storage(self, training_type, num_envs, num_transitions_per_env, obs, actions_shape): - # create rollout storage - self.storage = RolloutStorage( - training_type, - num_envs, - num_transitions_per_env, - obs, - actions_shape, - self.device, - ) - - def act(self, obs): - # compute the actions - self.transition.actions = self.policy.act(obs).detach() - self.transition.privileged_actions = self.policy.evaluate(obs).detach() - # record the observations - self.transition.observations = obs - return self.transition.actions - - def process_env_step(self, obs, rewards, dones, extras): - # update the normalizers - self.policy.update_normalization(obs) - - # record the rewards and dones - self.transition.rewards = rewards - self.transition.dones = dones - # record the transition - self.storage.add_transitions(self.transition) - self.transition.clear() - self.policy.reset(dones) - - def update(self): - self.num_updates += 1 - mean_behavior_loss = 0 - loss = 0 - cnt = 0 - - for epoch in range(self.num_learning_epochs): - self.policy.reset(hidden_states=self.last_hidden_states) - self.policy.detach_hidden_states() - for obs, _, privileged_actions, dones in self.storage.generator(): - - # inference the student for gradient computation - actions = self.policy.act_inference(obs) - - # behavior cloning loss - behavior_loss = self.loss_fn(actions, privileged_actions) - - # total loss - loss = loss + behavior_loss - mean_behavior_loss += behavior_loss.item() - cnt += 1 - - # gradient step - if cnt % self.gradient_length == 0: - self.optimizer.zero_grad() - loss.backward() - if self.is_multi_gpu: - self.reduce_parameters() - if self.max_grad_norm: - nn.utils.clip_grad_norm_(self.policy.student.parameters(), self.max_grad_norm) - self.optimizer.step() - self.policy.detach_hidden_states() - loss = 0 - - # reset dones - self.policy.reset(dones.view(-1)) - self.policy.detach_hidden_states(dones.view(-1)) - - mean_behavior_loss /= cnt - self.storage.clear() - self.last_hidden_states = self.policy.get_hidden_states() - self.policy.detach_hidden_states() - - # construct the loss dictionary - loss_dict = {"behavior": mean_behavior_loss} - - return loss_dict - - """ - Helper functions - """ - - def broadcast_parameters(self): - """Broadcast model parameters to all GPUs.""" - # obtain the model parameters on current GPU - model_params = [self.policy.state_dict()] - # broadcast the model parameters - torch.distributed.broadcast_object_list(model_params, src=0) - # load the model parameters on all GPUs from source GPU - self.policy.load_state_dict(model_params[0]) - - def reduce_parameters(self): - """Collect gradients from all GPUs and average them. - - This function is called after the backward pass to synchronize the gradients across all GPUs. - """ - # Create a tensor to store the gradients - grads = [param.grad.view(-1) for param in self.policy.parameters() if param.grad is not None] - all_grads = torch.cat(grads) - # Average the gradients across all GPUs - torch.distributed.all_reduce(all_grads, op=torch.distributed.ReduceOp.SUM) - all_grads /= self.gpu_world_size - # Update the gradients for all parameters with the reduced gradients - offset = 0 - for param in self.policy.parameters(): - if param.grad is not None: - numel = param.numel() - # copy data back from shared buffer - param.grad.data.copy_(all_grads[offset : offset + numel].view_as(param.grad.data)) - # update the offset for the next parameter - offset += numel diff --git a/rsl_rl/build/lib/rsl_rl/algorithms/ppo.py b/rsl_rl/build/lib/rsl_rl/algorithms/ppo.py deleted file mode 100644 index 6c21fc5..0000000 --- a/rsl_rl/build/lib/rsl_rl/algorithms/ppo.py +++ /dev/null @@ -1,469 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -from __future__ import annotations - -import torch -import torch.nn as nn -import torch.optim as optim -from itertools import chain - -from rsl_rl.modules import ActorCritic -from rsl_rl.modules.rnd import RandomNetworkDistillation -from rsl_rl.storage import RolloutStorage -from rsl_rl.utils import string_to_callable - - -class PPO: - """Proximal Policy Optimization algorithm (https://arxiv.org/abs/1707.06347).""" - - policy: ActorCritic - """The actor critic module.""" - - def __init__( - self, - policy, - num_learning_epochs=5, - num_mini_batches=4, - clip_param=0.2, - gamma=0.99, - lam=0.95, - value_loss_coef=1.0, - entropy_coef=0.01, - learning_rate=0.001, - max_grad_norm=1.0, - use_clipped_value_loss=True, - schedule="adaptive", - desired_kl=0.01, - device="cpu", - normalize_advantage_per_mini_batch=False, - # RND parameters - rnd_cfg: dict | None = None, - # Symmetry parameters - symmetry_cfg: dict | None = None, - # Distributed training parameters - multi_gpu_cfg: dict | None = None, - ): - # device-related parameters - self.device = device - self.is_multi_gpu = multi_gpu_cfg is not None - # Multi-GPU parameters - if multi_gpu_cfg is not None: - self.gpu_global_rank = multi_gpu_cfg["global_rank"] - self.gpu_world_size = multi_gpu_cfg["world_size"] - else: - self.gpu_global_rank = 0 - self.gpu_world_size = 1 - - # RND components - if rnd_cfg is not None: - # Extract parameters used in ppo - rnd_lr = rnd_cfg.pop("learning_rate", 1e-3) - # Create RND module - self.rnd = RandomNetworkDistillation(device=self.device, **rnd_cfg) - # Create RND optimizer - params = self.rnd.predictor.parameters() - self.rnd_optimizer = optim.Adam(params, lr=rnd_lr) - else: - self.rnd = None - self.rnd_optimizer = None - - # Symmetry components - if symmetry_cfg is not None: - # Check if symmetry is enabled - use_symmetry = symmetry_cfg["use_data_augmentation"] or symmetry_cfg["use_mirror_loss"] - # Print that we are not using symmetry - if not use_symmetry: - print("Symmetry not used for learning. We will use it for logging instead.") - # If function is a string then resolve it to a function - if isinstance(symmetry_cfg["data_augmentation_func"], str): - symmetry_cfg["data_augmentation_func"] = string_to_callable(symmetry_cfg["data_augmentation_func"]) - # Check valid configuration - if symmetry_cfg["use_data_augmentation"] and not callable(symmetry_cfg["data_augmentation_func"]): - raise ValueError( - "Data augmentation enabled but the function is not callable:" - f" {symmetry_cfg['data_augmentation_func']}" - ) - # Store symmetry configuration - self.symmetry = symmetry_cfg - else: - self.symmetry = None - - # PPO components - self.policy = policy - self.policy.to(self.device) - # Create optimizer - self.optimizer = optim.Adam(self.policy.parameters(), lr=learning_rate) - # Create rollout storage - self.storage: RolloutStorage = None # type: ignore - self.transition = RolloutStorage.Transition() - - # PPO parameters - self.clip_param = clip_param - self.num_learning_epochs = num_learning_epochs - self.num_mini_batches = num_mini_batches - self.value_loss_coef = value_loss_coef - self.entropy_coef = entropy_coef - self.gamma = gamma - self.lam = lam - self.max_grad_norm = max_grad_norm - self.use_clipped_value_loss = use_clipped_value_loss - self.desired_kl = desired_kl - self.schedule = schedule - self.learning_rate = learning_rate - self.normalize_advantage_per_mini_batch = normalize_advantage_per_mini_batch - - def init_storage(self, training_type, num_envs, num_transitions_per_env, obs, actions_shape): - # create rollout storage - self.storage = RolloutStorage( - training_type, - num_envs, - num_transitions_per_env, - obs, - actions_shape, - self.device, - ) - - def act(self, obs): - if self.policy.is_recurrent: - self.transition.hidden_states = self.policy.get_hidden_states() - # compute the actions and values - self.transition.actions = self.policy.act(obs).detach() - self.transition.values = self.policy.evaluate(obs).detach() - self.transition.actions_log_prob = self.policy.get_actions_log_prob(self.transition.actions).detach() - self.transition.action_mean = self.policy.action_mean.detach() - self.transition.action_sigma = self.policy.action_std.detach() - # need to record obs before env.step() - self.transition.observations = obs - return self.transition.actions - - def process_env_step(self, obs, rewards, dones, extras): - # update the normalizers - self.policy.update_normalization(obs) - if self.rnd: - self.rnd.update_normalization(obs) - - # Record the rewards and dones - # Note: we clone here because later on we bootstrap the rewards based on timeouts - self.transition.rewards = rewards.clone() - self.transition.dones = dones - - # Compute the intrinsic rewards and add to extrinsic rewards - if self.rnd: - # Compute the intrinsic rewards - self.intrinsic_rewards = self.rnd.get_intrinsic_reward(obs) - # Add intrinsic rewards to extrinsic rewards - self.transition.rewards += self.intrinsic_rewards - - # Bootstrapping on time outs - if "time_outs" in extras: - self.transition.rewards += self.gamma * torch.squeeze( - self.transition.values * extras["time_outs"].unsqueeze(1).to(self.device), 1 - ) - - # record the transition - self.storage.add_transitions(self.transition) - self.transition.clear() - self.policy.reset(dones) - - def compute_returns(self, obs): - # compute value for the last step - last_values = self.policy.evaluate(obs).detach() - self.storage.compute_returns( - last_values, self.gamma, self.lam, normalize_advantage=not self.normalize_advantage_per_mini_batch - ) - - def update(self): # noqa: C901 - mean_value_loss = 0 - mean_surrogate_loss = 0 - mean_entropy = 0 - # -- RND loss - if self.rnd: - mean_rnd_loss = 0 - else: - mean_rnd_loss = None - # -- Symmetry loss - if self.symmetry: - mean_symmetry_loss = 0 - else: - mean_symmetry_loss = None - - # generator for mini batches - if self.policy.is_recurrent: - generator = self.storage.recurrent_mini_batch_generator(self.num_mini_batches, self.num_learning_epochs) - else: - generator = self.storage.mini_batch_generator(self.num_mini_batches, self.num_learning_epochs) - - # iterate over batches - for ( - obs_batch, - actions_batch, - target_values_batch, - advantages_batch, - returns_batch, - old_actions_log_prob_batch, - old_mu_batch, - old_sigma_batch, - hid_states_batch, - masks_batch, - ) in generator: - - # number of augmentations per sample - # we start with 1 and increase it if we use symmetry augmentation - num_aug = 1 - # original batch size - # we assume policy group is always there and needs augmentation - original_batch_size = obs_batch.batch_size[0] - - # check if we should normalize advantages per mini batch - if self.normalize_advantage_per_mini_batch: - with torch.no_grad(): - advantages_batch = (advantages_batch - advantages_batch.mean()) / (advantages_batch.std() + 1e-8) - - # Perform symmetric augmentation - if self.symmetry and self.symmetry["use_data_augmentation"]: - # augmentation using symmetry - data_augmentation_func = self.symmetry["data_augmentation_func"] - # returned shape: [batch_size * num_aug, ...] - obs_batch, actions_batch = data_augmentation_func( - obs=obs_batch, - actions=actions_batch, - env=self.symmetry["_env"], - ) - # compute number of augmentations per sample - # we assume policy group is always there and needs augmentation - num_aug = int(obs_batch.batch_size[0] / original_batch_size) - # repeat the rest of the batch - # -- actor - old_actions_log_prob_batch = old_actions_log_prob_batch.repeat(num_aug, 1) - # -- critic - target_values_batch = target_values_batch.repeat(num_aug, 1) - advantages_batch = advantages_batch.repeat(num_aug, 1) - returns_batch = returns_batch.repeat(num_aug, 1) - - # Recompute actions log prob and entropy for current batch of transitions - # Note: we need to do this because we updated the policy with the new parameters - # -- actor - self.policy.act(obs_batch, masks=masks_batch, hidden_states=hid_states_batch[0]) - actions_log_prob_batch = self.policy.get_actions_log_prob(actions_batch) - # -- critic - value_batch = self.policy.evaluate(obs_batch, masks=masks_batch, hidden_states=hid_states_batch[1]) - # -- entropy - # we only keep the entropy of the first augmentation (the original one) - mu_batch = self.policy.action_mean[:original_batch_size] - sigma_batch = self.policy.action_std[:original_batch_size] - entropy_batch = self.policy.entropy[:original_batch_size] - - # KL - if self.desired_kl is not None and self.schedule == "adaptive": - with torch.inference_mode(): - kl = torch.sum( - torch.log(sigma_batch / old_sigma_batch + 1.0e-5) - + (torch.square(old_sigma_batch) + torch.square(old_mu_batch - mu_batch)) - / (2.0 * torch.square(sigma_batch)) - - 0.5, - axis=-1, - ) - kl_mean = torch.mean(kl) - - # Reduce the KL divergence across all GPUs - if self.is_multi_gpu: - torch.distributed.all_reduce(kl_mean, op=torch.distributed.ReduceOp.SUM) - kl_mean /= self.gpu_world_size - - # Update the learning rate - # Perform this adaptation only on the main process - # TODO: Is this needed? If KL-divergence is the "same" across all GPUs, - # then the learning rate should be the same across all GPUs. - if self.gpu_global_rank == 0: - if kl_mean > self.desired_kl * 2.0: - self.learning_rate = max(1e-5, self.learning_rate / 1.5) - elif kl_mean < self.desired_kl / 2.0 and kl_mean > 0.0: - self.learning_rate = min(1e-2, self.learning_rate * 1.5) - - # Update the learning rate for all GPUs - if self.is_multi_gpu: - lr_tensor = torch.tensor(self.learning_rate, device=self.device) - torch.distributed.broadcast(lr_tensor, src=0) - self.learning_rate = lr_tensor.item() - - # Update the learning rate for all parameter groups - for param_group in self.optimizer.param_groups: - param_group["lr"] = self.learning_rate - - # Surrogate loss - ratio = torch.exp(actions_log_prob_batch - torch.squeeze(old_actions_log_prob_batch)) - surrogate = -torch.squeeze(advantages_batch) * ratio - surrogate_clipped = -torch.squeeze(advantages_batch) * torch.clamp( - ratio, 1.0 - self.clip_param, 1.0 + self.clip_param - ) - surrogate_loss = torch.max(surrogate, surrogate_clipped).mean() - - # Value function loss - if self.use_clipped_value_loss: - value_clipped = target_values_batch + (value_batch - target_values_batch).clamp( - -self.clip_param, self.clip_param - ) - value_losses = (value_batch - returns_batch).pow(2) - value_losses_clipped = (value_clipped - returns_batch).pow(2) - value_loss = torch.max(value_losses, value_losses_clipped).mean() - else: - value_loss = (returns_batch - value_batch).pow(2).mean() - - loss = surrogate_loss + self.value_loss_coef * value_loss - self.entropy_coef * entropy_batch.mean() - - # Symmetry loss - if self.symmetry: - # obtain the symmetric actions - # if we did augmentation before then we don't need to augment again - if not self.symmetry["use_data_augmentation"]: - data_augmentation_func = self.symmetry["data_augmentation_func"] - obs_batch, _ = data_augmentation_func(obs=obs_batch, actions=None, env=self.symmetry["_env"]) - # compute number of augmentations per sample - num_aug = int(obs_batch.shape[0] / original_batch_size) - - # actions predicted by the actor for symmetrically-augmented observations - mean_actions_batch = self.policy.act_inference(obs_batch.detach().clone()) - - # compute the symmetrically augmented actions - # note: we are assuming the first augmentation is the original one. - # We do not use the action_batch from earlier since that action was sampled from the distribution. - # However, the symmetry loss is computed using the mean of the distribution. - action_mean_orig = mean_actions_batch[:original_batch_size] - _, actions_mean_symm_batch = data_augmentation_func( - obs=None, actions=action_mean_orig, env=self.symmetry["_env"] - ) - - # compute the loss (we skip the first augmentation as it is the original one) - mse_loss = torch.nn.MSELoss() - symmetry_loss = mse_loss( - mean_actions_batch[original_batch_size:], actions_mean_symm_batch.detach()[original_batch_size:] - ) - # add the loss to the total loss - if self.symmetry["use_mirror_loss"]: - loss += self.symmetry["mirror_loss_coeff"] * symmetry_loss - else: - symmetry_loss = symmetry_loss.detach() - - # Random Network Distillation loss - # TODO: Move this processing to inside RND module. - if self.rnd: - # extract the rnd_state - # TODO: Check if we still need torch no grad. It is just an affine transformation. - with torch.no_grad(): - rnd_state_batch = self.rnd.get_rnd_state(obs_batch[:original_batch_size]) - rnd_state_batch = self.rnd.state_normalizer(rnd_state_batch) - # predict the embedding and the target - predicted_embedding = self.rnd.predictor(rnd_state_batch) - target_embedding = self.rnd.target(rnd_state_batch).detach() - # compute the loss as the mean squared error - mseloss = torch.nn.MSELoss() - rnd_loss = mseloss(predicted_embedding, target_embedding) - - # Compute the gradients - # -- For PPO - self.optimizer.zero_grad() - loss.backward() - # -- For RND - if self.rnd: - self.rnd_optimizer.zero_grad() # type: ignore - rnd_loss.backward() - - # Collect gradients from all GPUs - if self.is_multi_gpu: - self.reduce_parameters() - - # Apply the gradients - # -- For PPO - nn.utils.clip_grad_norm_(self.policy.parameters(), self.max_grad_norm) - self.optimizer.step() - # -- For RND - if self.rnd_optimizer: - self.rnd_optimizer.step() - - # Store the losses - mean_value_loss += value_loss.item() - mean_surrogate_loss += surrogate_loss.item() - mean_entropy += entropy_batch.mean().item() - # -- RND loss - if mean_rnd_loss is not None: - mean_rnd_loss += rnd_loss.item() - # -- Symmetry loss - if mean_symmetry_loss is not None: - mean_symmetry_loss += symmetry_loss.item() - - # -- For PPO - num_updates = self.num_learning_epochs * self.num_mini_batches - mean_value_loss /= num_updates - mean_surrogate_loss /= num_updates - mean_entropy /= num_updates - # -- For RND - if mean_rnd_loss is not None: - mean_rnd_loss /= num_updates - # -- For Symmetry - if mean_symmetry_loss is not None: - mean_symmetry_loss /= num_updates - # -- Clear the storage - self.storage.clear() - - # construct the loss dictionary - loss_dict = { - "value_function": mean_value_loss, - "surrogate": mean_surrogate_loss, - "entropy": mean_entropy, - } - if self.rnd: - loss_dict["rnd"] = mean_rnd_loss - if self.symmetry: - loss_dict["symmetry"] = mean_symmetry_loss - - return loss_dict - - """ - Helper functions - """ - - def broadcast_parameters(self): - """Broadcast model parameters to all GPUs.""" - # obtain the model parameters on current GPU - model_params = [self.policy.state_dict()] - if self.rnd: - model_params.append(self.rnd.predictor.state_dict()) - # broadcast the model parameters - torch.distributed.broadcast_object_list(model_params, src=0) - # load the model parameters on all GPUs from source GPU - self.policy.load_state_dict(model_params[0]) - if self.rnd: - self.rnd.predictor.load_state_dict(model_params[1]) - - def reduce_parameters(self): - """Collect gradients from all GPUs and average them. - - This function is called after the backward pass to synchronize the gradients across all GPUs. - """ - # Create a tensor to store the gradients - grads = [param.grad.view(-1) for param in self.policy.parameters() if param.grad is not None] - if self.rnd: - grads += [param.grad.view(-1) for param in self.rnd.parameters() if param.grad is not None] - all_grads = torch.cat(grads) - - # Average the gradients across all GPUs - torch.distributed.all_reduce(all_grads, op=torch.distributed.ReduceOp.SUM) - all_grads /= self.gpu_world_size - - # Get all parameters - all_params = self.policy.parameters() - if self.rnd: - all_params = chain(all_params, self.rnd.parameters()) - - # Update the gradients for all parameters with the reduced gradients - offset = 0 - for param in all_params: - if param.grad is not None: - numel = param.numel() - # copy data back from shared buffer - param.grad.data.copy_(all_grads[offset : offset + numel].view_as(param.grad.data)) - # update the offset for the next parameter - offset += numel diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/__init__.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/__init__.py deleted file mode 100644 index ebc2e20..0000000 --- a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -"""Main module for the rsl_rl package.""" diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/algorithms/__init__.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/algorithms/__init__.py deleted file mode 100644 index 0ef9163..0000000 --- a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/algorithms/__init__.py +++ /dev/null @@ -1,11 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -"""Implementation of different RL agents.""" - -from .distillation import Distillation -from .ppo import PPO -from .amp_ppo import AMP_PPO -__all__ = ["PPO", "Distillation", "AMP_PPO"] diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/algorithms/amp_ppo.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/algorithms/amp_ppo.py deleted file mode 100644 index a2164eb..0000000 --- a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/algorithms/amp_ppo.py +++ /dev/null @@ -1,571 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -from __future__ import annotations -from torch._tensor import Tensor -from torch._tensor import Tensor -from typing import Any -from copy import deepcopy - -import torch -import torch.nn as nn -import torch.optim as optim -from itertools import chain - -from rsl_rl.modules import ActorCritic -from rsl_rl.modules.rnd import RandomNetworkDistillation -from rsl_rl.storage import RolloutStorage, ReplayBufferMulti -from rsl_rl.utils import string_to_callable - - -class AMP_PPO: - """Proximal Policy Optimization algorithm (https://arxiv.org/abs/1707.06347).""" - - policy: ActorCritic - """The actor critic module.""" - - def __init__( - self, - policy, - discriminator, - amp_data, - amp_normalizer, - amp_num_frames=1, - amp_replay_buffer_size=100000, - num_learning_epochs=5, - num_mini_batches=4, - clip_param=0.2, - gamma=0.99, - lam=0.95, - value_loss_coef=1.0, - entropy_coef=0.01, - learning_rate=0.001, - max_grad_norm=1.0, - use_clipped_value_loss=True, - schedule="adaptive", - desired_kl=0.01, - device="cpu", - normalize_advantage_per_mini_batch=False, - # RND parameters - rnd_cfg: dict | None = None, - # Symmetry parameters - symmetry_cfg: dict | None = None, - # Distributed training parameters - multi_gpu_cfg: dict | None = None, - ): - # device-related parameters - self.device = device - self.is_multi_gpu = multi_gpu_cfg is not None - # Multi-GPU parameters - if multi_gpu_cfg is not None: - self.gpu_global_rank = multi_gpu_cfg["global_rank"] - self.gpu_world_size = multi_gpu_cfg["world_size"] - else: - self.gpu_global_rank = 0 - self.gpu_world_size = 1 - - # RND components - if rnd_cfg is not None: - # Extract parameters used in ppo - rnd_lr = rnd_cfg.pop("learning_rate", 1e-3) - # Create RND module - self.rnd = RandomNetworkDistillation(device=self.device, **rnd_cfg) - # Create RND optimizer - params = self.rnd.predictor.parameters() - self.rnd_optimizer = optim.Adam(params, lr=rnd_lr) - else: - self.rnd = None - self.rnd_optimizer = None - - # Symmetry components - if symmetry_cfg is not None: - # Check if symmetry is enabled - use_symmetry = symmetry_cfg["use_data_augmentation"] or symmetry_cfg["use_mirror_loss"] - # Print that we are not using symmetry - if not use_symmetry: - print("Symmetry not used for learning. We will use it for logging instead.") - # If function is a string then resolve it to a function - if isinstance(symmetry_cfg["data_augmentation_func"], str): - symmetry_cfg["data_augmentation_func"] = string_to_callable(symmetry_cfg["data_augmentation_func"]) - # Check valid configuration - if symmetry_cfg["use_data_augmentation"] and not callable(symmetry_cfg["data_augmentation_func"]): - raise ValueError( - "Data augmentation enabled but the function is not callable:" - f" {symmetry_cfg['data_augmentation_func']}" - ) - # Store symmetry configuration - self.symmetry = symmetry_cfg - else: - self.symmetry = None - - ## AMP components - self.discriminator = discriminator - self.discriminator.to(self.device) - - self.amp_storage = ReplayBufferMulti(discriminator.state_dim, amp_replay_buffer_size, amp_num_frames, device) - self.amp_data = amp_data - self.amp_normalizer = amp_normalizer - - # PPO components - self.policy = policy - self.policy.to(self.device) - - # Create rollout storage - self.storage: RolloutStorage = None # type: ignore - self.transition = RolloutStorage.Transition() - self.amp_transition = RolloutStorage.Transition() - params = [ - {'params': self.policy.parameters(), 'name': 'policy'}, - ] - - params.append({ - 'params': self.discriminator.trunk.parameters(), - 'weight_decay': 10e-4, - 'name': f'amp_trunk' - }) - params.append({ - 'params': self.discriminator.amp_linear.parameters(), - 'weight_decay': 10e-2, - 'name': f'amp_head' - }) - - # Create optimizer - self.optimizer = optim.Adam(params, lr=learning_rate) - - # PPO parameters - self.clip_param = clip_param - self.num_learning_epochs = num_learning_epochs - self.num_mini_batches = num_mini_batches - self.value_loss_coef = value_loss_coef - self.entropy_coef = entropy_coef - self.gamma = gamma - self.lam = lam - self.max_grad_norm = max_grad_norm - self.use_clipped_value_loss = use_clipped_value_loss - self.desired_kl = desired_kl - self.schedule = schedule - self.learning_rate = learning_rate - self.normalize_advantage_per_mini_batch = normalize_advantage_per_mini_batch - - - def init_storage(self, training_type, num_envs, num_transitions_per_env, obs, actions_shape): - # create rollout storage - self.storage = RolloutStorage( - training_type, - num_envs, - num_transitions_per_env, - obs, - actions_shape, - self.device, - ) - - def act(self, obs, amp_obs): - if self.policy.is_recurrent: - self.transition.hidden_states = self.policy.get_hidden_states() - # compute the actions and values - self.transition.actions = self.policy.act(obs).detach() - self.transition.values = self.policy.evaluate(obs).detach() - self.transition.actions_log_prob = self.policy.get_actions_log_prob(self.transition.actions).detach() - self.transition.action_mean = self.policy.action_mean.detach() - self.transition.action_sigma = self.policy.action_std.detach() - # need to record obs before env.step() - self.transition.observations = obs - self.amp_transition.observations = amp_obs - return self.transition.actions - - def process_env_step(self, obs, rewards, dones, extras,amp_obs, amp_obs_frames=None): - # update the normalizers - self.policy.update_normalization(obs) - if self.rnd: - self.rnd.update_normalization(obs) - - # Record the rewards and dones - # Note: we clone here because later on we bootstrap the rewards based on timeouts - self.transition.rewards = rewards.clone() - self.transition.dones = dones - - # Compute the intrinsic rewards and add to extrinsic rewards - if self.rnd: - # Compute the intrinsic rewards - self.intrinsic_rewards = self.rnd.get_intrinsic_reward(obs) - # Add intrinsic rewards to extrinsic rewards - self.transition.rewards += self.intrinsic_rewards - - # Bootstrapping on time outs - if "time_outs" in extras: - self.transition.rewards += self.gamma * torch.squeeze( - self.transition.values * extras["time_outs"].unsqueeze(1).to(self.device), 1 - ) - - if amp_obs_frames is not None: - self.amp_storage.insert(amp_obs_frames) - else: - self.amp_storage.insert(self.amp_transition.observations, amp_obs) - - # record the transition - self.storage.add_transitions(self.transition) - self.transition.clear() - self.amp_transition.clear() - self.policy.reset(dones) - - def compute_returns(self, obs): - # compute value for the last step - last_values = self.policy.evaluate(obs).detach() - self.storage.compute_returns( - last_values, self.gamma, self.lam, normalize_advantage=not self.normalize_advantage_per_mini_batch - ) - - def update(self): # noqa: C901 - mean_value_loss = 0 - mean_surrogate_loss = 0 - mean_entropy = 0 - mean_amp_loss = 0 - mean_grad_pen_loss = 0 - mean_policy_pred = 0 - mean_expert_pred = 0 - # -- RND loss - if self.rnd: - mean_rnd_loss = 0 - else: - mean_rnd_loss = None - # -- Symmetry loss - if self.symmetry: - mean_symmetry_loss = 0 - else: - mean_symmetry_loss = None - - # generator for mini batches - if self.policy.is_recurrent: - generator = self.storage.recurrent_mini_batch_generator(self.num_mini_batches, self.num_learning_epochs) - else: - generator = self.storage.mini_batch_generator(self.num_mini_batches, self.num_learning_epochs) - - - amp_policy_generator = self.amp_storage.feed_forward_generator( - self.num_learning_epochs * self.num_mini_batches, - self.storage.num_envs * self.storage.num_transitions_per_env // self.num_mini_batches, - ) - - amp_expert_generator = self.amp_data.feed_forward_generator_23dof_multi( - self.num_learning_epochs * self.num_mini_batches, - self.storage.num_envs * self.storage.num_transitions_per_env // self.num_mini_batches, - ) - - # iterate over batches - for sample, sample_amp_policy, sample_amp_expert in zip(generator, amp_policy_generator, amp_expert_generator): - ( - obs_batch, - actions_batch, - target_values_batch, - advantages_batch, - returns_batch, - old_actions_log_prob_batch, - old_mu_batch, - old_sigma_batch, - hid_states_batch, - masks_batch, - ) = sample - - # number of augmentations per sample - # we start with 1 and increase it if we use symmetry augmentation - num_aug = 1 - # original batch size - # we assume policy group is always there and needs augmentation - original_batch_size = obs_batch.batch_size[0] - - # check if we should normalize advantages per mini batch - if self.normalize_advantage_per_mini_batch: - with torch.no_grad(): - advantages_batch = (advantages_batch - advantages_batch.mean()) / (advantages_batch.std() + 1e-8) - - # Perform symmetric augmentation - if self.symmetry and self.symmetry["use_data_augmentation"]: - # augmentation using symmetry - data_augmentation_func = self.symmetry["data_augmentation_func"] - # returned shape: [batch_size * num_aug, ...] - obs_batch, actions_batch = data_augmentation_func( - obs=obs_batch, - actions=actions_batch, - env=self.symmetry["_env"], - ) - # compute number of augmentations per sample - # we assume policy group is always there and needs augmentation - num_aug = int(obs_batch.batch_size[0] / original_batch_size) - # repeat the rest of the batch - # -- actor - old_actions_log_prob_batch = old_actions_log_prob_batch.repeat(num_aug, 1) - # -- critic - target_values_batch = target_values_batch.repeat(num_aug, 1) - advantages_batch = advantages_batch.repeat(num_aug, 1) - returns_batch = returns_batch.repeat(num_aug, 1) - - # Recompute actions log prob and entropy for current batch of transitions - # Note: we need to do this because we updated the policy with the new parameters - # -- actor - self.policy.act(obs_batch, masks=masks_batch, hidden_states=hid_states_batch[0]) - actions_log_prob_batch = self.policy.get_actions_log_prob(actions_batch) - # -- critic - value_batch = self.policy.evaluate(obs_batch, masks=masks_batch, hidden_states=hid_states_batch[1]) - # -- entropy - # we only keep the entropy of the first augmentation (the original one) - mu_batch = self.policy.action_mean[:original_batch_size] - sigma_batch = self.policy.action_std[:original_batch_size] - entropy_batch = self.policy.entropy[:original_batch_size] - - # KL - if self.desired_kl is not None and self.schedule == "adaptive": - with torch.inference_mode(): - kl = torch.sum( - torch.log(sigma_batch / old_sigma_batch + 1.0e-5) - + (torch.square(old_sigma_batch) + torch.square(old_mu_batch - mu_batch)) - / (2.0 * torch.square(sigma_batch)) - - 0.5, - axis=-1, - ) - kl_mean = torch.mean(kl) - - # Reduce the KL divergence across all GPUs - if self.is_multi_gpu: - torch.distributed.all_reduce(kl_mean, op=torch.distributed.ReduceOp.SUM) - kl_mean /= self.gpu_world_size - - # Update the learning rate - # Perform this adaptation only on the main process - # TODO: Is this needed? If KL-divergence is the "same" across all GPUs, - # then the learning rate should be the same across all GPUs. - if self.gpu_global_rank == 0: - if kl_mean > self.desired_kl * 2.0: - self.learning_rate = max(1e-5, self.learning_rate / 1.5) - elif kl_mean < self.desired_kl / 2.0 and kl_mean > 0.0: - self.learning_rate = min(1e-2, self.learning_rate * 1.5) - - # Update the learning rate for all GPUs - if self.is_multi_gpu: - lr_tensor = torch.tensor(self.learning_rate, device=self.device) - torch.distributed.broadcast(lr_tensor, src=0) - self.learning_rate = lr_tensor.item() - - # Update the learning rate for all parameter groups - for param_group in self.optimizer.param_groups: - param_group["lr"] = self.learning_rate - - # Surrogate loss - ratio = torch.exp(actions_log_prob_batch - torch.squeeze(old_actions_log_prob_batch)) - surrogate = -torch.squeeze(advantages_batch) * ratio - surrogate_clipped = -torch.squeeze(advantages_batch) * torch.clamp( - ratio, 1.0 - self.clip_param, 1.0 + self.clip_param - ) - surrogate_loss = torch.max(surrogate, surrogate_clipped).mean() - - # Value function loss - if self.use_clipped_value_loss: - value_clipped = target_values_batch + (value_batch - target_values_batch).clamp( - -self.clip_param, self.clip_param - ) - value_losses = (value_batch - returns_batch).pow(2) - value_losses_clipped = (value_clipped - returns_batch).pow(2) - value_loss = torch.max(value_losses, value_losses_clipped).mean() - else: - value_loss = (returns_batch - value_batch).pow(2).mean() - - loss = surrogate_loss + self.value_loss_coef * value_loss - self.entropy_coef * entropy_batch.mean() - - # Symmetry loss - if self.symmetry: - # obtain the symmetric actions - # if we did augmentation before then we don't need to augment again - if not self.symmetry["use_data_augmentation"]: - data_augmentation_func = self.symmetry["data_augmentation_func"] - obs_batch, _ = data_augmentation_func(obs=obs_batch, actions=None, env=self.symmetry["_env"]) - # compute number of augmentations per sample - num_aug = int(obs_batch.shape[0] / original_batch_size) - - # actions predicted by the actor for symmetrically-augmented observations - mean_actions_batch = self.policy.act_inference(obs_batch.detach().clone()) - - # compute the symmetrically augmented actions - # note: we are assuming the first augmentation is the original one. - # We do not use the action_batch from earlier since that action was sampled from the distribution. - # However, the symmetry loss is computed using the mean of the distribution. - action_mean_orig = mean_actions_batch[:original_batch_size] - _, actions_mean_symm_batch = data_augmentation_func( - obs=None, actions=action_mean_orig, env=self.symmetry["_env"] - ) - - # compute the loss (we skip the first augmentation as it is the original one) - mse_loss = torch.nn.MSELoss() - symmetry_loss = mse_loss( - mean_actions_batch[original_batch_size:], actions_mean_symm_batch.detach()[original_batch_size:] - ) - # add the loss to the total loss - if self.symmetry["use_mirror_loss"]: - loss += self.symmetry["mirror_loss_coeff"] * symmetry_loss - else: - symmetry_loss = symmetry_loss.detach() - - # Random Network Distillation loss - # TODO: Move this processing to inside RND module. - if self.rnd: - # extract the rnd_state - # TODO: Check if we still need torch no grad. It is just an affine transformation. - with torch.no_grad(): - rnd_state_batch = self.rnd.get_rnd_state(obs_batch[:original_batch_size]) - rnd_state_batch = self.rnd.state_normalizer(rnd_state_batch) - # predict the embedding and the target - predicted_embedding = self.rnd.predictor(rnd_state_batch) - target_embedding = self.rnd.target(rnd_state_batch).detach() - # compute the loss as the mean squared error - mseloss = torch.nn.MSELoss() - rnd_loss = mseloss(predicted_embedding, target_embedding) - - expert_states = sample_amp_expert - policy_states = sample_amp_policy - - with torch.no_grad(): - expert_states = self.amp_normalizer.normalize_torch(expert_states.to(self.device), self.device) - policy_states = self.amp_normalizer.normalize_torch(policy_states, self.device) - - contact_phase_push = obs_batch['critic'][:, -4] - mask_push = contact_phase_push == 1. - - if mask_push.any(): - policy_d = self.discriminator(policy_states.flatten(1)) - expert_states = expert_states.to(self.device) - expert_d = self.discriminator(expert_states.flatten(1)) - - expert_loss = torch.nn.MSELoss()(expert_d, torch.ones(expert_d.size(), device=self.device)) - policy_loss = torch.nn.MSELoss()(policy_d, -1 * torch.ones(policy_d.size(), device=self.device)) - amp_loss = 0.5 * (expert_loss + policy_loss) - - # grad penalty - grad_pen_loss = self.discriminator.compute_grad_pen(expert_states, lambda_=5) - else: - amp_loss = torch.tensor(0.0, device=self.device) - grad_pen_loss = torch.tensor(0.0, device=self.device) - expert_loss = torch.tensor(0.0, device=self.device) - policy_loss = torch.tensor(0.0, device=self.device) - - loss += (amp_loss + grad_pen_loss) - self.amp_normalizer.update(policy_states.cpu().numpy()) - self.amp_normalizer.update(expert_states.cpu().numpy()) - - # Compute the gradients - # -- For PPO - self.optimizer.zero_grad() - loss.backward() - # -- For RND - if self.rnd: - self.rnd_optimizer.zero_grad() # type: ignore - rnd_loss.backward() - - # Collect gradients from all GPUs - if self.is_multi_gpu: - self.reduce_parameters() - - # Apply the gradients - # -- For PPO - nn.utils.clip_grad_norm_(self.policy.parameters(), self.max_grad_norm) - self.optimizer.step() - # -- For RND - if self.rnd_optimizer: - self.rnd_optimizer.step() - - # Store the losses - mean_value_loss += value_loss.item() - mean_surrogate_loss += surrogate_loss.item() - mean_entropy += entropy_batch.mean().item() - mean_amp_loss += amp_loss.item() - mean_grad_pen_loss += grad_pen_loss.item() - mean_policy_pred += policy_loss.mean().item() - mean_expert_pred += expert_loss.mean().item() - # -- RND loss - if mean_rnd_loss is not None: - mean_rnd_loss += rnd_loss.item() - # -- Symmetry loss - if mean_symmetry_loss is not None: - mean_symmetry_loss += symmetry_loss.item() - - # -- For PPO - num_updates = self.num_learning_epochs * self.num_mini_batches - mean_value_loss /= num_updates - mean_surrogate_loss /= num_updates - mean_entropy /= num_updates - mean_amp_loss /= num_updates - mean_grad_pen_loss /= num_updates - mean_policy_pred /= num_updates - mean_expert_pred /= num_updates - - # -- For RND - if mean_rnd_loss is not None: - mean_rnd_loss /= num_updates - # -- For Symmetry - if mean_symmetry_loss is not None: - mean_symmetry_loss /= num_updates - # -- Clear the storage - self.storage.clear() - - # construct the loss dictionary - loss_dict = { - "value_function": mean_value_loss, - "surrogate": mean_surrogate_loss, - "entropy": mean_entropy, - "amp": mean_amp_loss, - "amp_grad_pen": mean_grad_pen_loss, - "amp_policy_pred": mean_policy_pred, - "amp_expert_pred": mean_expert_pred, - } - if self.rnd: - loss_dict["rnd"] = mean_rnd_loss - if self.symmetry: - loss_dict["symmetry"] = mean_symmetry_loss - - return loss_dict - - """ - Helper functions - """ - - def broadcast_parameters(self): - """Broadcast model parameters to all GPUs.""" - # obtain the model parameters on current GPU - model_params = [self.policy.state_dict()] - if self.rnd: - model_params.append(self.rnd.predictor.state_dict()) - # broadcast the model parameters - torch.distributed.broadcast_object_list(model_params, src=0) - # load the model parameters on all GPUs from source GPU - self.policy.load_state_dict(model_params[0]) - if self.rnd: - self.rnd.predictor.load_state_dict(model_params[1]) - - def reduce_parameters(self): - """Collect gradients from all GPUs and average them. - - This function is called after the backward pass to synchronize the gradients across all GPUs. - """ - # Create a tensor to store the gradients - grads = [param.grad.view(-1) for param in self.policy.parameters() if param.grad is not None] - if self.rnd: - grads += [param.grad.view(-1) for param in self.rnd.parameters() if param.grad is not None] - all_grads = torch.cat(grads) - - # Average the gradients across all GPUs - torch.distributed.all_reduce(all_grads, op=torch.distributed.ReduceOp.SUM) - all_grads /= self.gpu_world_size - - # Get all parameters - all_params = self.policy.parameters() - if self.rnd: - all_params = chain(all_params, self.rnd.parameters()) - - # Update the gradients for all parameters with the reduced gradients - offset = 0 - for param in all_params: - if param.grad is not None: - numel = param.numel() - # copy data back from shared buffer - param.grad.data.copy_(all_grads[offset : offset + numel].view_as(param.grad.data)) - # update the offset for the next parameter - offset += numel diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/algorithms/distillation.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/algorithms/distillation.py deleted file mode 100644 index 3a86e00..0000000 --- a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/algorithms/distillation.py +++ /dev/null @@ -1,185 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -import torch -import torch.nn as nn - -from rsl_rl.modules import StudentTeacher, StudentTeacherRecurrent -from rsl_rl.storage import RolloutStorage -from rsl_rl.utils import resolve_optimizer - - -class Distillation: - """Distillation algorithm for training a student model to mimic a teacher model.""" - - policy: StudentTeacher | StudentTeacherRecurrent - """The student teacher model.""" - - def __init__( - self, - policy, - num_learning_epochs=1, - gradient_length=15, - learning_rate=1e-3, - max_grad_norm=None, - loss_type="mse", - optimizer="adam", - device="cpu", - # Distributed training parameters - multi_gpu_cfg: dict | None = None, - ): - # device-related parameters - self.device = device - self.is_multi_gpu = multi_gpu_cfg is not None - # Multi-GPU parameters - if multi_gpu_cfg is not None: - self.gpu_global_rank = multi_gpu_cfg["global_rank"] - self.gpu_world_size = multi_gpu_cfg["world_size"] - else: - self.gpu_global_rank = 0 - self.gpu_world_size = 1 - - # distillation components - self.policy = policy - self.policy.to(self.device) - self.storage = None # initialized later - - # initialize the optimizer - self.optimizer = resolve_optimizer(optimizer)(self.policy.parameters(), lr=learning_rate) - - # initialize the transition - self.transition = RolloutStorage.Transition() - self.last_hidden_states = None - - # distillation parameters - self.num_learning_epochs = num_learning_epochs - self.gradient_length = gradient_length - self.learning_rate = learning_rate - self.max_grad_norm = max_grad_norm - - # initialize the loss function - loss_fn_dict = { - "mse": nn.functional.mse_loss, - "huber": nn.functional.huber_loss, - } - if loss_type in loss_fn_dict: - self.loss_fn = loss_fn_dict[loss_type] - else: - raise ValueError(f"Unknown loss type: {loss_type}. Supported types are: {list(loss_fn_dict.keys())}") - - self.num_updates = 0 - - def init_storage(self, training_type, num_envs, num_transitions_per_env, obs, actions_shape): - # create rollout storage - self.storage = RolloutStorage( - training_type, - num_envs, - num_transitions_per_env, - obs, - actions_shape, - self.device, - ) - - def act(self, obs): - # compute the actions - self.transition.actions = self.policy.act(obs).detach() - self.transition.privileged_actions = self.policy.evaluate(obs).detach() - # record the observations - self.transition.observations = obs - return self.transition.actions - - def process_env_step(self, obs, rewards, dones, extras): - # update the normalizers - self.policy.update_normalization(obs) - - # record the rewards and dones - self.transition.rewards = rewards - self.transition.dones = dones - # record the transition - self.storage.add_transitions(self.transition) - self.transition.clear() - self.policy.reset(dones) - - def update(self): - self.num_updates += 1 - mean_behavior_loss = 0 - loss = 0 - cnt = 0 - - for epoch in range(self.num_learning_epochs): - self.policy.reset(hidden_states=self.last_hidden_states) - self.policy.detach_hidden_states() - for obs, _, privileged_actions, dones in self.storage.generator(): - - # inference the student for gradient computation - actions = self.policy.act_inference(obs) - - # behavior cloning loss - behavior_loss = self.loss_fn(actions, privileged_actions) - - # total loss - loss = loss + behavior_loss - mean_behavior_loss += behavior_loss.item() - cnt += 1 - - # gradient step - if cnt % self.gradient_length == 0: - self.optimizer.zero_grad() - loss.backward() - if self.is_multi_gpu: - self.reduce_parameters() - if self.max_grad_norm: - nn.utils.clip_grad_norm_(self.policy.student.parameters(), self.max_grad_norm) - self.optimizer.step() - self.policy.detach_hidden_states() - loss = 0 - - # reset dones - self.policy.reset(dones.view(-1)) - self.policy.detach_hidden_states(dones.view(-1)) - - mean_behavior_loss /= cnt - self.storage.clear() - self.last_hidden_states = self.policy.get_hidden_states() - self.policy.detach_hidden_states() - - # construct the loss dictionary - loss_dict = {"behavior": mean_behavior_loss} - - return loss_dict - - """ - Helper functions - """ - - def broadcast_parameters(self): - """Broadcast model parameters to all GPUs.""" - # obtain the model parameters on current GPU - model_params = [self.policy.state_dict()] - # broadcast the model parameters - torch.distributed.broadcast_object_list(model_params, src=0) - # load the model parameters on all GPUs from source GPU - self.policy.load_state_dict(model_params[0]) - - def reduce_parameters(self): - """Collect gradients from all GPUs and average them. - - This function is called after the backward pass to synchronize the gradients across all GPUs. - """ - # Create a tensor to store the gradients - grads = [param.grad.view(-1) for param in self.policy.parameters() if param.grad is not None] - all_grads = torch.cat(grads) - # Average the gradients across all GPUs - torch.distributed.all_reduce(all_grads, op=torch.distributed.ReduceOp.SUM) - all_grads /= self.gpu_world_size - # Update the gradients for all parameters with the reduced gradients - offset = 0 - for param in self.policy.parameters(): - if param.grad is not None: - numel = param.numel() - # copy data back from shared buffer - param.grad.data.copy_(all_grads[offset : offset + numel].view_as(param.grad.data)) - # update the offset for the next parameter - offset += numel diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/algorithms/ppo.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/algorithms/ppo.py deleted file mode 100644 index 6c21fc5..0000000 --- a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/algorithms/ppo.py +++ /dev/null @@ -1,469 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -from __future__ import annotations - -import torch -import torch.nn as nn -import torch.optim as optim -from itertools import chain - -from rsl_rl.modules import ActorCritic -from rsl_rl.modules.rnd import RandomNetworkDistillation -from rsl_rl.storage import RolloutStorage -from rsl_rl.utils import string_to_callable - - -class PPO: - """Proximal Policy Optimization algorithm (https://arxiv.org/abs/1707.06347).""" - - policy: ActorCritic - """The actor critic module.""" - - def __init__( - self, - policy, - num_learning_epochs=5, - num_mini_batches=4, - clip_param=0.2, - gamma=0.99, - lam=0.95, - value_loss_coef=1.0, - entropy_coef=0.01, - learning_rate=0.001, - max_grad_norm=1.0, - use_clipped_value_loss=True, - schedule="adaptive", - desired_kl=0.01, - device="cpu", - normalize_advantage_per_mini_batch=False, - # RND parameters - rnd_cfg: dict | None = None, - # Symmetry parameters - symmetry_cfg: dict | None = None, - # Distributed training parameters - multi_gpu_cfg: dict | None = None, - ): - # device-related parameters - self.device = device - self.is_multi_gpu = multi_gpu_cfg is not None - # Multi-GPU parameters - if multi_gpu_cfg is not None: - self.gpu_global_rank = multi_gpu_cfg["global_rank"] - self.gpu_world_size = multi_gpu_cfg["world_size"] - else: - self.gpu_global_rank = 0 - self.gpu_world_size = 1 - - # RND components - if rnd_cfg is not None: - # Extract parameters used in ppo - rnd_lr = rnd_cfg.pop("learning_rate", 1e-3) - # Create RND module - self.rnd = RandomNetworkDistillation(device=self.device, **rnd_cfg) - # Create RND optimizer - params = self.rnd.predictor.parameters() - self.rnd_optimizer = optim.Adam(params, lr=rnd_lr) - else: - self.rnd = None - self.rnd_optimizer = None - - # Symmetry components - if symmetry_cfg is not None: - # Check if symmetry is enabled - use_symmetry = symmetry_cfg["use_data_augmentation"] or symmetry_cfg["use_mirror_loss"] - # Print that we are not using symmetry - if not use_symmetry: - print("Symmetry not used for learning. We will use it for logging instead.") - # If function is a string then resolve it to a function - if isinstance(symmetry_cfg["data_augmentation_func"], str): - symmetry_cfg["data_augmentation_func"] = string_to_callable(symmetry_cfg["data_augmentation_func"]) - # Check valid configuration - if symmetry_cfg["use_data_augmentation"] and not callable(symmetry_cfg["data_augmentation_func"]): - raise ValueError( - "Data augmentation enabled but the function is not callable:" - f" {symmetry_cfg['data_augmentation_func']}" - ) - # Store symmetry configuration - self.symmetry = symmetry_cfg - else: - self.symmetry = None - - # PPO components - self.policy = policy - self.policy.to(self.device) - # Create optimizer - self.optimizer = optim.Adam(self.policy.parameters(), lr=learning_rate) - # Create rollout storage - self.storage: RolloutStorage = None # type: ignore - self.transition = RolloutStorage.Transition() - - # PPO parameters - self.clip_param = clip_param - self.num_learning_epochs = num_learning_epochs - self.num_mini_batches = num_mini_batches - self.value_loss_coef = value_loss_coef - self.entropy_coef = entropy_coef - self.gamma = gamma - self.lam = lam - self.max_grad_norm = max_grad_norm - self.use_clipped_value_loss = use_clipped_value_loss - self.desired_kl = desired_kl - self.schedule = schedule - self.learning_rate = learning_rate - self.normalize_advantage_per_mini_batch = normalize_advantage_per_mini_batch - - def init_storage(self, training_type, num_envs, num_transitions_per_env, obs, actions_shape): - # create rollout storage - self.storage = RolloutStorage( - training_type, - num_envs, - num_transitions_per_env, - obs, - actions_shape, - self.device, - ) - - def act(self, obs): - if self.policy.is_recurrent: - self.transition.hidden_states = self.policy.get_hidden_states() - # compute the actions and values - self.transition.actions = self.policy.act(obs).detach() - self.transition.values = self.policy.evaluate(obs).detach() - self.transition.actions_log_prob = self.policy.get_actions_log_prob(self.transition.actions).detach() - self.transition.action_mean = self.policy.action_mean.detach() - self.transition.action_sigma = self.policy.action_std.detach() - # need to record obs before env.step() - self.transition.observations = obs - return self.transition.actions - - def process_env_step(self, obs, rewards, dones, extras): - # update the normalizers - self.policy.update_normalization(obs) - if self.rnd: - self.rnd.update_normalization(obs) - - # Record the rewards and dones - # Note: we clone here because later on we bootstrap the rewards based on timeouts - self.transition.rewards = rewards.clone() - self.transition.dones = dones - - # Compute the intrinsic rewards and add to extrinsic rewards - if self.rnd: - # Compute the intrinsic rewards - self.intrinsic_rewards = self.rnd.get_intrinsic_reward(obs) - # Add intrinsic rewards to extrinsic rewards - self.transition.rewards += self.intrinsic_rewards - - # Bootstrapping on time outs - if "time_outs" in extras: - self.transition.rewards += self.gamma * torch.squeeze( - self.transition.values * extras["time_outs"].unsqueeze(1).to(self.device), 1 - ) - - # record the transition - self.storage.add_transitions(self.transition) - self.transition.clear() - self.policy.reset(dones) - - def compute_returns(self, obs): - # compute value for the last step - last_values = self.policy.evaluate(obs).detach() - self.storage.compute_returns( - last_values, self.gamma, self.lam, normalize_advantage=not self.normalize_advantage_per_mini_batch - ) - - def update(self): # noqa: C901 - mean_value_loss = 0 - mean_surrogate_loss = 0 - mean_entropy = 0 - # -- RND loss - if self.rnd: - mean_rnd_loss = 0 - else: - mean_rnd_loss = None - # -- Symmetry loss - if self.symmetry: - mean_symmetry_loss = 0 - else: - mean_symmetry_loss = None - - # generator for mini batches - if self.policy.is_recurrent: - generator = self.storage.recurrent_mini_batch_generator(self.num_mini_batches, self.num_learning_epochs) - else: - generator = self.storage.mini_batch_generator(self.num_mini_batches, self.num_learning_epochs) - - # iterate over batches - for ( - obs_batch, - actions_batch, - target_values_batch, - advantages_batch, - returns_batch, - old_actions_log_prob_batch, - old_mu_batch, - old_sigma_batch, - hid_states_batch, - masks_batch, - ) in generator: - - # number of augmentations per sample - # we start with 1 and increase it if we use symmetry augmentation - num_aug = 1 - # original batch size - # we assume policy group is always there and needs augmentation - original_batch_size = obs_batch.batch_size[0] - - # check if we should normalize advantages per mini batch - if self.normalize_advantage_per_mini_batch: - with torch.no_grad(): - advantages_batch = (advantages_batch - advantages_batch.mean()) / (advantages_batch.std() + 1e-8) - - # Perform symmetric augmentation - if self.symmetry and self.symmetry["use_data_augmentation"]: - # augmentation using symmetry - data_augmentation_func = self.symmetry["data_augmentation_func"] - # returned shape: [batch_size * num_aug, ...] - obs_batch, actions_batch = data_augmentation_func( - obs=obs_batch, - actions=actions_batch, - env=self.symmetry["_env"], - ) - # compute number of augmentations per sample - # we assume policy group is always there and needs augmentation - num_aug = int(obs_batch.batch_size[0] / original_batch_size) - # repeat the rest of the batch - # -- actor - old_actions_log_prob_batch = old_actions_log_prob_batch.repeat(num_aug, 1) - # -- critic - target_values_batch = target_values_batch.repeat(num_aug, 1) - advantages_batch = advantages_batch.repeat(num_aug, 1) - returns_batch = returns_batch.repeat(num_aug, 1) - - # Recompute actions log prob and entropy for current batch of transitions - # Note: we need to do this because we updated the policy with the new parameters - # -- actor - self.policy.act(obs_batch, masks=masks_batch, hidden_states=hid_states_batch[0]) - actions_log_prob_batch = self.policy.get_actions_log_prob(actions_batch) - # -- critic - value_batch = self.policy.evaluate(obs_batch, masks=masks_batch, hidden_states=hid_states_batch[1]) - # -- entropy - # we only keep the entropy of the first augmentation (the original one) - mu_batch = self.policy.action_mean[:original_batch_size] - sigma_batch = self.policy.action_std[:original_batch_size] - entropy_batch = self.policy.entropy[:original_batch_size] - - # KL - if self.desired_kl is not None and self.schedule == "adaptive": - with torch.inference_mode(): - kl = torch.sum( - torch.log(sigma_batch / old_sigma_batch + 1.0e-5) - + (torch.square(old_sigma_batch) + torch.square(old_mu_batch - mu_batch)) - / (2.0 * torch.square(sigma_batch)) - - 0.5, - axis=-1, - ) - kl_mean = torch.mean(kl) - - # Reduce the KL divergence across all GPUs - if self.is_multi_gpu: - torch.distributed.all_reduce(kl_mean, op=torch.distributed.ReduceOp.SUM) - kl_mean /= self.gpu_world_size - - # Update the learning rate - # Perform this adaptation only on the main process - # TODO: Is this needed? If KL-divergence is the "same" across all GPUs, - # then the learning rate should be the same across all GPUs. - if self.gpu_global_rank == 0: - if kl_mean > self.desired_kl * 2.0: - self.learning_rate = max(1e-5, self.learning_rate / 1.5) - elif kl_mean < self.desired_kl / 2.0 and kl_mean > 0.0: - self.learning_rate = min(1e-2, self.learning_rate * 1.5) - - # Update the learning rate for all GPUs - if self.is_multi_gpu: - lr_tensor = torch.tensor(self.learning_rate, device=self.device) - torch.distributed.broadcast(lr_tensor, src=0) - self.learning_rate = lr_tensor.item() - - # Update the learning rate for all parameter groups - for param_group in self.optimizer.param_groups: - param_group["lr"] = self.learning_rate - - # Surrogate loss - ratio = torch.exp(actions_log_prob_batch - torch.squeeze(old_actions_log_prob_batch)) - surrogate = -torch.squeeze(advantages_batch) * ratio - surrogate_clipped = -torch.squeeze(advantages_batch) * torch.clamp( - ratio, 1.0 - self.clip_param, 1.0 + self.clip_param - ) - surrogate_loss = torch.max(surrogate, surrogate_clipped).mean() - - # Value function loss - if self.use_clipped_value_loss: - value_clipped = target_values_batch + (value_batch - target_values_batch).clamp( - -self.clip_param, self.clip_param - ) - value_losses = (value_batch - returns_batch).pow(2) - value_losses_clipped = (value_clipped - returns_batch).pow(2) - value_loss = torch.max(value_losses, value_losses_clipped).mean() - else: - value_loss = (returns_batch - value_batch).pow(2).mean() - - loss = surrogate_loss + self.value_loss_coef * value_loss - self.entropy_coef * entropy_batch.mean() - - # Symmetry loss - if self.symmetry: - # obtain the symmetric actions - # if we did augmentation before then we don't need to augment again - if not self.symmetry["use_data_augmentation"]: - data_augmentation_func = self.symmetry["data_augmentation_func"] - obs_batch, _ = data_augmentation_func(obs=obs_batch, actions=None, env=self.symmetry["_env"]) - # compute number of augmentations per sample - num_aug = int(obs_batch.shape[0] / original_batch_size) - - # actions predicted by the actor for symmetrically-augmented observations - mean_actions_batch = self.policy.act_inference(obs_batch.detach().clone()) - - # compute the symmetrically augmented actions - # note: we are assuming the first augmentation is the original one. - # We do not use the action_batch from earlier since that action was sampled from the distribution. - # However, the symmetry loss is computed using the mean of the distribution. - action_mean_orig = mean_actions_batch[:original_batch_size] - _, actions_mean_symm_batch = data_augmentation_func( - obs=None, actions=action_mean_orig, env=self.symmetry["_env"] - ) - - # compute the loss (we skip the first augmentation as it is the original one) - mse_loss = torch.nn.MSELoss() - symmetry_loss = mse_loss( - mean_actions_batch[original_batch_size:], actions_mean_symm_batch.detach()[original_batch_size:] - ) - # add the loss to the total loss - if self.symmetry["use_mirror_loss"]: - loss += self.symmetry["mirror_loss_coeff"] * symmetry_loss - else: - symmetry_loss = symmetry_loss.detach() - - # Random Network Distillation loss - # TODO: Move this processing to inside RND module. - if self.rnd: - # extract the rnd_state - # TODO: Check if we still need torch no grad. It is just an affine transformation. - with torch.no_grad(): - rnd_state_batch = self.rnd.get_rnd_state(obs_batch[:original_batch_size]) - rnd_state_batch = self.rnd.state_normalizer(rnd_state_batch) - # predict the embedding and the target - predicted_embedding = self.rnd.predictor(rnd_state_batch) - target_embedding = self.rnd.target(rnd_state_batch).detach() - # compute the loss as the mean squared error - mseloss = torch.nn.MSELoss() - rnd_loss = mseloss(predicted_embedding, target_embedding) - - # Compute the gradients - # -- For PPO - self.optimizer.zero_grad() - loss.backward() - # -- For RND - if self.rnd: - self.rnd_optimizer.zero_grad() # type: ignore - rnd_loss.backward() - - # Collect gradients from all GPUs - if self.is_multi_gpu: - self.reduce_parameters() - - # Apply the gradients - # -- For PPO - nn.utils.clip_grad_norm_(self.policy.parameters(), self.max_grad_norm) - self.optimizer.step() - # -- For RND - if self.rnd_optimizer: - self.rnd_optimizer.step() - - # Store the losses - mean_value_loss += value_loss.item() - mean_surrogate_loss += surrogate_loss.item() - mean_entropy += entropy_batch.mean().item() - # -- RND loss - if mean_rnd_loss is not None: - mean_rnd_loss += rnd_loss.item() - # -- Symmetry loss - if mean_symmetry_loss is not None: - mean_symmetry_loss += symmetry_loss.item() - - # -- For PPO - num_updates = self.num_learning_epochs * self.num_mini_batches - mean_value_loss /= num_updates - mean_surrogate_loss /= num_updates - mean_entropy /= num_updates - # -- For RND - if mean_rnd_loss is not None: - mean_rnd_loss /= num_updates - # -- For Symmetry - if mean_symmetry_loss is not None: - mean_symmetry_loss /= num_updates - # -- Clear the storage - self.storage.clear() - - # construct the loss dictionary - loss_dict = { - "value_function": mean_value_loss, - "surrogate": mean_surrogate_loss, - "entropy": mean_entropy, - } - if self.rnd: - loss_dict["rnd"] = mean_rnd_loss - if self.symmetry: - loss_dict["symmetry"] = mean_symmetry_loss - - return loss_dict - - """ - Helper functions - """ - - def broadcast_parameters(self): - """Broadcast model parameters to all GPUs.""" - # obtain the model parameters on current GPU - model_params = [self.policy.state_dict()] - if self.rnd: - model_params.append(self.rnd.predictor.state_dict()) - # broadcast the model parameters - torch.distributed.broadcast_object_list(model_params, src=0) - # load the model parameters on all GPUs from source GPU - self.policy.load_state_dict(model_params[0]) - if self.rnd: - self.rnd.predictor.load_state_dict(model_params[1]) - - def reduce_parameters(self): - """Collect gradients from all GPUs and average them. - - This function is called after the backward pass to synchronize the gradients across all GPUs. - """ - # Create a tensor to store the gradients - grads = [param.grad.view(-1) for param in self.policy.parameters() if param.grad is not None] - if self.rnd: - grads += [param.grad.view(-1) for param in self.rnd.parameters() if param.grad is not None] - all_grads = torch.cat(grads) - - # Average the gradients across all GPUs - torch.distributed.all_reduce(all_grads, op=torch.distributed.ReduceOp.SUM) - all_grads /= self.gpu_world_size - - # Get all parameters - all_params = self.policy.parameters() - if self.rnd: - all_params = chain(all_params, self.rnd.parameters()) - - # Update the gradients for all parameters with the reduced gradients - offset = 0 - for param in all_params: - if param.grad is not None: - numel = param.numel() - # copy data back from shared buffer - param.grad.data.copy_(all_grads[offset : offset + numel].view_as(param.grad.data)) - # update the offset for the next parameter - offset += numel diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/__init__.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/__init__.py deleted file mode 100644 index ebc2e20..0000000 --- a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -"""Main module for the rsl_rl package.""" diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/algorithms/__init__.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/algorithms/__init__.py deleted file mode 100644 index 0ef9163..0000000 --- a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/algorithms/__init__.py +++ /dev/null @@ -1,11 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -"""Implementation of different RL agents.""" - -from .distillation import Distillation -from .ppo import PPO -from .amp_ppo import AMP_PPO -__all__ = ["PPO", "Distillation", "AMP_PPO"] diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/algorithms/amp_ppo.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/algorithms/amp_ppo.py deleted file mode 100644 index a2164eb..0000000 --- a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/algorithms/amp_ppo.py +++ /dev/null @@ -1,571 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -from __future__ import annotations -from torch._tensor import Tensor -from torch._tensor import Tensor -from typing import Any -from copy import deepcopy - -import torch -import torch.nn as nn -import torch.optim as optim -from itertools import chain - -from rsl_rl.modules import ActorCritic -from rsl_rl.modules.rnd import RandomNetworkDistillation -from rsl_rl.storage import RolloutStorage, ReplayBufferMulti -from rsl_rl.utils import string_to_callable - - -class AMP_PPO: - """Proximal Policy Optimization algorithm (https://arxiv.org/abs/1707.06347).""" - - policy: ActorCritic - """The actor critic module.""" - - def __init__( - self, - policy, - discriminator, - amp_data, - amp_normalizer, - amp_num_frames=1, - amp_replay_buffer_size=100000, - num_learning_epochs=5, - num_mini_batches=4, - clip_param=0.2, - gamma=0.99, - lam=0.95, - value_loss_coef=1.0, - entropy_coef=0.01, - learning_rate=0.001, - max_grad_norm=1.0, - use_clipped_value_loss=True, - schedule="adaptive", - desired_kl=0.01, - device="cpu", - normalize_advantage_per_mini_batch=False, - # RND parameters - rnd_cfg: dict | None = None, - # Symmetry parameters - symmetry_cfg: dict | None = None, - # Distributed training parameters - multi_gpu_cfg: dict | None = None, - ): - # device-related parameters - self.device = device - self.is_multi_gpu = multi_gpu_cfg is not None - # Multi-GPU parameters - if multi_gpu_cfg is not None: - self.gpu_global_rank = multi_gpu_cfg["global_rank"] - self.gpu_world_size = multi_gpu_cfg["world_size"] - else: - self.gpu_global_rank = 0 - self.gpu_world_size = 1 - - # RND components - if rnd_cfg is not None: - # Extract parameters used in ppo - rnd_lr = rnd_cfg.pop("learning_rate", 1e-3) - # Create RND module - self.rnd = RandomNetworkDistillation(device=self.device, **rnd_cfg) - # Create RND optimizer - params = self.rnd.predictor.parameters() - self.rnd_optimizer = optim.Adam(params, lr=rnd_lr) - else: - self.rnd = None - self.rnd_optimizer = None - - # Symmetry components - if symmetry_cfg is not None: - # Check if symmetry is enabled - use_symmetry = symmetry_cfg["use_data_augmentation"] or symmetry_cfg["use_mirror_loss"] - # Print that we are not using symmetry - if not use_symmetry: - print("Symmetry not used for learning. We will use it for logging instead.") - # If function is a string then resolve it to a function - if isinstance(symmetry_cfg["data_augmentation_func"], str): - symmetry_cfg["data_augmentation_func"] = string_to_callable(symmetry_cfg["data_augmentation_func"]) - # Check valid configuration - if symmetry_cfg["use_data_augmentation"] and not callable(symmetry_cfg["data_augmentation_func"]): - raise ValueError( - "Data augmentation enabled but the function is not callable:" - f" {symmetry_cfg['data_augmentation_func']}" - ) - # Store symmetry configuration - self.symmetry = symmetry_cfg - else: - self.symmetry = None - - ## AMP components - self.discriminator = discriminator - self.discriminator.to(self.device) - - self.amp_storage = ReplayBufferMulti(discriminator.state_dim, amp_replay_buffer_size, amp_num_frames, device) - self.amp_data = amp_data - self.amp_normalizer = amp_normalizer - - # PPO components - self.policy = policy - self.policy.to(self.device) - - # Create rollout storage - self.storage: RolloutStorage = None # type: ignore - self.transition = RolloutStorage.Transition() - self.amp_transition = RolloutStorage.Transition() - params = [ - {'params': self.policy.parameters(), 'name': 'policy'}, - ] - - params.append({ - 'params': self.discriminator.trunk.parameters(), - 'weight_decay': 10e-4, - 'name': f'amp_trunk' - }) - params.append({ - 'params': self.discriminator.amp_linear.parameters(), - 'weight_decay': 10e-2, - 'name': f'amp_head' - }) - - # Create optimizer - self.optimizer = optim.Adam(params, lr=learning_rate) - - # PPO parameters - self.clip_param = clip_param - self.num_learning_epochs = num_learning_epochs - self.num_mini_batches = num_mini_batches - self.value_loss_coef = value_loss_coef - self.entropy_coef = entropy_coef - self.gamma = gamma - self.lam = lam - self.max_grad_norm = max_grad_norm - self.use_clipped_value_loss = use_clipped_value_loss - self.desired_kl = desired_kl - self.schedule = schedule - self.learning_rate = learning_rate - self.normalize_advantage_per_mini_batch = normalize_advantage_per_mini_batch - - - def init_storage(self, training_type, num_envs, num_transitions_per_env, obs, actions_shape): - # create rollout storage - self.storage = RolloutStorage( - training_type, - num_envs, - num_transitions_per_env, - obs, - actions_shape, - self.device, - ) - - def act(self, obs, amp_obs): - if self.policy.is_recurrent: - self.transition.hidden_states = self.policy.get_hidden_states() - # compute the actions and values - self.transition.actions = self.policy.act(obs).detach() - self.transition.values = self.policy.evaluate(obs).detach() - self.transition.actions_log_prob = self.policy.get_actions_log_prob(self.transition.actions).detach() - self.transition.action_mean = self.policy.action_mean.detach() - self.transition.action_sigma = self.policy.action_std.detach() - # need to record obs before env.step() - self.transition.observations = obs - self.amp_transition.observations = amp_obs - return self.transition.actions - - def process_env_step(self, obs, rewards, dones, extras,amp_obs, amp_obs_frames=None): - # update the normalizers - self.policy.update_normalization(obs) - if self.rnd: - self.rnd.update_normalization(obs) - - # Record the rewards and dones - # Note: we clone here because later on we bootstrap the rewards based on timeouts - self.transition.rewards = rewards.clone() - self.transition.dones = dones - - # Compute the intrinsic rewards and add to extrinsic rewards - if self.rnd: - # Compute the intrinsic rewards - self.intrinsic_rewards = self.rnd.get_intrinsic_reward(obs) - # Add intrinsic rewards to extrinsic rewards - self.transition.rewards += self.intrinsic_rewards - - # Bootstrapping on time outs - if "time_outs" in extras: - self.transition.rewards += self.gamma * torch.squeeze( - self.transition.values * extras["time_outs"].unsqueeze(1).to(self.device), 1 - ) - - if amp_obs_frames is not None: - self.amp_storage.insert(amp_obs_frames) - else: - self.amp_storage.insert(self.amp_transition.observations, amp_obs) - - # record the transition - self.storage.add_transitions(self.transition) - self.transition.clear() - self.amp_transition.clear() - self.policy.reset(dones) - - def compute_returns(self, obs): - # compute value for the last step - last_values = self.policy.evaluate(obs).detach() - self.storage.compute_returns( - last_values, self.gamma, self.lam, normalize_advantage=not self.normalize_advantage_per_mini_batch - ) - - def update(self): # noqa: C901 - mean_value_loss = 0 - mean_surrogate_loss = 0 - mean_entropy = 0 - mean_amp_loss = 0 - mean_grad_pen_loss = 0 - mean_policy_pred = 0 - mean_expert_pred = 0 - # -- RND loss - if self.rnd: - mean_rnd_loss = 0 - else: - mean_rnd_loss = None - # -- Symmetry loss - if self.symmetry: - mean_symmetry_loss = 0 - else: - mean_symmetry_loss = None - - # generator for mini batches - if self.policy.is_recurrent: - generator = self.storage.recurrent_mini_batch_generator(self.num_mini_batches, self.num_learning_epochs) - else: - generator = self.storage.mini_batch_generator(self.num_mini_batches, self.num_learning_epochs) - - - amp_policy_generator = self.amp_storage.feed_forward_generator( - self.num_learning_epochs * self.num_mini_batches, - self.storage.num_envs * self.storage.num_transitions_per_env // self.num_mini_batches, - ) - - amp_expert_generator = self.amp_data.feed_forward_generator_23dof_multi( - self.num_learning_epochs * self.num_mini_batches, - self.storage.num_envs * self.storage.num_transitions_per_env // self.num_mini_batches, - ) - - # iterate over batches - for sample, sample_amp_policy, sample_amp_expert in zip(generator, amp_policy_generator, amp_expert_generator): - ( - obs_batch, - actions_batch, - target_values_batch, - advantages_batch, - returns_batch, - old_actions_log_prob_batch, - old_mu_batch, - old_sigma_batch, - hid_states_batch, - masks_batch, - ) = sample - - # number of augmentations per sample - # we start with 1 and increase it if we use symmetry augmentation - num_aug = 1 - # original batch size - # we assume policy group is always there and needs augmentation - original_batch_size = obs_batch.batch_size[0] - - # check if we should normalize advantages per mini batch - if self.normalize_advantage_per_mini_batch: - with torch.no_grad(): - advantages_batch = (advantages_batch - advantages_batch.mean()) / (advantages_batch.std() + 1e-8) - - # Perform symmetric augmentation - if self.symmetry and self.symmetry["use_data_augmentation"]: - # augmentation using symmetry - data_augmentation_func = self.symmetry["data_augmentation_func"] - # returned shape: [batch_size * num_aug, ...] - obs_batch, actions_batch = data_augmentation_func( - obs=obs_batch, - actions=actions_batch, - env=self.symmetry["_env"], - ) - # compute number of augmentations per sample - # we assume policy group is always there and needs augmentation - num_aug = int(obs_batch.batch_size[0] / original_batch_size) - # repeat the rest of the batch - # -- actor - old_actions_log_prob_batch = old_actions_log_prob_batch.repeat(num_aug, 1) - # -- critic - target_values_batch = target_values_batch.repeat(num_aug, 1) - advantages_batch = advantages_batch.repeat(num_aug, 1) - returns_batch = returns_batch.repeat(num_aug, 1) - - # Recompute actions log prob and entropy for current batch of transitions - # Note: we need to do this because we updated the policy with the new parameters - # -- actor - self.policy.act(obs_batch, masks=masks_batch, hidden_states=hid_states_batch[0]) - actions_log_prob_batch = self.policy.get_actions_log_prob(actions_batch) - # -- critic - value_batch = self.policy.evaluate(obs_batch, masks=masks_batch, hidden_states=hid_states_batch[1]) - # -- entropy - # we only keep the entropy of the first augmentation (the original one) - mu_batch = self.policy.action_mean[:original_batch_size] - sigma_batch = self.policy.action_std[:original_batch_size] - entropy_batch = self.policy.entropy[:original_batch_size] - - # KL - if self.desired_kl is not None and self.schedule == "adaptive": - with torch.inference_mode(): - kl = torch.sum( - torch.log(sigma_batch / old_sigma_batch + 1.0e-5) - + (torch.square(old_sigma_batch) + torch.square(old_mu_batch - mu_batch)) - / (2.0 * torch.square(sigma_batch)) - - 0.5, - axis=-1, - ) - kl_mean = torch.mean(kl) - - # Reduce the KL divergence across all GPUs - if self.is_multi_gpu: - torch.distributed.all_reduce(kl_mean, op=torch.distributed.ReduceOp.SUM) - kl_mean /= self.gpu_world_size - - # Update the learning rate - # Perform this adaptation only on the main process - # TODO: Is this needed? If KL-divergence is the "same" across all GPUs, - # then the learning rate should be the same across all GPUs. - if self.gpu_global_rank == 0: - if kl_mean > self.desired_kl * 2.0: - self.learning_rate = max(1e-5, self.learning_rate / 1.5) - elif kl_mean < self.desired_kl / 2.0 and kl_mean > 0.0: - self.learning_rate = min(1e-2, self.learning_rate * 1.5) - - # Update the learning rate for all GPUs - if self.is_multi_gpu: - lr_tensor = torch.tensor(self.learning_rate, device=self.device) - torch.distributed.broadcast(lr_tensor, src=0) - self.learning_rate = lr_tensor.item() - - # Update the learning rate for all parameter groups - for param_group in self.optimizer.param_groups: - param_group["lr"] = self.learning_rate - - # Surrogate loss - ratio = torch.exp(actions_log_prob_batch - torch.squeeze(old_actions_log_prob_batch)) - surrogate = -torch.squeeze(advantages_batch) * ratio - surrogate_clipped = -torch.squeeze(advantages_batch) * torch.clamp( - ratio, 1.0 - self.clip_param, 1.0 + self.clip_param - ) - surrogate_loss = torch.max(surrogate, surrogate_clipped).mean() - - # Value function loss - if self.use_clipped_value_loss: - value_clipped = target_values_batch + (value_batch - target_values_batch).clamp( - -self.clip_param, self.clip_param - ) - value_losses = (value_batch - returns_batch).pow(2) - value_losses_clipped = (value_clipped - returns_batch).pow(2) - value_loss = torch.max(value_losses, value_losses_clipped).mean() - else: - value_loss = (returns_batch - value_batch).pow(2).mean() - - loss = surrogate_loss + self.value_loss_coef * value_loss - self.entropy_coef * entropy_batch.mean() - - # Symmetry loss - if self.symmetry: - # obtain the symmetric actions - # if we did augmentation before then we don't need to augment again - if not self.symmetry["use_data_augmentation"]: - data_augmentation_func = self.symmetry["data_augmentation_func"] - obs_batch, _ = data_augmentation_func(obs=obs_batch, actions=None, env=self.symmetry["_env"]) - # compute number of augmentations per sample - num_aug = int(obs_batch.shape[0] / original_batch_size) - - # actions predicted by the actor for symmetrically-augmented observations - mean_actions_batch = self.policy.act_inference(obs_batch.detach().clone()) - - # compute the symmetrically augmented actions - # note: we are assuming the first augmentation is the original one. - # We do not use the action_batch from earlier since that action was sampled from the distribution. - # However, the symmetry loss is computed using the mean of the distribution. - action_mean_orig = mean_actions_batch[:original_batch_size] - _, actions_mean_symm_batch = data_augmentation_func( - obs=None, actions=action_mean_orig, env=self.symmetry["_env"] - ) - - # compute the loss (we skip the first augmentation as it is the original one) - mse_loss = torch.nn.MSELoss() - symmetry_loss = mse_loss( - mean_actions_batch[original_batch_size:], actions_mean_symm_batch.detach()[original_batch_size:] - ) - # add the loss to the total loss - if self.symmetry["use_mirror_loss"]: - loss += self.symmetry["mirror_loss_coeff"] * symmetry_loss - else: - symmetry_loss = symmetry_loss.detach() - - # Random Network Distillation loss - # TODO: Move this processing to inside RND module. - if self.rnd: - # extract the rnd_state - # TODO: Check if we still need torch no grad. It is just an affine transformation. - with torch.no_grad(): - rnd_state_batch = self.rnd.get_rnd_state(obs_batch[:original_batch_size]) - rnd_state_batch = self.rnd.state_normalizer(rnd_state_batch) - # predict the embedding and the target - predicted_embedding = self.rnd.predictor(rnd_state_batch) - target_embedding = self.rnd.target(rnd_state_batch).detach() - # compute the loss as the mean squared error - mseloss = torch.nn.MSELoss() - rnd_loss = mseloss(predicted_embedding, target_embedding) - - expert_states = sample_amp_expert - policy_states = sample_amp_policy - - with torch.no_grad(): - expert_states = self.amp_normalizer.normalize_torch(expert_states.to(self.device), self.device) - policy_states = self.amp_normalizer.normalize_torch(policy_states, self.device) - - contact_phase_push = obs_batch['critic'][:, -4] - mask_push = contact_phase_push == 1. - - if mask_push.any(): - policy_d = self.discriminator(policy_states.flatten(1)) - expert_states = expert_states.to(self.device) - expert_d = self.discriminator(expert_states.flatten(1)) - - expert_loss = torch.nn.MSELoss()(expert_d, torch.ones(expert_d.size(), device=self.device)) - policy_loss = torch.nn.MSELoss()(policy_d, -1 * torch.ones(policy_d.size(), device=self.device)) - amp_loss = 0.5 * (expert_loss + policy_loss) - - # grad penalty - grad_pen_loss = self.discriminator.compute_grad_pen(expert_states, lambda_=5) - else: - amp_loss = torch.tensor(0.0, device=self.device) - grad_pen_loss = torch.tensor(0.0, device=self.device) - expert_loss = torch.tensor(0.0, device=self.device) - policy_loss = torch.tensor(0.0, device=self.device) - - loss += (amp_loss + grad_pen_loss) - self.amp_normalizer.update(policy_states.cpu().numpy()) - self.amp_normalizer.update(expert_states.cpu().numpy()) - - # Compute the gradients - # -- For PPO - self.optimizer.zero_grad() - loss.backward() - # -- For RND - if self.rnd: - self.rnd_optimizer.zero_grad() # type: ignore - rnd_loss.backward() - - # Collect gradients from all GPUs - if self.is_multi_gpu: - self.reduce_parameters() - - # Apply the gradients - # -- For PPO - nn.utils.clip_grad_norm_(self.policy.parameters(), self.max_grad_norm) - self.optimizer.step() - # -- For RND - if self.rnd_optimizer: - self.rnd_optimizer.step() - - # Store the losses - mean_value_loss += value_loss.item() - mean_surrogate_loss += surrogate_loss.item() - mean_entropy += entropy_batch.mean().item() - mean_amp_loss += amp_loss.item() - mean_grad_pen_loss += grad_pen_loss.item() - mean_policy_pred += policy_loss.mean().item() - mean_expert_pred += expert_loss.mean().item() - # -- RND loss - if mean_rnd_loss is not None: - mean_rnd_loss += rnd_loss.item() - # -- Symmetry loss - if mean_symmetry_loss is not None: - mean_symmetry_loss += symmetry_loss.item() - - # -- For PPO - num_updates = self.num_learning_epochs * self.num_mini_batches - mean_value_loss /= num_updates - mean_surrogate_loss /= num_updates - mean_entropy /= num_updates - mean_amp_loss /= num_updates - mean_grad_pen_loss /= num_updates - mean_policy_pred /= num_updates - mean_expert_pred /= num_updates - - # -- For RND - if mean_rnd_loss is not None: - mean_rnd_loss /= num_updates - # -- For Symmetry - if mean_symmetry_loss is not None: - mean_symmetry_loss /= num_updates - # -- Clear the storage - self.storage.clear() - - # construct the loss dictionary - loss_dict = { - "value_function": mean_value_loss, - "surrogate": mean_surrogate_loss, - "entropy": mean_entropy, - "amp": mean_amp_loss, - "amp_grad_pen": mean_grad_pen_loss, - "amp_policy_pred": mean_policy_pred, - "amp_expert_pred": mean_expert_pred, - } - if self.rnd: - loss_dict["rnd"] = mean_rnd_loss - if self.symmetry: - loss_dict["symmetry"] = mean_symmetry_loss - - return loss_dict - - """ - Helper functions - """ - - def broadcast_parameters(self): - """Broadcast model parameters to all GPUs.""" - # obtain the model parameters on current GPU - model_params = [self.policy.state_dict()] - if self.rnd: - model_params.append(self.rnd.predictor.state_dict()) - # broadcast the model parameters - torch.distributed.broadcast_object_list(model_params, src=0) - # load the model parameters on all GPUs from source GPU - self.policy.load_state_dict(model_params[0]) - if self.rnd: - self.rnd.predictor.load_state_dict(model_params[1]) - - def reduce_parameters(self): - """Collect gradients from all GPUs and average them. - - This function is called after the backward pass to synchronize the gradients across all GPUs. - """ - # Create a tensor to store the gradients - grads = [param.grad.view(-1) for param in self.policy.parameters() if param.grad is not None] - if self.rnd: - grads += [param.grad.view(-1) for param in self.rnd.parameters() if param.grad is not None] - all_grads = torch.cat(grads) - - # Average the gradients across all GPUs - torch.distributed.all_reduce(all_grads, op=torch.distributed.ReduceOp.SUM) - all_grads /= self.gpu_world_size - - # Get all parameters - all_params = self.policy.parameters() - if self.rnd: - all_params = chain(all_params, self.rnd.parameters()) - - # Update the gradients for all parameters with the reduced gradients - offset = 0 - for param in all_params: - if param.grad is not None: - numel = param.numel() - # copy data back from shared buffer - param.grad.data.copy_(all_grads[offset : offset + numel].view_as(param.grad.data)) - # update the offset for the next parameter - offset += numel diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/algorithms/distillation.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/algorithms/distillation.py deleted file mode 100644 index 3a86e00..0000000 --- a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/algorithms/distillation.py +++ /dev/null @@ -1,185 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -import torch -import torch.nn as nn - -from rsl_rl.modules import StudentTeacher, StudentTeacherRecurrent -from rsl_rl.storage import RolloutStorage -from rsl_rl.utils import resolve_optimizer - - -class Distillation: - """Distillation algorithm for training a student model to mimic a teacher model.""" - - policy: StudentTeacher | StudentTeacherRecurrent - """The student teacher model.""" - - def __init__( - self, - policy, - num_learning_epochs=1, - gradient_length=15, - learning_rate=1e-3, - max_grad_norm=None, - loss_type="mse", - optimizer="adam", - device="cpu", - # Distributed training parameters - multi_gpu_cfg: dict | None = None, - ): - # device-related parameters - self.device = device - self.is_multi_gpu = multi_gpu_cfg is not None - # Multi-GPU parameters - if multi_gpu_cfg is not None: - self.gpu_global_rank = multi_gpu_cfg["global_rank"] - self.gpu_world_size = multi_gpu_cfg["world_size"] - else: - self.gpu_global_rank = 0 - self.gpu_world_size = 1 - - # distillation components - self.policy = policy - self.policy.to(self.device) - self.storage = None # initialized later - - # initialize the optimizer - self.optimizer = resolve_optimizer(optimizer)(self.policy.parameters(), lr=learning_rate) - - # initialize the transition - self.transition = RolloutStorage.Transition() - self.last_hidden_states = None - - # distillation parameters - self.num_learning_epochs = num_learning_epochs - self.gradient_length = gradient_length - self.learning_rate = learning_rate - self.max_grad_norm = max_grad_norm - - # initialize the loss function - loss_fn_dict = { - "mse": nn.functional.mse_loss, - "huber": nn.functional.huber_loss, - } - if loss_type in loss_fn_dict: - self.loss_fn = loss_fn_dict[loss_type] - else: - raise ValueError(f"Unknown loss type: {loss_type}. Supported types are: {list(loss_fn_dict.keys())}") - - self.num_updates = 0 - - def init_storage(self, training_type, num_envs, num_transitions_per_env, obs, actions_shape): - # create rollout storage - self.storage = RolloutStorage( - training_type, - num_envs, - num_transitions_per_env, - obs, - actions_shape, - self.device, - ) - - def act(self, obs): - # compute the actions - self.transition.actions = self.policy.act(obs).detach() - self.transition.privileged_actions = self.policy.evaluate(obs).detach() - # record the observations - self.transition.observations = obs - return self.transition.actions - - def process_env_step(self, obs, rewards, dones, extras): - # update the normalizers - self.policy.update_normalization(obs) - - # record the rewards and dones - self.transition.rewards = rewards - self.transition.dones = dones - # record the transition - self.storage.add_transitions(self.transition) - self.transition.clear() - self.policy.reset(dones) - - def update(self): - self.num_updates += 1 - mean_behavior_loss = 0 - loss = 0 - cnt = 0 - - for epoch in range(self.num_learning_epochs): - self.policy.reset(hidden_states=self.last_hidden_states) - self.policy.detach_hidden_states() - for obs, _, privileged_actions, dones in self.storage.generator(): - - # inference the student for gradient computation - actions = self.policy.act_inference(obs) - - # behavior cloning loss - behavior_loss = self.loss_fn(actions, privileged_actions) - - # total loss - loss = loss + behavior_loss - mean_behavior_loss += behavior_loss.item() - cnt += 1 - - # gradient step - if cnt % self.gradient_length == 0: - self.optimizer.zero_grad() - loss.backward() - if self.is_multi_gpu: - self.reduce_parameters() - if self.max_grad_norm: - nn.utils.clip_grad_norm_(self.policy.student.parameters(), self.max_grad_norm) - self.optimizer.step() - self.policy.detach_hidden_states() - loss = 0 - - # reset dones - self.policy.reset(dones.view(-1)) - self.policy.detach_hidden_states(dones.view(-1)) - - mean_behavior_loss /= cnt - self.storage.clear() - self.last_hidden_states = self.policy.get_hidden_states() - self.policy.detach_hidden_states() - - # construct the loss dictionary - loss_dict = {"behavior": mean_behavior_loss} - - return loss_dict - - """ - Helper functions - """ - - def broadcast_parameters(self): - """Broadcast model parameters to all GPUs.""" - # obtain the model parameters on current GPU - model_params = [self.policy.state_dict()] - # broadcast the model parameters - torch.distributed.broadcast_object_list(model_params, src=0) - # load the model parameters on all GPUs from source GPU - self.policy.load_state_dict(model_params[0]) - - def reduce_parameters(self): - """Collect gradients from all GPUs and average them. - - This function is called after the backward pass to synchronize the gradients across all GPUs. - """ - # Create a tensor to store the gradients - grads = [param.grad.view(-1) for param in self.policy.parameters() if param.grad is not None] - all_grads = torch.cat(grads) - # Average the gradients across all GPUs - torch.distributed.all_reduce(all_grads, op=torch.distributed.ReduceOp.SUM) - all_grads /= self.gpu_world_size - # Update the gradients for all parameters with the reduced gradients - offset = 0 - for param in self.policy.parameters(): - if param.grad is not None: - numel = param.numel() - # copy data back from shared buffer - param.grad.data.copy_(all_grads[offset : offset + numel].view_as(param.grad.data)) - # update the offset for the next parameter - offset += numel diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/algorithms/ppo.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/algorithms/ppo.py deleted file mode 100644 index 6c21fc5..0000000 --- a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/algorithms/ppo.py +++ /dev/null @@ -1,469 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -from __future__ import annotations - -import torch -import torch.nn as nn -import torch.optim as optim -from itertools import chain - -from rsl_rl.modules import ActorCritic -from rsl_rl.modules.rnd import RandomNetworkDistillation -from rsl_rl.storage import RolloutStorage -from rsl_rl.utils import string_to_callable - - -class PPO: - """Proximal Policy Optimization algorithm (https://arxiv.org/abs/1707.06347).""" - - policy: ActorCritic - """The actor critic module.""" - - def __init__( - self, - policy, - num_learning_epochs=5, - num_mini_batches=4, - clip_param=0.2, - gamma=0.99, - lam=0.95, - value_loss_coef=1.0, - entropy_coef=0.01, - learning_rate=0.001, - max_grad_norm=1.0, - use_clipped_value_loss=True, - schedule="adaptive", - desired_kl=0.01, - device="cpu", - normalize_advantage_per_mini_batch=False, - # RND parameters - rnd_cfg: dict | None = None, - # Symmetry parameters - symmetry_cfg: dict | None = None, - # Distributed training parameters - multi_gpu_cfg: dict | None = None, - ): - # device-related parameters - self.device = device - self.is_multi_gpu = multi_gpu_cfg is not None - # Multi-GPU parameters - if multi_gpu_cfg is not None: - self.gpu_global_rank = multi_gpu_cfg["global_rank"] - self.gpu_world_size = multi_gpu_cfg["world_size"] - else: - self.gpu_global_rank = 0 - self.gpu_world_size = 1 - - # RND components - if rnd_cfg is not None: - # Extract parameters used in ppo - rnd_lr = rnd_cfg.pop("learning_rate", 1e-3) - # Create RND module - self.rnd = RandomNetworkDistillation(device=self.device, **rnd_cfg) - # Create RND optimizer - params = self.rnd.predictor.parameters() - self.rnd_optimizer = optim.Adam(params, lr=rnd_lr) - else: - self.rnd = None - self.rnd_optimizer = None - - # Symmetry components - if symmetry_cfg is not None: - # Check if symmetry is enabled - use_symmetry = symmetry_cfg["use_data_augmentation"] or symmetry_cfg["use_mirror_loss"] - # Print that we are not using symmetry - if not use_symmetry: - print("Symmetry not used for learning. We will use it for logging instead.") - # If function is a string then resolve it to a function - if isinstance(symmetry_cfg["data_augmentation_func"], str): - symmetry_cfg["data_augmentation_func"] = string_to_callable(symmetry_cfg["data_augmentation_func"]) - # Check valid configuration - if symmetry_cfg["use_data_augmentation"] and not callable(symmetry_cfg["data_augmentation_func"]): - raise ValueError( - "Data augmentation enabled but the function is not callable:" - f" {symmetry_cfg['data_augmentation_func']}" - ) - # Store symmetry configuration - self.symmetry = symmetry_cfg - else: - self.symmetry = None - - # PPO components - self.policy = policy - self.policy.to(self.device) - # Create optimizer - self.optimizer = optim.Adam(self.policy.parameters(), lr=learning_rate) - # Create rollout storage - self.storage: RolloutStorage = None # type: ignore - self.transition = RolloutStorage.Transition() - - # PPO parameters - self.clip_param = clip_param - self.num_learning_epochs = num_learning_epochs - self.num_mini_batches = num_mini_batches - self.value_loss_coef = value_loss_coef - self.entropy_coef = entropy_coef - self.gamma = gamma - self.lam = lam - self.max_grad_norm = max_grad_norm - self.use_clipped_value_loss = use_clipped_value_loss - self.desired_kl = desired_kl - self.schedule = schedule - self.learning_rate = learning_rate - self.normalize_advantage_per_mini_batch = normalize_advantage_per_mini_batch - - def init_storage(self, training_type, num_envs, num_transitions_per_env, obs, actions_shape): - # create rollout storage - self.storage = RolloutStorage( - training_type, - num_envs, - num_transitions_per_env, - obs, - actions_shape, - self.device, - ) - - def act(self, obs): - if self.policy.is_recurrent: - self.transition.hidden_states = self.policy.get_hidden_states() - # compute the actions and values - self.transition.actions = self.policy.act(obs).detach() - self.transition.values = self.policy.evaluate(obs).detach() - self.transition.actions_log_prob = self.policy.get_actions_log_prob(self.transition.actions).detach() - self.transition.action_mean = self.policy.action_mean.detach() - self.transition.action_sigma = self.policy.action_std.detach() - # need to record obs before env.step() - self.transition.observations = obs - return self.transition.actions - - def process_env_step(self, obs, rewards, dones, extras): - # update the normalizers - self.policy.update_normalization(obs) - if self.rnd: - self.rnd.update_normalization(obs) - - # Record the rewards and dones - # Note: we clone here because later on we bootstrap the rewards based on timeouts - self.transition.rewards = rewards.clone() - self.transition.dones = dones - - # Compute the intrinsic rewards and add to extrinsic rewards - if self.rnd: - # Compute the intrinsic rewards - self.intrinsic_rewards = self.rnd.get_intrinsic_reward(obs) - # Add intrinsic rewards to extrinsic rewards - self.transition.rewards += self.intrinsic_rewards - - # Bootstrapping on time outs - if "time_outs" in extras: - self.transition.rewards += self.gamma * torch.squeeze( - self.transition.values * extras["time_outs"].unsqueeze(1).to(self.device), 1 - ) - - # record the transition - self.storage.add_transitions(self.transition) - self.transition.clear() - self.policy.reset(dones) - - def compute_returns(self, obs): - # compute value for the last step - last_values = self.policy.evaluate(obs).detach() - self.storage.compute_returns( - last_values, self.gamma, self.lam, normalize_advantage=not self.normalize_advantage_per_mini_batch - ) - - def update(self): # noqa: C901 - mean_value_loss = 0 - mean_surrogate_loss = 0 - mean_entropy = 0 - # -- RND loss - if self.rnd: - mean_rnd_loss = 0 - else: - mean_rnd_loss = None - # -- Symmetry loss - if self.symmetry: - mean_symmetry_loss = 0 - else: - mean_symmetry_loss = None - - # generator for mini batches - if self.policy.is_recurrent: - generator = self.storage.recurrent_mini_batch_generator(self.num_mini_batches, self.num_learning_epochs) - else: - generator = self.storage.mini_batch_generator(self.num_mini_batches, self.num_learning_epochs) - - # iterate over batches - for ( - obs_batch, - actions_batch, - target_values_batch, - advantages_batch, - returns_batch, - old_actions_log_prob_batch, - old_mu_batch, - old_sigma_batch, - hid_states_batch, - masks_batch, - ) in generator: - - # number of augmentations per sample - # we start with 1 and increase it if we use symmetry augmentation - num_aug = 1 - # original batch size - # we assume policy group is always there and needs augmentation - original_batch_size = obs_batch.batch_size[0] - - # check if we should normalize advantages per mini batch - if self.normalize_advantage_per_mini_batch: - with torch.no_grad(): - advantages_batch = (advantages_batch - advantages_batch.mean()) / (advantages_batch.std() + 1e-8) - - # Perform symmetric augmentation - if self.symmetry and self.symmetry["use_data_augmentation"]: - # augmentation using symmetry - data_augmentation_func = self.symmetry["data_augmentation_func"] - # returned shape: [batch_size * num_aug, ...] - obs_batch, actions_batch = data_augmentation_func( - obs=obs_batch, - actions=actions_batch, - env=self.symmetry["_env"], - ) - # compute number of augmentations per sample - # we assume policy group is always there and needs augmentation - num_aug = int(obs_batch.batch_size[0] / original_batch_size) - # repeat the rest of the batch - # -- actor - old_actions_log_prob_batch = old_actions_log_prob_batch.repeat(num_aug, 1) - # -- critic - target_values_batch = target_values_batch.repeat(num_aug, 1) - advantages_batch = advantages_batch.repeat(num_aug, 1) - returns_batch = returns_batch.repeat(num_aug, 1) - - # Recompute actions log prob and entropy for current batch of transitions - # Note: we need to do this because we updated the policy with the new parameters - # -- actor - self.policy.act(obs_batch, masks=masks_batch, hidden_states=hid_states_batch[0]) - actions_log_prob_batch = self.policy.get_actions_log_prob(actions_batch) - # -- critic - value_batch = self.policy.evaluate(obs_batch, masks=masks_batch, hidden_states=hid_states_batch[1]) - # -- entropy - # we only keep the entropy of the first augmentation (the original one) - mu_batch = self.policy.action_mean[:original_batch_size] - sigma_batch = self.policy.action_std[:original_batch_size] - entropy_batch = self.policy.entropy[:original_batch_size] - - # KL - if self.desired_kl is not None and self.schedule == "adaptive": - with torch.inference_mode(): - kl = torch.sum( - torch.log(sigma_batch / old_sigma_batch + 1.0e-5) - + (torch.square(old_sigma_batch) + torch.square(old_mu_batch - mu_batch)) - / (2.0 * torch.square(sigma_batch)) - - 0.5, - axis=-1, - ) - kl_mean = torch.mean(kl) - - # Reduce the KL divergence across all GPUs - if self.is_multi_gpu: - torch.distributed.all_reduce(kl_mean, op=torch.distributed.ReduceOp.SUM) - kl_mean /= self.gpu_world_size - - # Update the learning rate - # Perform this adaptation only on the main process - # TODO: Is this needed? If KL-divergence is the "same" across all GPUs, - # then the learning rate should be the same across all GPUs. - if self.gpu_global_rank == 0: - if kl_mean > self.desired_kl * 2.0: - self.learning_rate = max(1e-5, self.learning_rate / 1.5) - elif kl_mean < self.desired_kl / 2.0 and kl_mean > 0.0: - self.learning_rate = min(1e-2, self.learning_rate * 1.5) - - # Update the learning rate for all GPUs - if self.is_multi_gpu: - lr_tensor = torch.tensor(self.learning_rate, device=self.device) - torch.distributed.broadcast(lr_tensor, src=0) - self.learning_rate = lr_tensor.item() - - # Update the learning rate for all parameter groups - for param_group in self.optimizer.param_groups: - param_group["lr"] = self.learning_rate - - # Surrogate loss - ratio = torch.exp(actions_log_prob_batch - torch.squeeze(old_actions_log_prob_batch)) - surrogate = -torch.squeeze(advantages_batch) * ratio - surrogate_clipped = -torch.squeeze(advantages_batch) * torch.clamp( - ratio, 1.0 - self.clip_param, 1.0 + self.clip_param - ) - surrogate_loss = torch.max(surrogate, surrogate_clipped).mean() - - # Value function loss - if self.use_clipped_value_loss: - value_clipped = target_values_batch + (value_batch - target_values_batch).clamp( - -self.clip_param, self.clip_param - ) - value_losses = (value_batch - returns_batch).pow(2) - value_losses_clipped = (value_clipped - returns_batch).pow(2) - value_loss = torch.max(value_losses, value_losses_clipped).mean() - else: - value_loss = (returns_batch - value_batch).pow(2).mean() - - loss = surrogate_loss + self.value_loss_coef * value_loss - self.entropy_coef * entropy_batch.mean() - - # Symmetry loss - if self.symmetry: - # obtain the symmetric actions - # if we did augmentation before then we don't need to augment again - if not self.symmetry["use_data_augmentation"]: - data_augmentation_func = self.symmetry["data_augmentation_func"] - obs_batch, _ = data_augmentation_func(obs=obs_batch, actions=None, env=self.symmetry["_env"]) - # compute number of augmentations per sample - num_aug = int(obs_batch.shape[0] / original_batch_size) - - # actions predicted by the actor for symmetrically-augmented observations - mean_actions_batch = self.policy.act_inference(obs_batch.detach().clone()) - - # compute the symmetrically augmented actions - # note: we are assuming the first augmentation is the original one. - # We do not use the action_batch from earlier since that action was sampled from the distribution. - # However, the symmetry loss is computed using the mean of the distribution. - action_mean_orig = mean_actions_batch[:original_batch_size] - _, actions_mean_symm_batch = data_augmentation_func( - obs=None, actions=action_mean_orig, env=self.symmetry["_env"] - ) - - # compute the loss (we skip the first augmentation as it is the original one) - mse_loss = torch.nn.MSELoss() - symmetry_loss = mse_loss( - mean_actions_batch[original_batch_size:], actions_mean_symm_batch.detach()[original_batch_size:] - ) - # add the loss to the total loss - if self.symmetry["use_mirror_loss"]: - loss += self.symmetry["mirror_loss_coeff"] * symmetry_loss - else: - symmetry_loss = symmetry_loss.detach() - - # Random Network Distillation loss - # TODO: Move this processing to inside RND module. - if self.rnd: - # extract the rnd_state - # TODO: Check if we still need torch no grad. It is just an affine transformation. - with torch.no_grad(): - rnd_state_batch = self.rnd.get_rnd_state(obs_batch[:original_batch_size]) - rnd_state_batch = self.rnd.state_normalizer(rnd_state_batch) - # predict the embedding and the target - predicted_embedding = self.rnd.predictor(rnd_state_batch) - target_embedding = self.rnd.target(rnd_state_batch).detach() - # compute the loss as the mean squared error - mseloss = torch.nn.MSELoss() - rnd_loss = mseloss(predicted_embedding, target_embedding) - - # Compute the gradients - # -- For PPO - self.optimizer.zero_grad() - loss.backward() - # -- For RND - if self.rnd: - self.rnd_optimizer.zero_grad() # type: ignore - rnd_loss.backward() - - # Collect gradients from all GPUs - if self.is_multi_gpu: - self.reduce_parameters() - - # Apply the gradients - # -- For PPO - nn.utils.clip_grad_norm_(self.policy.parameters(), self.max_grad_norm) - self.optimizer.step() - # -- For RND - if self.rnd_optimizer: - self.rnd_optimizer.step() - - # Store the losses - mean_value_loss += value_loss.item() - mean_surrogate_loss += surrogate_loss.item() - mean_entropy += entropy_batch.mean().item() - # -- RND loss - if mean_rnd_loss is not None: - mean_rnd_loss += rnd_loss.item() - # -- Symmetry loss - if mean_symmetry_loss is not None: - mean_symmetry_loss += symmetry_loss.item() - - # -- For PPO - num_updates = self.num_learning_epochs * self.num_mini_batches - mean_value_loss /= num_updates - mean_surrogate_loss /= num_updates - mean_entropy /= num_updates - # -- For RND - if mean_rnd_loss is not None: - mean_rnd_loss /= num_updates - # -- For Symmetry - if mean_symmetry_loss is not None: - mean_symmetry_loss /= num_updates - # -- Clear the storage - self.storage.clear() - - # construct the loss dictionary - loss_dict = { - "value_function": mean_value_loss, - "surrogate": mean_surrogate_loss, - "entropy": mean_entropy, - } - if self.rnd: - loss_dict["rnd"] = mean_rnd_loss - if self.symmetry: - loss_dict["symmetry"] = mean_symmetry_loss - - return loss_dict - - """ - Helper functions - """ - - def broadcast_parameters(self): - """Broadcast model parameters to all GPUs.""" - # obtain the model parameters on current GPU - model_params = [self.policy.state_dict()] - if self.rnd: - model_params.append(self.rnd.predictor.state_dict()) - # broadcast the model parameters - torch.distributed.broadcast_object_list(model_params, src=0) - # load the model parameters on all GPUs from source GPU - self.policy.load_state_dict(model_params[0]) - if self.rnd: - self.rnd.predictor.load_state_dict(model_params[1]) - - def reduce_parameters(self): - """Collect gradients from all GPUs and average them. - - This function is called after the backward pass to synchronize the gradients across all GPUs. - """ - # Create a tensor to store the gradients - grads = [param.grad.view(-1) for param in self.policy.parameters() if param.grad is not None] - if self.rnd: - grads += [param.grad.view(-1) for param in self.rnd.parameters() if param.grad is not None] - all_grads = torch.cat(grads) - - # Average the gradients across all GPUs - torch.distributed.all_reduce(all_grads, op=torch.distributed.ReduceOp.SUM) - all_grads /= self.gpu_world_size - - # Get all parameters - all_params = self.policy.parameters() - if self.rnd: - all_params = chain(all_params, self.rnd.parameters()) - - # Update the gradients for all parameters with the reduced gradients - offset = 0 - for param in all_params: - if param.grad is not None: - numel = param.numel() - # copy data back from shared buffer - param.grad.data.copy_(all_grads[offset : offset + numel].view_as(param.grad.data)) - # update the offset for the next parameter - offset += numel diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/env/__init__.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/env/__init__.py deleted file mode 100644 index ab7c056..0000000 --- a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/env/__init__.py +++ /dev/null @@ -1,10 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -"""Submodule defining the environment definitions.""" - -from .vec_env import VecEnv - -__all__ = ["VecEnv"] diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/env/vec_env.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/env/vec_env.py deleted file mode 100644 index ae471cf..0000000 --- a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/env/vec_env.py +++ /dev/null @@ -1,113 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -from __future__ import annotations - -import torch -from abc import ABC, abstractmethod -from tensordict import TensorDict - - -class VecEnv(ABC): - """Abstract class for a vectorized environment. - - The vectorized environment is a collection of environments that are synchronized. This means that - the same type of action is applied to all environments and the same type of observation is returned from all - environments. - """ - - num_envs: int - """Number of environments.""" - - num_actions: int - """Number of actions.""" - - max_episode_length: int | torch.Tensor - - max_episode_length_s: float - """Maximum episode length. - - The maximum episode length can be a scalar or a tensor. If it is a scalar, it is the same for all environments. - If it is a tensor, it is the maximum episode length for each environment. This is useful for dynamic episode - lengths. - """ - - episode_length_buf: torch.Tensor - """Buffer for current episode lengths.""" - - device: torch.device | str - """Device to use.""" - - cfg: dict | object - """Configuration object.""" - - reset_env_ids: torch.Tensor | None = None - - contact_phase: torch.Tensor | None = None - """ - Operations. - """ - - @abstractmethod - def get_observations(self) -> TensorDict: - """Return the current observations. - - Returns: - observations (TensorDict): Observations from the environment. - """ - raise NotImplementedError - - @abstractmethod - def get_amp_observations(self) -> TensorDict: - """Return the current AMP observations. - - Returns: - observations (TensorDict): Observations from the environment. - """ - raise NotImplementedError - - @abstractmethod - def step(self, actions: torch.Tensor) -> tuple[TensorDict, torch.Tensor, torch.Tensor, dict]: - """Apply input action to the environment. - - Args: - actions (torch.Tensor): Input actions to apply. Shape: (num_envs, num_actions) - - Returns: - observations (TensorDict): Observations from the environment. - rewards (torch.Tensor): Rewards from the environment. Shape: (num_envs,) - dones (torch.Tensor): Done flags from the environment. Shape: (num_envs,) - extras (dict): Extra information from the environment. - - Observations: - - The observations TensorDict usually contains multiple observation groups. The `obs_groups` - dictionary of the runner configuration specifies which observation groups are used for which - purpose, i.e., it maps the available observation groups to observation sets. The observation sets - (keys of the `obs_groups` dictionary) currently used by rsl_rl are: - - - "policy": Specified observation groups are used as input to the actor/student network. - - "critic": Specified observation groups are used as input to the critic network. - - "teacher": Specified observation groups are used as input to the teacher network. - - "rnd_state": Specified observation groups are used as input to the RND network. - - Incomplete or incorrect configurations are handled in the `resolve_obs_groups()` function in - `rsl_rl/utils/utils.py`. - - Extras: - - The extras dictionary includes metrics such as the episode reward, episode length, etc. The following - dictionary keys are used by rsl_rl: - - - "time_outs" (torch.Tensor): Timeouts for the environments. These correspond to terminations that - happen due to time limits and not due to the environment reaching a terminal state. This is useful - for environments that have a fixed episode length. - - - "log" (dict[str, float | torch.Tensor]): Additional information for logging and debugging purposes. - The key should be a string and start with "/" for namespacing. The value can be a scalar or a - tensor. If it is a tensor, the mean of the tensor is used for logging. - """ - raise NotImplementedError - \ No newline at end of file diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/__init__.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/__init__.py deleted file mode 100644 index 9afe1a2..0000000 --- a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/__init__.py +++ /dev/null @@ -1,21 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -"""Definitions for neural-network components for RL-agents.""" - -from .actor_critic import ActorCritic -from .actor_critic_recurrent import ActorCriticRecurrent -from .rnd import * -from .student_teacher import StudentTeacher -from .student_teacher_recurrent import StudentTeacherRecurrent -from .symmetry import * -from .discriminator_multi import DiscriminatorMulti -__all__ = [ - "ActorCritic", - "ActorCriticRecurrent", - "StudentTeacher", - "StudentTeacherRecurrent", - "DiscriminatorMulti", -] diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/actor_critic.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/actor_critic.py deleted file mode 100644 index 0efc36b..0000000 --- a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/actor_critic.py +++ /dev/null @@ -1,195 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -from __future__ import annotations - -import torch -import torch.nn as nn -from torch.distributions import Normal - -from rsl_rl.networks import MLP, EmpiricalNormalization - - -class ActorCritic(nn.Module): - is_recurrent = False - - def __init__( - self, - obs, - obs_groups, - num_actions, - actor_obs_normalization=False, - critic_obs_normalization=False, - actor_hidden_dims=[256, 256, 256], - critic_hidden_dims=[256, 256, 256], - activation="elu", - init_noise_std=1.0, - noise_std_type: str = "scalar", - state_dependent_std=False, - **kwargs, - ): - if kwargs: - print( - "ActorCritic.__init__ got unexpected arguments, which will be ignored: " - + str([key for key in kwargs.keys()]) - ) - super().__init__() - - # get the observation dimensions - self.obs_groups = obs_groups - num_actor_obs = 0 - for obs_group in obs_groups["policy"]: - assert len(obs[obs_group].shape) == 2, "The ActorCritic module only supports 1D observations." - num_actor_obs += obs[obs_group].shape[-1] - num_critic_obs = 0 - for obs_group in obs_groups["critic"]: - assert len(obs[obs_group].shape) == 2, "The ActorCritic module only supports 1D observations." - num_critic_obs += obs[obs_group].shape[-1] - - self.state_dependent_std = state_dependent_std - # actor - if self.state_dependent_std: - self.actor = MLP(num_actor_obs, [2, num_actions], actor_hidden_dims, activation) - else: - self.actor = MLP(num_actor_obs, num_actions, actor_hidden_dims, activation) - # actor observation normalization - self.actor_obs_normalization = actor_obs_normalization - if actor_obs_normalization: - self.actor_obs_normalizer = EmpiricalNormalization(num_actor_obs) - else: - self.actor_obs_normalizer = torch.nn.Identity() - print(f"Actor MLP: {self.actor}") - - # critic - self.critic = MLP(num_critic_obs, 1, critic_hidden_dims, activation) - # critic observation normalization - self.critic_obs_normalization = critic_obs_normalization - if critic_obs_normalization: - self.critic_obs_normalizer = EmpiricalNormalization(num_critic_obs) - else: - self.critic_obs_normalizer = torch.nn.Identity() - print(f"Critic MLP: {self.critic}") - - # Action noise - self.noise_std_type = noise_std_type - if self.state_dependent_std: - torch.nn.init.zeros_(self.actor[-2].weight[num_actions:]) - if self.noise_std_type == "scalar": - torch.nn.init.constant_(self.actor[-2].bias[num_actions:], init_noise_std) - elif self.noise_std_type == "log": - torch.nn.init.constant_( - self.actor[-2].bias[num_actions:], torch.log(torch.tensor(init_noise_std + 1e-7)) - ) - else: - raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") - else: - if self.noise_std_type == "scalar": - self.std = nn.Parameter(init_noise_std * torch.ones(num_actions)) - elif self.noise_std_type == "log": - self.log_std = nn.Parameter(torch.log(init_noise_std * torch.ones(num_actions))) - else: - raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") - - # Action distribution (populated in update_distribution) - self.distribution = None - # disable args validation for speedup - Normal.set_default_validate_args(False) - - def reset(self, dones=None): - pass - - def forward(self): - raise NotImplementedError - - @property - def action_mean(self): - return self.distribution.mean - - @property - def action_std(self): - return self.distribution.stddev - - @property - def entropy(self): - return self.distribution.entropy().sum(dim=-1) - - def update_distribution(self, obs): - if self.state_dependent_std: - # compute mean and standard deviation - mean_and_std = self.actor(obs) - if self.noise_std_type == "scalar": - mean, std = torch.unbind(mean_and_std, dim=-2) - elif self.noise_std_type == "log": - mean, log_std = torch.unbind(mean_and_std, dim=-2) - std = torch.exp(log_std) - else: - raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") - else: - # compute mean - mean = self.actor(obs) - # compute standard deviation - if self.noise_std_type == "scalar": - std = self.std.expand_as(mean) - elif self.noise_std_type == "log": - std = torch.exp(self.log_std).expand_as(mean) - else: - raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") - # create distribution - self.distribution = Normal(mean, std) - - def act(self, obs, **kwargs): - obs = self.get_actor_obs(obs) - obs = self.actor_obs_normalizer(obs) - self.update_distribution(obs) - return self.distribution.sample() - - def act_inference(self, obs): - obs = self.get_actor_obs(obs) - obs = self.actor_obs_normalizer(obs) - return self.actor(obs) - - def evaluate(self, obs, **kwargs): - obs = self.get_critic_obs(obs) - obs = self.critic_obs_normalizer(obs) - return self.critic(obs) - - def get_actor_obs(self, obs): - obs_list = [] - for obs_group in self.obs_groups["policy"]: - obs_list.append(obs[obs_group]) - return torch.cat(obs_list, dim=-1) - - def get_critic_obs(self, obs): - obs_list = [] - for obs_group in self.obs_groups["critic"]: - obs_list.append(obs[obs_group]) - return torch.cat(obs_list, dim=-1) - - def get_actions_log_prob(self, actions): - return self.distribution.log_prob(actions).sum(dim=-1) - - def update_normalization(self, obs): - if self.actor_obs_normalization: - actor_obs = self.get_actor_obs(obs) - self.actor_obs_normalizer.update(actor_obs) - if self.critic_obs_normalization: - critic_obs = self.get_critic_obs(obs) - self.critic_obs_normalizer.update(critic_obs) - - def load_state_dict(self, state_dict, strict=True): - """Load the parameters of the actor-critic model. - - Args: - state_dict (dict): State dictionary of the model. - strict (bool): Whether to strictly enforce that the keys in state_dict match the keys returned by this - module's state_dict() function. - - Returns: - bool: Whether this training resumes a previous training. This flag is used by the `load()` function of - `OnPolicyRunner` to determine how to load further parameters (relevant for, e.g., distillation). - """ - - super().load_state_dict(state_dict, strict=strict) - return True # training resumes diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/actor_critic_recurrent.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/actor_critic_recurrent.py deleted file mode 100644 index bba46ca..0000000 --- a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/actor_critic_recurrent.py +++ /dev/null @@ -1,218 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -from __future__ import annotations - -import torch -import torch.nn as nn -import warnings -from torch.distributions import Normal - -from rsl_rl.networks import MLP, EmpiricalNormalization, Memory - - -class ActorCriticRecurrent(nn.Module): - is_recurrent = True - - def __init__( - self, - obs, - obs_groups, - num_actions, - actor_obs_normalization=False, - critic_obs_normalization=False, - actor_hidden_dims=[256, 256, 256], - critic_hidden_dims=[256, 256, 256], - activation="elu", - init_noise_std=1.0, - noise_std_type: str = "scalar", - state_dependent_std=False, - rnn_type="lstm", - rnn_hidden_dim=256, - rnn_num_layers=1, - **kwargs, - ): - if "rnn_hidden_size" in kwargs: - warnings.warn( - "The argument `rnn_hidden_size` is deprecated and will be removed in a future version. " - "Please use `rnn_hidden_dim` instead.", - DeprecationWarning, - ) - if rnn_hidden_dim == 256: # Only override if the new argument is at its default - rnn_hidden_dim = kwargs.pop("rnn_hidden_size") - if kwargs: - print( - "ActorCriticRecurrent.__init__ got unexpected arguments, which will be ignored: " + str(kwargs.keys()), - ) - super().__init__() - - # get the observation dimensions - self.obs_groups = obs_groups - num_actor_obs = 0 - for obs_group in obs_groups["policy"]: - assert len(obs[obs_group].shape) == 2, "The ActorCriticRecurrent module only supports 1D observations." - num_actor_obs += obs[obs_group].shape[-1] - num_critic_obs = 0 - for obs_group in obs_groups["critic"]: - assert len(obs[obs_group].shape) == 2, "The ActorCriticRecurrent module only supports 1D observations." - num_critic_obs += obs[obs_group].shape[-1] - - self.state_dependent_std = state_dependent_std - # actor - self.memory_a = Memory(num_actor_obs, type=rnn_type, num_layers=rnn_num_layers, hidden_size=rnn_hidden_dim) - if self.state_dependent_std: - self.actor = MLP(rnn_hidden_dim, [2, num_actions], actor_hidden_dims, activation) - else: - self.actor = MLP(rnn_hidden_dim, num_actions, actor_hidden_dims, activation) - - # actor observation normalization - self.actor_obs_normalization = actor_obs_normalization - if actor_obs_normalization: - self.actor_obs_normalizer = EmpiricalNormalization(num_actor_obs) - else: - self.actor_obs_normalizer = torch.nn.Identity() - print(f"Actor RNN: {self.memory_a}") - print(f"Actor MLP: {self.actor}") - - # critic - self.memory_c = Memory(num_critic_obs, type=rnn_type, num_layers=rnn_num_layers, hidden_size=rnn_hidden_dim) - self.critic = MLP(rnn_hidden_dim, 1, critic_hidden_dims, activation) - # critic observation normalization - self.critic_obs_normalization = critic_obs_normalization - if critic_obs_normalization: - self.critic_obs_normalizer = EmpiricalNormalization(num_critic_obs) - else: - self.critic_obs_normalizer = torch.nn.Identity() - print(f"Critic RNN: {self.memory_c}") - print(f"Critic MLP: {self.critic}") - - # Action noise - self.noise_std_type = noise_std_type - if self.state_dependent_std: - torch.nn.init.zeros_(self.actor[-2].weight[num_actions:]) - if self.noise_std_type == "scalar": - torch.nn.init.constant_(self.actor[-2].bias[num_actions:], init_noise_std) - elif self.noise_std_type == "log": - torch.nn.init.constant_( - self.actor[-2].bias[num_actions:], torch.log(torch.tensor(init_noise_std + 1e-7)) - ) - else: - raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") - else: - if self.noise_std_type == "scalar": - self.std = nn.Parameter(init_noise_std * torch.ones(num_actions)) - elif self.noise_std_type == "log": - self.log_std = nn.Parameter(torch.log(init_noise_std * torch.ones(num_actions))) - else: - raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") - - # Action distribution (populated in update_distribution) - self.distribution = None - # disable args validation for speedup - Normal.set_default_validate_args(False) - - @property - def action_mean(self): - return self.distribution.mean - - @property - def action_std(self): - return self.distribution.stddev - - @property - def entropy(self): - return self.distribution.entropy().sum(dim=-1) - - def reset(self, dones=None): - self.memory_a.reset(dones) - self.memory_c.reset(dones) - - def forward(self): - raise NotImplementedError - - def update_distribution(self, obs): - if self.state_dependent_std: - # compute mean and standard deviation - mean_and_std = self.actor(obs) - if self.noise_std_type == "scalar": - mean, std = torch.unbind(mean_and_std, dim=-2) - elif self.noise_std_type == "log": - mean, log_std = torch.unbind(mean_and_std, dim=-2) - std = torch.exp(log_std) - else: - raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") - else: - # compute mean - mean = self.actor(obs) - # compute standard deviation - if self.noise_std_type == "scalar": - std = self.std.expand_as(mean) - elif self.noise_std_type == "log": - std = torch.exp(self.log_std).expand_as(mean) - else: - raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") - # create distribution - self.distribution = Normal(mean, std) - - def act(self, obs, masks=None, hidden_states=None): - obs = self.get_actor_obs(obs) - obs = self.actor_obs_normalizer(obs) - out_mem = self.memory_a(obs, masks, hidden_states).squeeze(0) - self.update_distribution(out_mem) - return self.distribution.sample() - - def act_inference(self, obs): - obs = self.get_actor_obs(obs) - obs = self.actor_obs_normalizer(obs) - out_mem = self.memory_a(obs).squeeze(0) - return self.actor(out_mem) - - def evaluate(self, obs, masks=None, hidden_states=None): - obs = self.get_critic_obs(obs) - obs = self.critic_obs_normalizer(obs) - out_mem = self.memory_c(obs, masks, hidden_states).squeeze(0) - return self.critic(out_mem) - - def get_actor_obs(self, obs): - obs_list = [] - for obs_group in self.obs_groups["policy"]: - obs_list.append(obs[obs_group]) - return torch.cat(obs_list, dim=-1) - - def get_critic_obs(self, obs): - obs_list = [] - for obs_group in self.obs_groups["critic"]: - obs_list.append(obs[obs_group]) - return torch.cat(obs_list, dim=-1) - - def get_actions_log_prob(self, actions): - return self.distribution.log_prob(actions).sum(dim=-1) - - def get_hidden_states(self): - return self.memory_a.hidden_states, self.memory_c.hidden_states - - def update_normalization(self, obs): - if self.actor_obs_normalization: - actor_obs = self.get_actor_obs(obs) - self.actor_obs_normalizer.update(actor_obs) - if self.critic_obs_normalization: - critic_obs = self.get_critic_obs(obs) - self.critic_obs_normalizer.update(critic_obs) - - def load_state_dict(self, state_dict, strict=True): - """Load the parameters of the actor-critic model. - - Args: - state_dict (dict): State dictionary of the model. - strict (bool): Whether to strictly enforce that the keys in state_dict match the keys returned by this - module's state_dict() function. - - Returns: - bool: Whether this training resumes a previous training. This flag is used by the `load()` function of - `OnPolicyRunner` to determine how to load further parameters (relevant for, e.g., distillation). - """ - - super().load_state_dict(state_dict, strict=strict) - return True diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/discriminator_multi.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/discriminator_multi.py deleted file mode 100644 index 7899d23..0000000 --- a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/discriminator_multi.py +++ /dev/null @@ -1,102 +0,0 @@ -import torch -import torch.nn as nn -from torch import autograd -import torch.nn.utils.spectral_norm as spectral_norm - - -class DiscriminatorMulti(nn.Module): - def __init__( - self, state_dim, amp_reward_coef, hidden_layer_sizes, device, - num_frames=2, task_reward_lerp=0.0, use_lerp=True): - super(DiscriminatorMulti, self).__init__() - - self.device = device - self.state_dim = state_dim - self.use_lerp = use_lerp - self.num_frames = num_frames # 存储帧数参数 - - self.amp_reward_coef = amp_reward_coef - amp_layers = [] - - curr_in_dim = state_dim * num_frames - for hidden_dim in hidden_layer_sizes: - amp_layers.append(spectral_norm(nn.Linear(curr_in_dim, hidden_dim))) - amp_layers.append(nn.ReLU()) - curr_in_dim = hidden_dim - self.trunk = nn.Sequential(*amp_layers).to(device) - self.amp_linear = spectral_norm(nn.Linear(hidden_layer_sizes[-1], 1)).to(device) - - self.trunk.train() - self.amp_linear.train() - - self.task_reward_lerp = task_reward_lerp - - def forward(self, x): - h = self.trunk(x) - d = self.amp_linear(h) - return d - - def compute_grad_pen(self, - expert_states, # 改为接收多帧状态列表 - lambda_=10): - # 将多帧状态沿最后一个维度拼接 - expert_data = expert_states.flatten(1) - expert_data.requires_grad = True - - disc = self.amp_linear(self.trunk(expert_data)) - ones = torch.ones(disc.size(), device=disc.device) - grad = autograd.grad( - outputs=disc, inputs=expert_data, - grad_outputs=ones, create_graph=True, - retain_graph=True, only_inputs=True)[0] - - # Enforce that the grad norm approaches 0. - grad_pen = lambda_ * (grad.norm(2, dim=1) - 0).pow(2).mean() - return grad_pen - - - def get_disc_weights(self): - weights = [] - for m in self.trunk.modules(): - if isinstance(m, nn.Linear): - weights.append(torch.flatten(m.weight)) - - weights.append(torch.flatten(self.amp_linear.weight)) - return weights - - def get_disc_logit_weights(self): - return torch.flatten(self.amp_linear.weight) - - def predict_amp_reward( - self, states, # 改为接收多帧状态列表 - task_reward, normalizer=None): - """ - states: torch.Tensor, shape=(num_envs, num_frames, state_dim) - task_reward: torch.Tensor, shape=(num_envs, 1) - """ - # import ipdb; ipdb.set_trace() - with torch.no_grad(): - self.eval() - if normalizer is not None: - # 对每一帧状态进行归一化 - states = normalizer.normalize_torch(states, self.device) - - # 拼接多帧状态 - state_cat = states.flatten(1) - d = self.amp_linear(self.trunk(state_cat)) - disc_reward = self.amp_reward_coef * torch.clamp(1 - (1/4) * torch.square(d - 1), min=0) - - if self.use_lerp: - if self.task_reward_lerp > 0: - reward = self._lerp_reward(disc_reward, task_reward.unsqueeze(-1)) - self.train() - return reward.squeeze(), d, disc_reward.squeeze() * (1.0 - self.task_reward_lerp) - else: - disc_reward *= 0.02 - reward = task_reward.unsqueeze(-1) + disc_reward - self.train() - return reward.squeeze(), d, disc_reward.squeeze() - - def _lerp_reward(self, disc_r, task_r): - r = (1.0 - self.task_reward_lerp) * disc_r + self.task_reward_lerp * task_r - return r \ No newline at end of file diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/rnd.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/rnd.py deleted file mode 100644 index 8e65c43..0000000 --- a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/rnd.py +++ /dev/null @@ -1,209 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -from __future__ import annotations - -import torch -import torch.nn as nn - -from rsl_rl.networks import MLP, EmpiricalDiscountedVariationNormalization, EmpiricalNormalization - - -class RandomNetworkDistillation(nn.Module): - """Implementation of Random Network Distillation (RND) [1] - - References: - .. [1] Burda, Yuri, et al. "Exploration by random network distillation." arXiv preprint arXiv:1810.12894 (2018). - """ - - def __init__( - self, - num_states: int, - obs_groups: dict, - num_outputs: int, - predictor_hidden_dims: list[int], - target_hidden_dims: list[int], - activation: str = "elu", - weight: float = 0.0, - state_normalization: bool = False, - reward_normalization: bool = False, - device: str = "cpu", - weight_schedule: dict | None = None, - ): - """Initialize the RND module. - - - If :attr:`state_normalization` is True, then the input state is normalized using an Empirical Normalization layer. - - If :attr:`reward_normalization` is True, then the intrinsic reward is normalized using an Empirical Discounted - Variation Normalization layer. - - .. note:: - If the hidden dimensions are -1 in the predictor and target networks configuration, then the number of states - is used as the hidden dimension. - - Args: - num_states: Number of states/inputs to the predictor and target networks. - num_outputs: Number of outputs (embedding size) of the predictor and target networks. - predictor_hidden_dims: List of hidden dimensions of the predictor network. - target_hidden_dims: List of hidden dimensions of the target network. - activation: Activation function. Defaults to "elu". - weight: Scaling factor of the intrinsic reward. Defaults to 0.0. - state_normalization: Whether to normalize the input state. Defaults to False. - reward_normalization: Whether to normalize the intrinsic reward. Defaults to False. - device: Device to use. Defaults to "cpu". - weight_schedule: The type of schedule to use for the RND weight parameter. - Defaults to None, in which case the weight parameter is constant. - It is a dictionary with the following keys: - - - "mode": The type of schedule to use for the RND weight parameter. - - "constant": Constant weight schedule. - - "step": Step weight schedule. - - "linear": Linear weight schedule. - - For the "step" weight schedule, the following parameters are required: - - - "final_step": The step at which the weight parameter is set to the final value. - - "final_value": The final value of the weight parameter. - - For the "linear" weight schedule, the following parameters are required: - - "initial_step": The step at which the weight parameter is set to the initial value. - - "final_step": The step at which the weight parameter is set to the final value. - - "final_value": The final value of the weight parameter. - """ - # initialize parent class - super().__init__() - - # Store parameters - self.num_states = num_states - self.obs_groups = obs_groups - self.num_outputs = num_outputs - self.initial_weight = weight - self.device = device - self.state_normalization = state_normalization - self.reward_normalization = reward_normalization - - # Normalization of input gates - if state_normalization: - self.state_normalizer = EmpiricalNormalization(shape=[self.num_states], until=1.0e8).to(self.device) - else: - self.state_normalizer = torch.nn.Identity() - # Normalization of intrinsic reward - if reward_normalization: - self.reward_normalizer = EmpiricalDiscountedVariationNormalization(shape=[], until=1.0e8).to(self.device) - else: - self.reward_normalizer = torch.nn.Identity() - - # counter for the number of updates - self.update_counter = 0 - - # resolve weight schedule - if weight_schedule is not None: - self.weight_scheduler_params = weight_schedule - self.weight_scheduler = getattr(self, f"_{weight_schedule['mode']}_weight_schedule") - else: - self.weight_scheduler = None - # Create network architecture - self.predictor = MLP(num_states, num_outputs, predictor_hidden_dims, activation).to(self.device) - self.target = MLP(num_states, num_outputs, target_hidden_dims, activation).to(self.device) - - # make target network not trainable - self.target.eval() - - def get_intrinsic_reward(self, obs) -> torch.Tensor: - # Note: the counter is updated number of env steps per learning iteration - self.update_counter += 1 - # Extract the rnd state from the observation - rnd_state = self.get_rnd_state(obs) - rnd_state = self.state_normalizer(rnd_state) - # Obtain the embedding of the rnd state from the target and predictor networks - target_embedding = self.target(rnd_state).detach() - predictor_embedding = self.predictor(rnd_state).detach() - # Compute the intrinsic reward as the distance between the embeddings - intrinsic_reward = torch.linalg.norm(target_embedding - predictor_embedding, dim=1) - # Normalize intrinsic reward - intrinsic_reward = self.reward_normalizer(intrinsic_reward) - - # Check the weight schedule - if self.weight_scheduler is not None: - self.weight = self.weight_scheduler(step=self.update_counter, **self.weight_scheduler_params) - else: - self.weight = self.initial_weight - # Scale intrinsic reward - intrinsic_reward *= self.weight - - return intrinsic_reward - - def forward(self, *args, **kwargs): - raise RuntimeError("Forward method is not implemented. Use get_intrinsic_reward instead.") - - def train(self, mode: bool = True): - # sets module into training mode - self.predictor.train(mode) - if self.state_normalization: - self.state_normalizer.train(mode) - if self.reward_normalization: - self.reward_normalizer.train(mode) - return self - - def eval(self): - return self.train(False) - - def get_rnd_state(self, obs): - obs_list = [] - for obs_group in self.obs_groups["rnd_state"]: - obs_list.append(obs[obs_group]) - return torch.cat(obs_list, dim=-1) - - def update_normalization(self, obs): - # Normalize the state - if self.state_normalization: - rnd_state = self.get_rnd_state(obs) - self.state_normalizer.update(rnd_state) - - """ - Different weight schedules. - """ - - def _constant_weight_schedule(self, step: int, **kwargs): - return self.initial_weight - - def _step_weight_schedule(self, step: int, final_step: int, final_value: float, **kwargs): - return self.initial_weight if step < final_step else final_value - - def _linear_weight_schedule(self, step: int, initial_step: int, final_step: int, final_value: float, **kwargs): - if step < initial_step: - return self.initial_weight - elif step > final_step: - return final_value - else: - return self.initial_weight + (final_value - self.initial_weight) * (step - initial_step) / ( - final_step - initial_step - ) - - -def resolve_rnd_config(alg_cfg, obs, obs_groups, env): - """Resolve the RND configuration. - - Args: - alg_cfg: The algorithm configuration dictionary. - obs: The observation dictionary. - obs_groups: The observation groups dictionary. - env: The environment. - - Returns: - The resolved algorithm configuration dictionary. - """ - # resolve dimension of rnd gated state - if "rnd_cfg" in alg_cfg and alg_cfg["rnd_cfg"] is not None: - # get dimension of rnd gated state - num_rnd_state = 0 - for obs_group in obs_groups["rnd_state"]: - assert len(obs[obs_group].shape) == 2, "The RND module only supports 1D observations." - num_rnd_state += obs[obs_group].shape[-1] - # add rnd gated state to config - alg_cfg["rnd_cfg"]["num_states"] = num_rnd_state - alg_cfg["rnd_cfg"]["obs_groups"] = obs_groups - # scale down the rnd weight with timestep - alg_cfg["rnd_cfg"]["weight"] *= env.unwrapped.step_dt - return alg_cfg diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/student_teacher.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/student_teacher.py deleted file mode 100644 index 6bf1380..0000000 --- a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/student_teacher.py +++ /dev/null @@ -1,206 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -from __future__ import annotations - -import torch -import torch.nn as nn -from torch.distributions import Normal - -from rsl_rl.networks import MLP, EmpiricalNormalization - - -class StudentTeacher(nn.Module): - is_recurrent = False - - def __init__( - self, - obs, - obs_groups, - num_actions, - student_obs_normalization=False, - teacher_obs_normalization=False, - student_hidden_dims=[256, 256, 256], - teacher_hidden_dims=[256, 256, 256], - activation="elu", - init_noise_std=0.1, - noise_std_type: str = "scalar", - **kwargs, - ): - if kwargs: - print( - "StudentTeacher.__init__ got unexpected arguments, which will be ignored: " - + str([key for key in kwargs.keys()]) - ) - super().__init__() - - self.loaded_teacher = False # indicates if teacher has been loaded - - # get the observation dimensions - self.obs_groups = obs_groups - num_student_obs = 0 - for obs_group in obs_groups["policy"]: - assert len(obs[obs_group].shape) == 2, "The StudentTeacher module only supports 1D observations." - num_student_obs += obs[obs_group].shape[-1] - num_teacher_obs = 0 - for obs_group in obs_groups["teacher"]: - assert len(obs[obs_group].shape) == 2, "The StudentTeacher module only supports 1D observations." - num_teacher_obs += obs[obs_group].shape[-1] - - # student - self.student = MLP(num_student_obs, num_actions, student_hidden_dims, activation) - - # student observation normalization - self.student_obs_normalization = student_obs_normalization - if student_obs_normalization: - self.student_obs_normalizer = EmpiricalNormalization(num_student_obs) - else: - self.student_obs_normalizer = torch.nn.Identity() - - print(f"Student MLP: {self.student}") - - # teacher - self.teacher = MLP(num_teacher_obs, num_actions, teacher_hidden_dims, activation) - self.teacher.eval() - - # teacher observation normalization - self.teacher_obs_normalization = teacher_obs_normalization - if teacher_obs_normalization: - self.teacher_obs_normalizer = EmpiricalNormalization(num_teacher_obs) - else: - self.teacher_obs_normalizer = torch.nn.Identity() - - print(f"Teacher MLP: {self.teacher}") - - # action noise - self.noise_std_type = noise_std_type - if self.noise_std_type == "scalar": - self.std = nn.Parameter(init_noise_std * torch.ones(num_actions)) - elif self.noise_std_type == "log": - self.log_std = nn.Parameter(torch.log(init_noise_std * torch.ones(num_actions))) - else: - raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") - - # action distribution (populated in update_distribution) - self.distribution = None - # disable args validation for speedup - Normal.set_default_validate_args(False) - - def reset(self, dones=None, hidden_states=None): - pass - - def forward(self): - raise NotImplementedError - - @property - def action_mean(self): - return self.distribution.mean - - @property - def action_std(self): - return self.distribution.stddev - - @property - def entropy(self): - return self.distribution.entropy().sum(dim=-1) - - def update_distribution(self, obs): - # compute mean - mean = self.student(obs) - # compute standard deviation - if self.noise_std_type == "scalar": - std = self.std.expand_as(mean) - elif self.noise_std_type == "log": - std = torch.exp(self.log_std).expand_as(mean) - else: - raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") - # create distribution - self.distribution = Normal(mean, std) - - def act(self, obs): - obs = self.get_student_obs(obs) - obs = self.student_obs_normalizer(obs) - self.update_distribution(obs) - return self.distribution.sample() - - def act_inference(self, obs): - obs = self.get_student_obs(obs) - obs = self.student_obs_normalizer(obs) - return self.student(obs) - - def evaluate(self, obs): - obs = self.get_teacher_obs(obs) - obs = self.teacher_obs_normalizer(obs) - with torch.no_grad(): - return self.teacher(obs) - - def get_student_obs(self, obs): - obs_list = [] - for obs_group in self.obs_groups["policy"]: - obs_list.append(obs[obs_group]) - return torch.cat(obs_list, dim=-1) - - def get_teacher_obs(self, obs): - obs_list = [] - for obs_group in self.obs_groups["teacher"]: - obs_list.append(obs[obs_group]) - return torch.cat(obs_list, dim=-1) - - def get_hidden_states(self): - return None - - def detach_hidden_states(self, dones=None): - pass - - def train(self, mode=True): - super().train(mode) - # make sure teacher is in eval mode - self.teacher.eval() - self.teacher_obs_normalizer.eval() - - def update_normalization(self, obs): - if self.student_obs_normalization: - student_obs = self.get_student_obs(obs) - self.student_obs_normalizer.update(student_obs) - - def load_state_dict(self, state_dict, strict=True): - """Load the parameters of the student and teacher networks. - - Args: - state_dict (dict): State dictionary of the model. - strict (bool): Whether to strictly enforce that the keys in state_dict match the keys returned by this - module's state_dict() function. - - Returns: - bool: Whether this training resumes a previous training. This flag is used by the `load()` function of - `OnPolicyRunner` to determine how to load further parameters. - """ - - # check if state_dict contains teacher and student or just teacher parameters - if any("actor" in key for key in state_dict.keys()): # loading parameters from rl training - # rename keys to match teacher and remove critic parameters - teacher_state_dict = {} - teacher_obs_normalizer_state_dict = {} - for key, value in state_dict.items(): - if "actor." in key: - teacher_state_dict[key.replace("actor.", "")] = value - if "actor_obs_normalizer." in key: - teacher_obs_normalizer_state_dict[key.replace("actor_obs_normalizer.", "")] = value - self.teacher.load_state_dict(teacher_state_dict, strict=strict) - self.teacher_obs_normalizer.load_state_dict(teacher_obs_normalizer_state_dict, strict=strict) - # set flag for successfully loading the parameters - self.loaded_teacher = True - self.teacher.eval() - self.teacher_obs_normalizer.eval() - return False # training does not resume - elif any("student" in key for key in state_dict.keys()): # loading parameters from distillation training - super().load_state_dict(state_dict, strict=strict) - # set flag for successfully loading the parameters - self.loaded_teacher = True - self.teacher.eval() - self.teacher_obs_normalizer.eval() - return True # training resumes - else: - raise ValueError("state_dict does not contain student or teacher parameters") diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/student_teacher_recurrent.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/student_teacher_recurrent.py deleted file mode 100644 index 964a2dc..0000000 --- a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/student_teacher_recurrent.py +++ /dev/null @@ -1,249 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -from __future__ import annotations - -import torch -import torch.nn as nn -import warnings -from torch.distributions import Normal - -from rsl_rl.networks import MLP, EmpiricalNormalization, Memory - - -class StudentTeacherRecurrent(nn.Module): - is_recurrent = True - - def __init__( - self, - obs, - obs_groups, - num_actions, - student_obs_normalization=False, - teacher_obs_normalization=False, - student_hidden_dims=[256, 256, 256], - teacher_hidden_dims=[256, 256, 256], - activation="elu", - init_noise_std=0.1, - noise_std_type: str = "scalar", - rnn_type="lstm", - rnn_hidden_dim=256, - rnn_num_layers=1, - teacher_recurrent=False, - **kwargs, - ): - if "rnn_hidden_size" in kwargs: - warnings.warn( - "The argument `rnn_hidden_size` is deprecated and will be removed in a future version. " - "Please use `rnn_hidden_dim` instead.", - DeprecationWarning, - ) - if rnn_hidden_dim == 256: # Only override if the new argument is at its default - rnn_hidden_dim = kwargs.pop("rnn_hidden_size") - if kwargs: - print( - "StudentTeacherRecurrent.__init__ got unexpected arguments, which will be ignored: " - + str(kwargs.keys()), - ) - super().__init__() - - self.loaded_teacher = False # indicates if teacher has been loaded - self.teacher_recurrent = teacher_recurrent # indicates if teacher is recurrent too - - # get the observation dimensions - self.obs_groups = obs_groups - num_student_obs = 0 - for obs_group in obs_groups["policy"]: - assert len(obs[obs_group].shape) == 2, "The StudentTeacher module only supports 1D observations." - num_student_obs += obs[obs_group].shape[-1] - num_teacher_obs = 0 - for obs_group in obs_groups["teacher"]: - assert len(obs[obs_group].shape) == 2, "The StudentTeacher module only supports 1D observations." - num_teacher_obs += obs[obs_group].shape[-1] - - # student - self.memory_s = Memory(num_student_obs, type=rnn_type, num_layers=rnn_num_layers, hidden_size=rnn_hidden_dim) - self.student = MLP(rnn_hidden_dim, num_actions, student_hidden_dims, activation) - - # student observation normalization - self.student_obs_normalization = student_obs_normalization - if student_obs_normalization: - self.student_obs_normalizer = EmpiricalNormalization(num_student_obs) - else: - self.student_obs_normalizer = torch.nn.Identity() - - print(f"Student RNN: {self.memory_s}") - print(f"Student MLP: {self.student}") - - # teacher - if self.teacher_recurrent: - self.memory_t = Memory( - num_teacher_obs, type=rnn_type, num_layers=rnn_num_layers, hidden_size=rnn_hidden_dim - ) - num_teacher_obs = rnn_hidden_dim - self.teacher = MLP(num_teacher_obs, num_actions, teacher_hidden_dims, activation) - - # teacher observation normalization - self.teacher_obs_normalization = teacher_obs_normalization - if teacher_obs_normalization: - self.teacher_obs_normalizer = EmpiricalNormalization(num_teacher_obs) - else: - self.teacher_obs_normalizer = torch.nn.Identity() - - if self.teacher_recurrent: - print(f"Teacher RNN: {self.memory_t}") - print(f"Teacher MLP: {self.teacher}") - - # action noise - self.noise_std_type = noise_std_type - if self.noise_std_type == "scalar": - self.std = nn.Parameter(init_noise_std * torch.ones(num_actions)) - elif self.noise_std_type == "log": - self.log_std = nn.Parameter(torch.log(init_noise_std * torch.ones(num_actions))) - else: - raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") - - # action distribution (populated in update_distribution) - self.distribution = None - # disable args validation for speedup - Normal.set_default_validate_args(False) - - def reset(self, dones=None, hidden_states=None): - if hidden_states is None: - hidden_states = (None, None) - self.memory_s.reset(dones, hidden_states[0]) - if self.teacher_recurrent: - self.memory_t.reset(dones, hidden_states[1]) - - def forward(self): - raise NotImplementedError - - @property - def action_mean(self): - return self.distribution.mean - - @property - def action_std(self): - return self.distribution.stddev - - @property - def entropy(self): - return self.distribution.entropy().sum(dim=-1) - - def update_distribution(self, obs): - # compute mean - mean = self.student(obs) - # compute standard deviation - if self.noise_std_type == "scalar": - std = self.std.expand_as(mean) - elif self.noise_std_type == "log": - std = torch.exp(self.log_std).expand_as(mean) - else: - raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") - # create distribution - self.distribution = Normal(mean, std) - - def act(self, obs): - obs = self.get_student_obs(obs) - obs = self.student_obs_normalizer(obs) - out_mem = self.memory_s(obs).squeeze(0) - self.update_distribution(out_mem) - return self.distribution.sample() - - def act_inference(self, obs): - obs = self.get_student_obs(obs) - obs = self.student_obs_normalizer(obs) - out_mem = self.memory_s(obs).squeeze(0) - return self.student(out_mem) - - def evaluate(self, obs): - obs = self.get_teacher_obs(obs) - obs = self.teacher_obs_normalizer(obs) - with torch.no_grad(): - if self.teacher_recurrent: - self.memory_t.eval() - obs = self.memory_t(obs).squeeze(0) - return self.teacher(obs) - - def get_student_obs(self, obs): - obs_list = [] - for obs_group in self.obs_groups["policy"]: - obs_list.append(obs[obs_group]) - return torch.cat(obs_list, dim=-1) - - def get_teacher_obs(self, obs): - obs_list = [] - for obs_group in self.obs_groups["teacher"]: - obs_list.append(obs[obs_group]) - return torch.cat(obs_list, dim=-1) - - def get_hidden_states(self): - if self.teacher_recurrent: - return self.memory_s.hidden_states, self.memory_t.hidden_states - else: - return self.memory_s.hidden_states, None - - def detach_hidden_states(self, dones=None): - self.memory_s.detach_hidden_states(dones) - if self.teacher_recurrent: - self.memory_t.detach_hidden_states(dones) - - def train(self, mode=True): - super().train(mode) - # make sure teacher is in eval mode - self.teacher.eval() - self.teacher_obs_normalizer.eval() - - def update_normalization(self, obs): - if self.student_obs_normalization: - student_obs = self.get_student_obs(obs) - self.student_obs_normalizer.update(student_obs) - - def load_state_dict(self, state_dict, strict=True): - """Load the parameters of the student and teacher networks. - - Args: - state_dict (dict): State dictionary of the model. - strict (bool): Whether to strictly enforce that the keys in state_dict match the keys returned by this - module's state_dict() function. - - Returns: - bool: Whether this training resumes a previous training. This flag is used by the `load()` function of - `OnPolicyRunner` to determine how to load further parameters. - """ - - # check if state_dict contains teacher and student or just teacher parameters - if any("actor" in key for key in state_dict.keys()): # loading parameters from rl training - # rename keys to match teacher and remove critic parameters - teacher_state_dict = {} - teacher_obs_normalizer_state_dict = {} - for key, value in state_dict.items(): - if "actor." in key: - teacher_state_dict[key.replace("actor.", "")] = value - if "actor_obs_normalizer." in key: - teacher_obs_normalizer_state_dict[key.replace("actor_obs_normalizer.", "")] = value - self.teacher.load_state_dict(teacher_state_dict, strict=strict) - self.teacher_obs_normalizer.load_state_dict(teacher_obs_normalizer_state_dict, strict=strict) - # also load recurrent memory if teacher is recurrent - if self.teacher_recurrent: - memory_t_state_dict = {} - for key, value in state_dict.items(): - if "memory_a." in key: - memory_t_state_dict[key.replace("memory_a.", "")] = value - self.memory_t.load_state_dict(memory_t_state_dict, strict=strict) - # set flag for successfully loading the parameters - self.loaded_teacher = True - self.teacher.eval() - self.teacher_obs_normalizer.eval() - return False # training does not resume - elif any("student" in key for key in state_dict.keys()): # loading parameters from distillation training - super().load_state_dict(state_dict, strict=strict) - # set flag for successfully loading the parameters - self.loaded_teacher = True - self.teacher.eval() - self.teacher_obs_normalizer.eval() - return True # training resumes - else: - raise ValueError("state_dict does not contain student or teacher parameters") diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/symmetry.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/symmetry.py deleted file mode 100644 index b017515..0000000 --- a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/symmetry.py +++ /dev/null @@ -1,24 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -from __future__ import annotations - - -def resolve_symmetry_config(alg_cfg, env): - """Resolve the symmetry configuration. - - Args: - alg_cfg: The algorithm configuration dictionary. - env: The environment. - - Returns: - The resolved algorithm configuration dictionary. - """ - - # if using symmetry then pass the environment config object - if "symmetry_cfg" in alg_cfg and alg_cfg["symmetry_cfg"] is not None: - # this is used by the symmetry function for handling different observation terms - alg_cfg["symmetry_cfg"]["_env"] = env - return alg_cfg diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/networks/__init__.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/networks/__init__.py deleted file mode 100644 index c18f487..0000000 --- a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/networks/__init__.py +++ /dev/null @@ -1,10 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -"""Definitions for components of modules.""" - -from .memory import Memory -from .mlp import MLP -from .normalization import EmpiricalDiscountedVariationNormalization, EmpiricalNormalization diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/networks/memory.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/networks/memory.py deleted file mode 100644 index 7577357..0000000 --- a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/networks/memory.py +++ /dev/null @@ -1,70 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -from __future__ import annotations - -import torch.nn as nn - -from rsl_rl.utils import unpad_trajectories - - -class Memory(nn.Module): - """Memory module for recurrent networks. - - This module is used to store the hidden states of the policy. - Currently only supports GRU and LSTM. - """ - - def __init__(self, input_size, type="lstm", num_layers=1, hidden_size=256): - super().__init__() - # RNN - rnn_cls = nn.GRU if type.lower() == "gru" else nn.LSTM - self.rnn = rnn_cls(input_size=input_size, hidden_size=hidden_size, num_layers=num_layers) - self.hidden_states = None - - def forward(self, input, masks=None, hidden_states=None): - batch_mode = masks is not None - if batch_mode: - # batch mode: needs saved hidden states - if hidden_states is None: - raise ValueError("Hidden states not passed to memory module during policy update") - out, _ = self.rnn(input, hidden_states) - out = unpad_trajectories(out, masks) - else: - # inference/distillation mode: uses hidden states of last step - out, self.hidden_states = self.rnn(input.unsqueeze(0), self.hidden_states) - return out - - def reset(self, dones=None, hidden_states=None): - if dones is None: # reset all hidden states - if hidden_states is None: - self.hidden_states = None - else: - self.hidden_states = hidden_states - elif self.hidden_states is not None: # reset hidden states of done environments - if hidden_states is None: - if isinstance(self.hidden_states, tuple): # tuple in case of LSTM - for hidden_state in self.hidden_states: - hidden_state[..., dones == 1, :] = 0.0 - else: - self.hidden_states[..., dones == 1, :] = 0.0 - else: - NotImplementedError( - "Resetting hidden states of done environments with custom hidden states is not implemented" - ) - - def detach_hidden_states(self, dones=None): - if self.hidden_states is not None: - if dones is None: # detach all hidden states - if isinstance(self.hidden_states, tuple): # tuple in case of LSTM - self.hidden_states = tuple(hidden_state.detach() for hidden_state in self.hidden_states) - else: - self.hidden_states = self.hidden_states.detach() - else: # detach hidden states of done environments - if isinstance(self.hidden_states, tuple): # tuple in case of LSTM - for hidden_state in self.hidden_states: - hidden_state[..., dones == 1, :] = hidden_state[..., dones == 1, :].detach() - else: - self.hidden_states[..., dones == 1, :] = self.hidden_states[..., dones == 1, :].detach() diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/networks/mlp.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/networks/mlp.py deleted file mode 100644 index e91574e..0000000 --- a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/networks/mlp.py +++ /dev/null @@ -1,120 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -from __future__ import annotations - -import torch -import torch.nn as nn -from functools import reduce - -from rsl_rl.utils import resolve_nn_activation - - -class MLP(nn.Sequential): - """Multi-layer perceptron. - - The MLP network is a sequence of linear layers and activation functions. The - last layer is a linear layer that outputs the desired dimension unless the - last activation function is specified. - - It provides additional conveniences: - - - If the hidden dimensions have a value of ``-1``, the dimension is inferred - from the input dimension. - - If the output dimension is a tuple, the output is reshaped to the desired - shape. - - """ - - def __init__( - self, - input_dim: int, - output_dim: int | tuple[int] | list[int], - hidden_dims: tuple[int] | list[int], - activation: str = "elu", - last_activation: str | None = None, - ): - """Initialize the MLP. - - Args: - input_dim: Dimension of the input. - output_dim: Dimension of the output. - hidden_dims: Dimensions of the hidden layers. A value of ``-1`` indicates - that the dimension should be inferred from the input dimension. - activation: Activation function. Defaults to "elu". - last_activation: Activation function of the last layer. Defaults to None, - in which case the last layer is linear. - """ - super().__init__() - - # resolve activation functions - activation_mod = resolve_nn_activation(activation) - last_activation_mod = resolve_nn_activation(last_activation) if last_activation is not None else None - # resolve number of hidden dims if they are -1 - hidden_dims_processed = [input_dim if dim == -1 else dim for dim in hidden_dims] - - # create layers sequentially - layers = [] - layers.append(nn.Linear(input_dim, hidden_dims_processed[0])) - layers.append(activation_mod) - - for layer_index in range(len(hidden_dims_processed) - 1): - layers.append(nn.Linear(hidden_dims_processed[layer_index], hidden_dims_processed[layer_index + 1])) - layers.append(activation_mod) - - # add last layer - if isinstance(output_dim, int): - layers.append(nn.Linear(hidden_dims_processed[-1], output_dim)) - else: - # compute the total output dimension - total_out_dim = reduce(lambda x, y: x * y, output_dim) - # add a layer to reshape the output to the desired shape - layers.append(nn.Linear(hidden_dims_processed[-1], total_out_dim)) - layers.append(nn.Unflatten(dim=-1, unflattened_size=output_dim)) - - # add last activation function if specified - if last_activation_mod is not None: - layers.append(last_activation_mod) - - # register the layers - for idx, layer in enumerate(layers): - self.add_module(f"{idx}", layer) - - def init_weights(self, scales: float | tuple[float]): - """Initialize the weights of the MLP. - - Args: - scales: Scale factor for the weights. - """ - - def get_scale(idx) -> float: - """Get the scale factor for the weights of the MLP. - - Args: - idx: Index of the layer. - """ - return scales[idx] if isinstance(scales, (list, tuple)) else scales - - # initialize the weights - for idx, module in enumerate(self): - if isinstance(module, nn.Linear): - nn.init.orthogonal_(module.weight, gain=get_scale(idx)) - nn.init.zeros_(module.bias) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - """Forward pass of the MLP. - - Args: - x: Input tensor. - """ - for layer in self: - x = layer(x) - return x - - def reset(self, dones=None, hidden_states=None): - pass - - def detach_hidden_states(self, dones=None): - pass diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/networks/normalization.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/networks/normalization.py deleted file mode 100644 index 5fd9692..0000000 --- a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/networks/normalization.py +++ /dev/null @@ -1,130 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -# Copyright (c) 2020 Preferred Networks, Inc. - -from __future__ import annotations - -import torch -from torch import nn - - -class EmpiricalNormalization(nn.Module): - """Normalize mean and variance of values based on empirical values.""" - - def __init__(self, shape, eps=1e-2, until=None): - """Initialize EmpiricalNormalization module. - - Args: - shape (int or tuple of int): Shape of input values except batch axis. - eps (float): Small value for stability. - until (int or None): If this arg is specified, the module learns input values until the sum of batch sizes - exceeds it. - - Note: The normalization parameters are computed over the whole batch, not for each environment separately. - """ - super().__init__() - self.eps = eps - self.until = until - self.register_buffer("_mean", torch.zeros(shape).unsqueeze(0)) - self.register_buffer("_var", torch.ones(shape).unsqueeze(0)) - self.register_buffer("_std", torch.ones(shape).unsqueeze(0)) - self.register_buffer("count", torch.tensor(0, dtype=torch.long)) - - @property - def mean(self): - return self._mean.squeeze(0).clone() - - @property - def std(self): - return self._std.squeeze(0).clone() - - def forward(self, x): - """Normalize mean and variance of values based on empirical values.""" - - return (x - self._mean) / (self._std + self.eps) - - @torch.jit.unused - def update(self, x): - """Learn input values without computing the output values of them""" - - if not self.training: - return - if self.until is not None and self.count >= self.until: - return - - count_x = x.shape[0] - self.count += count_x - rate = count_x / self.count - var_x = torch.var(x, dim=0, unbiased=False, keepdim=True) - mean_x = torch.mean(x, dim=0, keepdim=True) - delta_mean = mean_x - self._mean - self._mean += rate * delta_mean - self._var += rate * (var_x - self._var + delta_mean * (mean_x - self._mean)) - self._std = torch.sqrt(self._var) - - @torch.jit.unused - def inverse(self, y): - """De-normalize values based on empirical values.""" - - return y * (self._std + self.eps) + self._mean - - -class EmpiricalDiscountedVariationNormalization(nn.Module): - """Reward normalization from Pathak's large scale study on PPO. - - Reward normalization. Since the reward function is non-stationary, it is useful to normalize - the scale of the rewards so that the value function can learn quickly. We did this by dividing - the rewards by a running estimate of the standard deviation of the sum of discounted rewards. - """ - - def __init__(self, shape, eps=1e-2, gamma=0.99, until=None): - super().__init__() - - self.emp_norm = EmpiricalNormalization(shape, eps, until) - self.disc_avg = _DiscountedAverage(gamma) - - def forward(self, rew): - if self.training: - # update discounted rewards - avg = self.disc_avg.update(rew) - # update moments from discounted rewards - self.emp_norm.update(avg) - - # normalize rewards with the empirical std - if self.emp_norm._std > 0: - return rew / self.emp_norm._std - else: - return rew - - -""" -Helper class. -""" - - -class _DiscountedAverage: - r"""Discounted average of rewards. - - The discounted average is defined as: - - .. math:: - - \bar{R}_t = \gamma \bar{R}_{t-1} + r_t - - Args: - gamma (float): Discount factor. - """ - - def __init__(self, gamma): - self.avg = None - self.gamma = gamma - - def update(self, rew: torch.Tensor) -> torch.Tensor: - if self.avg is None: - self.avg = rew - else: - self.avg = self.avg * self.gamma + rew - return self.avg diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/runners/__init__.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/runners/__init__.py deleted file mode 100644 index 61f1682..0000000 --- a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/runners/__init__.py +++ /dev/null @@ -1,12 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -"""Implementation of runners for environment-agent interaction.""" - -from .on_policy_runner import OnPolicyRunner # isort:skip -from .distillation_runner import DistillationRunner -from .amp_on_policy_runner import AMPOnPolicyRunner - -__all__ = ["OnPolicyRunner", "DistillationRunner", "AMPOnPolicyRunner"] diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/runners/amp_on_policy_runner.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/runners/amp_on_policy_runner.py deleted file mode 100644 index c0b9b9e..0000000 --- a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/runners/amp_on_policy_runner.py +++ /dev/null @@ -1,521 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -from __future__ import annotations - -import os -import statistics -import time -import torch -import warnings -from collections import deque - -import rsl_rl -from rsl_rl.algorithms import AMP_PPO -from rsl_rl.env import VecEnv -from rsl_rl.modules import ActorCritic, ActorCriticRecurrent,DiscriminatorMulti, resolve_rnd_config, resolve_symmetry_config -from rsl_rl.utils import resolve_obs_groups, store_code_state, Normalizer, G1_AMPLoader - - -class AMPOnPolicyRunner: - """On-policy runner for training and evaluation of actor-critic methods.""" - - def __init__(self, env: VecEnv, train_cfg: dict, log_dir: str | None = None, device="cpu"): - self.cfg = train_cfg - self.alg_cfg = train_cfg["algorithm"] - self.policy_cfg = train_cfg["policy"] - self.device = device - self.env = env - - # check if multi-gpu is enabled - self._configure_multi_gpu() - - # store training configuration - self.num_steps_per_env = self.cfg["num_steps_per_env"] - self.save_interval = self.cfg["save_interval"] - - # query observations from environment for algorithm construction - obs = self.env.get_observations() - default_sets = ["critic"] - if "rnd_cfg" in self.alg_cfg and self.alg_cfg["rnd_cfg"] is not None: - default_sets.append("rnd_state") - self.cfg["obs_groups"] = resolve_obs_groups(obs, self.cfg["obs_groups"], default_sets) - - self.amp_data = G1_AMPLoader( - device, - time_between_frames=1/50.0, - preload_transitions=True, - num_preload_transitions=train_cfg["amp_num_preload_transitions"], - motion_files=train_cfg["amp_motion_files"], - num_frames=train_cfg['amp_num_frames'] - ) - - self.amp_observation_dim = self.amp_data.observation_dim if self.cfg["amp_num_obs"] == 0 else self.cfg["amp_num_obs"] - self.amp_num_frames = 0 if self.cfg["amp_num_frames"] == 0 else self.cfg["amp_num_frames"] - self.amp_normalizer = Normalizer(self.amp_observation_dim) - self.discriminator = DiscriminatorMulti( - self.amp_observation_dim, - train_cfg["amp_reward_coef"], - train_cfg["amp_discr_hidden_dims"], - device, - train_cfg["amp_num_frames"], - train_cfg["amp_task_reward_lerp"], - train_cfg['use_lerp'], - ).to(self.device) - - - # create the algorithm - self.alg = self._construct_algorithm(obs) - - # Decide whether to disable logging - # We only log from the process with rank 0 (main process) - self.disable_logs = self.is_distributed and self.gpu_global_rank != 0 - - # Logging - self.log_dir = log_dir - self.writer = None - self.tot_timesteps = 0 - self.tot_time = 0 - self.current_learning_iteration = 0 - self.git_status_repos = [rsl_rl.__file__] - - def learn(self, num_learning_iterations: int, init_at_random_ep_len: bool = False): # noqa: C901 - # initialize writer - self._prepare_logging_writer() - - # randomize initial episode lengths (for exploration) - if init_at_random_ep_len: - self.env.episode_length_buf = torch.randint_like( - self.env.episode_length_buf, high=int(self.env.max_episode_length) - ) - - # start learning - obs = self.env.get_observations().to(self.device) - self.train_mode() # switch to train mode (for dropout for example) - amp_obs = self.env.get_amp_observations() - amp_obs = amp_obs.to(self.device) - if self.amp_num_frames != 0: - self.amp_obs_frames = torch.zeros(size=(self.env.num_envs, self.amp_num_frames, self.amp_observation_dim), device=self.device) - self.amp_obs_frames = torch.concat((self.amp_obs_frames[:, 1:], amp_obs.unsqueeze(1)), dim=1) - - # Book keeping - ep_infos = [] - rewbuffer = deque(maxlen=100) - lenbuffer = deque(maxlen=100) - step_discrewbuffer = deque(maxlen=100) - - cur_reward_sum = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) - cur_episode_length = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) - cur_single_step_disc_rew = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) - # create buffers for logging extrinsic and intrinsic rewards - if self.alg.rnd: - erewbuffer = deque(maxlen=100) - irewbuffer = deque(maxlen=100) - cur_ereward_sum = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) - cur_ireward_sum = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) - - # Ensure all parameters are in-synced - if self.is_distributed: - print(f"Synchronizing parameters for rank {self.gpu_global_rank}...") - self.alg.broadcast_parameters() - - # Start training - start_iter = self.current_learning_iteration - tot_iter = start_iter + num_learning_iterations - for it in range(start_iter, tot_iter): - start = time.time() - # Rollout - with torch.inference_mode(): - for _ in range(self.num_steps_per_env): - # Sample actions - actions = self.alg.act(obs,amp_obs) - # Step the environment - obs, rewards, dones, extras = self.env.step(actions.to(self.env.device)) - # Move to device - obs, rewards, dones = (obs.to(self.device), rewards.to(self.device), dones.to(self.device)) - - next_amp_obs = self.env.get_amp_observations() - next_amp_obs = next_amp_obs.to(self.device) - next_amp_obs_with_term = torch.clone(next_amp_obs) - - reset_env_ids = self.env.reset_env_ids - terminal_amp_states = self.env.get_amp_observations()[reset_env_ids] - next_amp_obs_with_term[reset_env_ids] = terminal_amp_states - self.amp_obs_frames = torch.concat((self.amp_obs_frames[:, 1:], next_amp_obs_with_term.unsqueeze(1)), dim=1) - - amp_reward = torch.zeros(self.env.num_envs, device=obs.device) - - mask = self.env.contact_phase[:, 0] == 1.0 - if mask.any(): - rewards[mask], logit, disc_reward = self.alg.discriminator.predict_amp_reward( - self.amp_obs_frames[mask], rewards[mask], normalizer=self.alg.amp_normalizer - ) - amp_reward[mask] += disc_reward - - # process the step - self.alg.process_env_step(obs, rewards, dones, extras, next_amp_obs_with_term, self.amp_obs_frames) - self.amp_obs_frames[reset_env_ids] = 0 - - amp_obs = torch.clone(next_amp_obs) - # Extract intrinsic rewards (only for logging) - intrinsic_rewards = self.alg.intrinsic_rewards if self.alg.rnd else None - # book keeping - if self.log_dir is not None: - if "episode" in extras: - ep_infos.append(extras["episode"]) - elif "log" in extras: - ep_infos.append(extras["log"]) - # Update rewards - if self.alg.rnd: - cur_ereward_sum += rewards - cur_ireward_sum += intrinsic_rewards # type: ignore - cur_reward_sum += rewards + intrinsic_rewards - else: - cur_reward_sum += rewards - # Update episode length - cur_episode_length += 1 - cur_single_step_disc_rew += amp_reward - # Clear data for completed episodes - # -- common - new_ids = (dones > 0).nonzero(as_tuple=False) - rewbuffer.extend(cur_reward_sum[new_ids][:, 0].cpu().numpy().tolist()) - lenbuffer.extend(cur_episode_length[new_ids][:, 0].cpu().numpy().tolist()) - cur_reward_sum[new_ids] = 0 - cur_episode_length[new_ids] = 0 - to_extend_disc = (cur_single_step_disc_rew[new_ids] / self.env.max_episode_length_s)[:, 0].cpu().numpy() - step_discrewbuffer.extend(to_extend_disc.tolist()) - cur_single_step_disc_rew[new_ids] = 0 - # -- intrinsic and extrinsic rewards - if self.alg.rnd: - erewbuffer.extend(cur_ereward_sum[new_ids][:, 0].cpu().numpy().tolist()) - irewbuffer.extend(cur_ireward_sum[new_ids][:, 0].cpu().numpy().tolist()) - cur_ereward_sum[new_ids] = 0 - cur_ireward_sum[new_ids] = 0 - - stop = time.time() - collection_time = stop - start - start = stop - - # compute returns - self.alg.compute_returns(obs) - - # update policy - loss_dict = self.alg.update() - - stop = time.time() - learn_time = stop - start - self.current_learning_iteration = it - # log info - if self.log_dir is not None and not self.disable_logs: - # Log information - self.log(locals()) - # Save model - if it % self.save_interval == 0: - self.save(os.path.join(self.log_dir, f"model_{it}.pt")) - - # Clear episode infos - ep_infos.clear() - # Save code state - if it == start_iter and not self.disable_logs: - # obtain all the diff files - git_file_paths = store_code_state(self.log_dir, self.git_status_repos) - # if possible store them to wandb - if self.logger_type in ["wandb", "neptune"] and git_file_paths: - for path in git_file_paths: - self.writer.save_file(path) - - # Save the final model after training - if self.log_dir is not None and not self.disable_logs: - self.save(os.path.join(self.log_dir, f"model_{self.current_learning_iteration}.pt")) - - def log(self, locs: dict, width: int = 80, pad: int = 35): - # Compute the collection size - collection_size = self.num_steps_per_env * self.env.num_envs * self.gpu_world_size - # Update total time-steps and time - self.tot_timesteps += collection_size - self.tot_time += locs["collection_time"] + locs["learn_time"] - iteration_time = locs["collection_time"] + locs["learn_time"] - - # -- Episode info - ep_string = "" - if locs["ep_infos"]: - for key in locs["ep_infos"][0]: - infotensor = torch.tensor([], device=self.device) - for ep_info in locs["ep_infos"]: - # handle scalar and zero dimensional tensor infos - if key not in ep_info: - continue - if not isinstance(ep_info[key], torch.Tensor): - ep_info[key] = torch.Tensor([ep_info[key]]) - if len(ep_info[key].shape) == 0: - ep_info[key] = ep_info[key].unsqueeze(0) - infotensor = torch.cat((infotensor, ep_info[key].to(self.device))) - value = torch.mean(infotensor) - # log to logger and terminal - if "/" in key: - self.writer.add_scalar(key, value, locs["it"]) - ep_string += f"""{f'{key}:':>{pad}} {value:.4f}\n""" - else: - self.writer.add_scalar("Episode/" + key, value, locs["it"]) - ep_string += f"""{f'Mean episode {key}:':>{pad}} {value:.4f}\n""" - - mean_std = self.alg.policy.action_std.mean() - fps = int(collection_size / (locs["collection_time"] + locs["learn_time"])) - - # -- Losses - for key, value in locs["loss_dict"].items(): - self.writer.add_scalar(f"Loss/{key}", value, locs["it"]) - self.writer.add_scalar("Loss/learning_rate", self.alg.learning_rate, locs["it"]) - - # -- Policy - self.writer.add_scalar("Policy/mean_noise_std", mean_std.item(), locs["it"]) - - # -- Performance - self.writer.add_scalar("Perf/total_fps", fps, locs["it"]) - self.writer.add_scalar("Perf/collection time", locs["collection_time"], locs["it"]) - self.writer.add_scalar("Perf/learning_time", locs["learn_time"], locs["it"]) - - # -- Training - if len(locs["rewbuffer"]) > 0: - # separate logging for intrinsic and extrinsic rewards - if hasattr(self.alg, "rnd") and self.alg.rnd: - self.writer.add_scalar("Rnd/mean_extrinsic_reward", statistics.mean(locs["erewbuffer"]), locs["it"]) - self.writer.add_scalar("Rnd/mean_intrinsic_reward", statistics.mean(locs["irewbuffer"]), locs["it"]) - self.writer.add_scalar("Rnd/weight", self.alg.rnd.weight, locs["it"]) - # everything else - self.writer.add_scalar("Train/mean_reward", statistics.mean(locs["rewbuffer"]), locs["it"]) - self.writer.add_scalar("Train/mean_episode_length", statistics.mean(locs["lenbuffer"]), locs["it"]) - self.writer.add_scalar('Train/mean_step_disc_reward', statistics.mean(locs['step_discrewbuffer']), locs['it']) - if self.logger_type != "wandb": # wandb does not support non-integer x-axis logging - self.writer.add_scalar("Train/mean_reward/time", statistics.mean(locs["rewbuffer"]), self.tot_time) - self.writer.add_scalar( - "Train/mean_episode_length/time", statistics.mean(locs["lenbuffer"]), self.tot_time - ) - - str = f" \033[1m Learning iteration {locs['it']}/{locs['tot_iter']} \033[0m " - - if len(locs["rewbuffer"]) > 0: - log_string = ( - f"""{'#' * width}\n""" - f"""{str.center(width, ' ')}\n\n""" - f"""{'Computation:':>{pad}} {fps:.0f} steps/s (collection: {locs[ - 'collection_time']:.3f}s, learning {locs['learn_time']:.3f}s)\n""" - f"""{'Mean action noise std:':>{pad}} {mean_std.item():.2f}\n""" - f"""{'Step disc reward:':>{pad}} {statistics.mean(locs['step_discrewbuffer']):.2f}\n""" - ) - # -- Losses - for key, value in locs["loss_dict"].items(): - log_string += f"""{f'Mean {key} loss:':>{pad}} {value:.4f}\n""" - # -- Rewards - if hasattr(self.alg, "rnd") and self.alg.rnd: - log_string += ( - f"""{'Mean extrinsic reward:':>{pad}} {statistics.mean(locs['erewbuffer']):.2f}\n""" - f"""{'Mean intrinsic reward:':>{pad}} {statistics.mean(locs['irewbuffer']):.2f}\n""" - ) - log_string += f"""{'Mean reward:':>{pad}} {statistics.mean(locs['rewbuffer']):.2f}\n""" - # -- episode info - log_string += f"""{'Mean episode length:':>{pad}} {statistics.mean(locs['lenbuffer']):.2f}\n""" - else: - log_string = ( - f"""{'#' * width}\n""" - f"""{str.center(width, ' ')}\n\n""" - f"""{'Computation:':>{pad}} {fps:.0f} steps/s (collection: {locs[ - 'collection_time']:.3f}s, learning {locs['learn_time']:.3f}s)\n""" - f"""{'Mean action noise std:':>{pad}} {mean_std.item():.2f}\n""" - ) - for key, value in locs["loss_dict"].items(): - log_string += f"""{f'{key}:':>{pad}} {value:.4f}\n""" - - log_string += ep_string - log_string += ( - f"""{'-' * width}\n""" - f"""{'Total timesteps:':>{pad}} {self.tot_timesteps}\n""" - f"""{'Iteration time:':>{pad}} {iteration_time:.2f}s\n""" - f"""{'Time elapsed:':>{pad}} {time.strftime("%H:%M:%S", time.gmtime(self.tot_time))}\n""" - f"""{'ETA:':>{pad}} {time.strftime( - "%H:%M:%S", - time.gmtime( - self.tot_time / (locs['it'] - locs['start_iter'] + 1) - * (locs['start_iter'] + locs['num_learning_iterations'] - locs['it']) - ) - )}\n""" - ) - print(log_string) - - def save(self, path: str, infos=None): - # -- Save model - saved_dict = { - "model_state_dict": self.alg.policy.state_dict(), - "optimizer_state_dict": self.alg.optimizer.state_dict(), - "iter": self.current_learning_iteration, - "infos": infos, - } - # -- Save RND model if used - if hasattr(self.alg, "rnd") and self.alg.rnd: - saved_dict["rnd_state_dict"] = self.alg.rnd.state_dict() - saved_dict["rnd_optimizer_state_dict"] = self.alg.rnd_optimizer.state_dict() - torch.save(saved_dict, path) - - # upload model to external logging service - if self.logger_type in ["neptune", "wandb"] and not self.disable_logs: - self.writer.save_model(path, self.current_learning_iteration) - - def load(self, path: str, load_optimizer: bool = True, map_location: str | None = None): - loaded_dict = torch.load(path, weights_only=False, map_location=map_location) - # -- Load model - resumed_training = self.alg.policy.load_state_dict(loaded_dict["model_state_dict"]) - # -- Load RND model if used - if hasattr(self.alg, "rnd") and self.alg.rnd: - self.alg.rnd.load_state_dict(loaded_dict["rnd_state_dict"]) - # -- load optimizer if used - if load_optimizer and resumed_training: - # -- algorithm optimizer - self.alg.optimizer.load_state_dict(loaded_dict["optimizer_state_dict"]) - # -- RND optimizer if used - if hasattr(self.alg, "rnd") and self.alg.rnd: - self.alg.rnd_optimizer.load_state_dict(loaded_dict["rnd_optimizer_state_dict"]) - # -- load current learning iteration - if resumed_training: - self.current_learning_iteration = loaded_dict["iter"] - return loaded_dict["infos"] - - def get_inference_policy(self, device=None): - self.eval_mode() # switch to evaluation mode (dropout for example) - if device is not None: - self.alg.policy.to(device) - return self.alg.policy.act_inference - - def train_mode(self): - # -- PPO - self.alg.policy.train() - self.alg.discriminator.train() - # -- RND - if hasattr(self.alg, "rnd") and self.alg.rnd: - self.alg.rnd.train() - - def eval_mode(self): - # -- PPO - self.alg.policy.eval() - self.alg.discriminator.eval() - # -- RND - if hasattr(self.alg, "rnd") and self.alg.rnd: - self.alg.rnd.eval() - - def add_git_repo_to_log(self, repo_file_path): - self.git_status_repos.append(repo_file_path) - - """ - Helper functions. - """ - - def _configure_multi_gpu(self): - """Configure multi-gpu training.""" - # check if distributed training is enabled - self.gpu_world_size = int(os.getenv("WORLD_SIZE", "1")) - self.is_distributed = self.gpu_world_size > 1 - - # if not distributed training, set local and global rank to 0 and return - if not self.is_distributed: - self.gpu_local_rank = 0 - self.gpu_global_rank = 0 - self.multi_gpu_cfg = None - return - - # get rank and world size - self.gpu_local_rank = int(os.getenv("LOCAL_RANK", "0")) - self.gpu_global_rank = int(os.getenv("RANK", "0")) - - # make a configuration dictionary - self.multi_gpu_cfg = { - "global_rank": self.gpu_global_rank, # rank of the main process - "local_rank": self.gpu_local_rank, # rank of the current process - "world_size": self.gpu_world_size, # total number of processes - } - - # check if user has device specified for local rank - if self.device != f"cuda:{self.gpu_local_rank}": - raise ValueError( - f"Device '{self.device}' does not match expected device for local rank '{self.gpu_local_rank}'." - ) - # validate multi-gpu configuration - if self.gpu_local_rank >= self.gpu_world_size: - raise ValueError( - f"Local rank '{self.gpu_local_rank}' is greater than or equal to world size '{self.gpu_world_size}'." - ) - if self.gpu_global_rank >= self.gpu_world_size: - raise ValueError( - f"Global rank '{self.gpu_global_rank}' is greater than or equal to world size '{self.gpu_world_size}'." - ) - - # initialize torch distributed - torch.distributed.init_process_group(backend="nccl", rank=self.gpu_global_rank, world_size=self.gpu_world_size) - # set device to the local rank - torch.cuda.set_device(self.gpu_local_rank) - - def _construct_algorithm(self, obs) -> AMP_PPO: - """Construct the actor-critic algorithm.""" - # resolve RND config - self.alg_cfg = resolve_rnd_config(self.alg_cfg, obs, self.cfg["obs_groups"], self.env) - - # resolve symmetry config - self.alg_cfg = resolve_symmetry_config(self.alg_cfg, self.env) - - # resolve deprecated normalization config - if self.cfg.get("empirical_normalization") is not None: - warnings.warn( - "The `empirical_normalization` parameter is deprecated. Please set `actor_obs_normalization` and " - "`critic_obs_normalization` as part of the `policy` configuration instead.", - DeprecationWarning, - ) - if self.policy_cfg.get("actor_obs_normalization") is None: - self.policy_cfg["actor_obs_normalization"] = self.cfg["empirical_normalization"] - if self.policy_cfg.get("critic_obs_normalization") is None: - self.policy_cfg["critic_obs_normalization"] = self.cfg["empirical_normalization"] - - # initialize the actor-critic - actor_critic_class = eval(self.policy_cfg.pop("class_name")) - actor_critic: ActorCritic | ActorCriticRecurrent = actor_critic_class( - obs, self.cfg["obs_groups"], self.env.num_actions, **self.policy_cfg - ).to(self.device) - - # initialize the algorithm - alg_class = eval(self.alg_cfg.pop("class_name")) - - alg: AMP_PPO = alg_class(actor_critic, self.discriminator, self.amp_data, self.amp_normalizer, self.amp_num_frames, device=self.device, **self.alg_cfg, multi_gpu_cfg=self.multi_gpu_cfg) - - # initialize the storage - alg.init_storage( - "rl", - self.env.num_envs, - self.num_steps_per_env, - obs, - [self.env.num_actions], - ) - - return alg - - def _prepare_logging_writer(self): - """Prepares the logging writers.""" - if self.log_dir is not None and self.writer is None and not self.disable_logs: - # Launch either Tensorboard or Neptune & Tensorboard summary writer(s), default: Tensorboard. - self.logger_type = self.cfg.get("logger", "tensorboard") - self.logger_type = self.logger_type.lower() - - if self.logger_type == "neptune": - from rsl_rl.utils.neptune_utils import NeptuneSummaryWriter - - self.writer = NeptuneSummaryWriter(log_dir=self.log_dir, flush_secs=10, cfg=self.cfg) - self.writer.log_config(self.env.cfg, self.cfg, self.alg_cfg, self.policy_cfg) - elif self.logger_type == "wandb": - from rsl_rl.utils.wandb_utils import WandbSummaryWriter - - self.writer = WandbSummaryWriter(log_dir=self.log_dir, flush_secs=10, cfg=self.cfg) - self.writer.log_config(self.env.cfg, self.cfg, self.alg_cfg, self.policy_cfg) - elif self.logger_type == "tensorboard": - from torch.utils.tensorboard import SummaryWriter - - self.writer = SummaryWriter(log_dir=self.log_dir, flush_secs=10) - else: - raise ValueError("Logger type not found. Please choose 'neptune', 'wandb' or 'tensorboard'.") diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/runners/distillation_runner.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/runners/distillation_runner.py deleted file mode 100644 index 9cc6a8b..0000000 --- a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/runners/distillation_runner.py +++ /dev/null @@ -1,179 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -from __future__ import annotations - -import os -import time -import torch -from collections import deque - -import rsl_rl -from rsl_rl.algorithms import Distillation -from rsl_rl.env import VecEnv -from rsl_rl.modules import StudentTeacher, StudentTeacherRecurrent -from rsl_rl.runners import OnPolicyRunner -from rsl_rl.utils import resolve_obs_groups, store_code_state - - -class DistillationRunner(OnPolicyRunner): - """On-policy runner for training and evaluation of teacher-student training.""" - - def __init__(self, env: VecEnv, train_cfg: dict, log_dir: str | None = None, device="cpu"): - self.cfg = train_cfg - self.alg_cfg = train_cfg["algorithm"] - self.policy_cfg = train_cfg["policy"] - self.device = device - self.env = env - - # check if multi-gpu is enabled - self._configure_multi_gpu() - - # store training configuration - self.num_steps_per_env = self.cfg["num_steps_per_env"] - self.save_interval = self.cfg["save_interval"] - - # query observations from environment for algorithm construction - obs = self.env.get_observations() - self.cfg["obs_groups"] = resolve_obs_groups(obs, self.cfg["obs_groups"], default_sets=["teacher"]) - - # create the algorithm - self.alg = self._construct_algorithm(obs) - - # Decide whether to disable logging - # We only log from the process with rank 0 (main process) - self.disable_logs = self.is_distributed and self.gpu_global_rank != 0 - - # Logging - self.log_dir = log_dir - self.writer = None - self.tot_timesteps = 0 - self.tot_time = 0 - self.current_learning_iteration = 0 - self.git_status_repos = [rsl_rl.__file__] - - def learn(self, num_learning_iterations: int, init_at_random_ep_len: bool = False): # noqa: C901 - # initialize writer - self._prepare_logging_writer() - # check if teacher is loaded - if not self.alg.policy.loaded_teacher: - raise ValueError("Teacher model parameters not loaded. Please load a teacher model to distill.") - - # randomize initial episode lengths (for exploration) - if init_at_random_ep_len: - self.env.episode_length_buf = torch.randint_like( - self.env.episode_length_buf, high=int(self.env.max_episode_length) - ) - - # start learning - obs = self.env.get_observations().to(self.device) - self.train_mode() # switch to train mode (for dropout for example) - - # Book keeping - ep_infos = [] - rewbuffer = deque(maxlen=100) - lenbuffer = deque(maxlen=100) - cur_reward_sum = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) - cur_episode_length = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) - - # Ensure all parameters are in-synced - if self.is_distributed: - print(f"Synchronizing parameters for rank {self.gpu_global_rank}...") - self.alg.broadcast_parameters() - - # Start training - start_iter = self.current_learning_iteration - tot_iter = start_iter + num_learning_iterations - for it in range(start_iter, tot_iter): - start = time.time() - # Rollout - with torch.inference_mode(): - for _ in range(self.num_steps_per_env): - # Sample actions - actions = self.alg.act(obs) - # Step the environment - obs, rewards, dones, extras = self.env.step(actions.to(self.env.device)) - # Move to device - obs, rewards, dones = (obs.to(self.device), rewards.to(self.device), dones.to(self.device)) - # process the step - self.alg.process_env_step(obs, rewards, dones, extras) - # book keeping - if self.log_dir is not None: - if "episode" in extras: - ep_infos.append(extras["episode"]) - elif "log" in extras: - ep_infos.append(extras["log"]) - # Update rewards - cur_reward_sum += rewards - # Update episode length - cur_episode_length += 1 - # Clear data for completed episodes - new_ids = (dones > 0).nonzero(as_tuple=False) - rewbuffer.extend(cur_reward_sum[new_ids][:, 0].cpu().numpy().tolist()) - lenbuffer.extend(cur_episode_length[new_ids][:, 0].cpu().numpy().tolist()) - cur_reward_sum[new_ids] = 0 - cur_episode_length[new_ids] = 0 - - stop = time.time() - collection_time = stop - start - start = stop - - # update policy - loss_dict = self.alg.update() - - stop = time.time() - learn_time = stop - start - self.current_learning_iteration = it - # log info - if self.log_dir is not None and not self.disable_logs: - # Log information - self.log(locals()) - # Save model - if it % self.save_interval == 0: - self.save(os.path.join(self.log_dir, f"model_{it}.pt")) - - # Clear episode infos - ep_infos.clear() - # Save code state - if it == start_iter and not self.disable_logs: - # obtain all the diff files - git_file_paths = store_code_state(self.log_dir, self.git_status_repos) - # if possible store them to wandb - if self.logger_type in ["wandb", "neptune"] and git_file_paths: - for path in git_file_paths: - self.writer.save_file(path) - - # Save the final model after training - if self.log_dir is not None and not self.disable_logs: - self.save(os.path.join(self.log_dir, f"model_{self.current_learning_iteration}.pt")) - - """ - Helper methods. - """ - - def _construct_algorithm(self, obs) -> Distillation: - """Construct the distillation algorithm.""" - # initialize the actor-critic - student_teacher_class = eval(self.policy_cfg.pop("class_name")) - student_teacher: StudentTeacher | StudentTeacherRecurrent = student_teacher_class( - obs, self.cfg["obs_groups"], self.env.num_actions, **self.policy_cfg - ).to(self.device) - - # initialize the algorithm - alg_class = eval(self.alg_cfg.pop("class_name")) - alg: Distillation = alg_class( - student_teacher, device=self.device, **self.alg_cfg, multi_gpu_cfg=self.multi_gpu_cfg - ) - - # initialize the storage - alg.init_storage( - "distillation", - self.env.num_envs, - self.num_steps_per_env, - obs, - [self.env.num_actions], - ) - - return alg diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/runners/on_policy_runner.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/runners/on_policy_runner.py deleted file mode 100644 index 36f11f3..0000000 --- a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/runners/on_policy_runner.py +++ /dev/null @@ -1,460 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -from __future__ import annotations - -import os -import statistics -import time -import torch -import warnings -from collections import deque - -import rsl_rl -from rsl_rl.algorithms import PPO -from rsl_rl.env import VecEnv -from rsl_rl.modules import ActorCritic, ActorCriticRecurrent, resolve_rnd_config, resolve_symmetry_config -from rsl_rl.utils import resolve_obs_groups, store_code_state - - -class OnPolicyRunner: - """On-policy runner for training and evaluation of actor-critic methods.""" - - def __init__(self, env: VecEnv, train_cfg: dict, log_dir: str | None = None, device="cpu"): - self.cfg = train_cfg - self.alg_cfg = train_cfg["algorithm"] - self.policy_cfg = train_cfg["policy"] - self.device = device - self.env = env - - # check if multi-gpu is enabled - self._configure_multi_gpu() - - # store training configuration - self.num_steps_per_env = self.cfg["num_steps_per_env"] - self.save_interval = self.cfg["save_interval"] - - # query observations from environment for algorithm construction - obs = self.env.get_observations() - default_sets = ["critic"] - if "rnd_cfg" in self.alg_cfg and self.alg_cfg["rnd_cfg"] is not None: - default_sets.append("rnd_state") - self.cfg["obs_groups"] = resolve_obs_groups(obs, self.cfg["obs_groups"], default_sets) - - # create the algorithm - self.alg = self._construct_algorithm(obs) - - # Decide whether to disable logging - # We only log from the process with rank 0 (main process) - self.disable_logs = self.is_distributed and self.gpu_global_rank != 0 - - # Logging - self.log_dir = log_dir - self.writer = None - self.tot_timesteps = 0 - self.tot_time = 0 - self.current_learning_iteration = 0 - self.git_status_repos = [rsl_rl.__file__] - - def learn(self, num_learning_iterations: int, init_at_random_ep_len: bool = False): # noqa: C901 - # initialize writer - self._prepare_logging_writer() - - # randomize initial episode lengths (for exploration) - if init_at_random_ep_len: - self.env.episode_length_buf = torch.randint_like( - self.env.episode_length_buf, high=int(self.env.max_episode_length) - ) - - # start learning - obs = self.env.get_observations().to(self.device) - self.train_mode() # switch to train mode (for dropout for example) - - # Book keeping - ep_infos = [] - rewbuffer = deque(maxlen=100) - lenbuffer = deque(maxlen=100) - cur_reward_sum = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) - cur_episode_length = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) - - # create buffers for logging extrinsic and intrinsic rewards - if self.alg.rnd: - erewbuffer = deque(maxlen=100) - irewbuffer = deque(maxlen=100) - cur_ereward_sum = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) - cur_ireward_sum = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) - - # Ensure all parameters are in-synced - if self.is_distributed: - print(f"Synchronizing parameters for rank {self.gpu_global_rank}...") - self.alg.broadcast_parameters() - - # Start training - start_iter = self.current_learning_iteration - tot_iter = start_iter + num_learning_iterations - for it in range(start_iter, tot_iter): - start = time.time() - # Rollout - with torch.inference_mode(): - for _ in range(self.num_steps_per_env): - # Sample actions - actions = self.alg.act(obs) - # Step the environment - obs, rewards, dones, extras = self.env.step(actions.to(self.env.device)) - # Move to device - obs, rewards, dones = (obs.to(self.device), rewards.to(self.device), dones.to(self.device)) - # process the step - self.alg.process_env_step(obs, rewards, dones, extras) - # Extract intrinsic rewards (only for logging) - intrinsic_rewards = self.alg.intrinsic_rewards if self.alg.rnd else None - # book keeping - if self.log_dir is not None: - if "episode" in extras: - ep_infos.append(extras["episode"]) - elif "log" in extras: - ep_infos.append(extras["log"]) - # Update rewards - if self.alg.rnd: - cur_ereward_sum += rewards - cur_ireward_sum += intrinsic_rewards # type: ignore - cur_reward_sum += rewards + intrinsic_rewards - else: - cur_reward_sum += rewards - # Update episode length - cur_episode_length += 1 - # Clear data for completed episodes - # -- common - new_ids = (dones > 0).nonzero(as_tuple=False) - rewbuffer.extend(cur_reward_sum[new_ids][:, 0].cpu().numpy().tolist()) - lenbuffer.extend(cur_episode_length[new_ids][:, 0].cpu().numpy().tolist()) - cur_reward_sum[new_ids] = 0 - cur_episode_length[new_ids] = 0 - # -- intrinsic and extrinsic rewards - if self.alg.rnd: - erewbuffer.extend(cur_ereward_sum[new_ids][:, 0].cpu().numpy().tolist()) - irewbuffer.extend(cur_ireward_sum[new_ids][:, 0].cpu().numpy().tolist()) - cur_ereward_sum[new_ids] = 0 - cur_ireward_sum[new_ids] = 0 - - stop = time.time() - collection_time = stop - start - start = stop - - # compute returns - self.alg.compute_returns(obs) - - # update policy - loss_dict = self.alg.update() - - stop = time.time() - learn_time = stop - start - self.current_learning_iteration = it - # log info - if self.log_dir is not None and not self.disable_logs: - # Log information - self.log(locals()) - # Save model - if it % self.save_interval == 0: - self.save(os.path.join(self.log_dir, f"model_{it}.pt")) - - # Clear episode infos - ep_infos.clear() - # Save code state - if it == start_iter and not self.disable_logs: - # obtain all the diff files - git_file_paths = store_code_state(self.log_dir, self.git_status_repos) - # if possible store them to wandb - if self.logger_type in ["wandb", "neptune"] and git_file_paths: - for path in git_file_paths: - self.writer.save_file(path) - - # Save the final model after training - if self.log_dir is not None and not self.disable_logs: - self.save(os.path.join(self.log_dir, f"model_{self.current_learning_iteration}.pt")) - - def log(self, locs: dict, width: int = 80, pad: int = 35): - # Compute the collection size - collection_size = self.num_steps_per_env * self.env.num_envs * self.gpu_world_size - # Update total time-steps and time - self.tot_timesteps += collection_size - self.tot_time += locs["collection_time"] + locs["learn_time"] - iteration_time = locs["collection_time"] + locs["learn_time"] - - # -- Episode info - ep_string = "" - if locs["ep_infos"]: - for key in locs["ep_infos"][0]: - infotensor = torch.tensor([], device=self.device) - for ep_info in locs["ep_infos"]: - # handle scalar and zero dimensional tensor infos - if key not in ep_info: - continue - if not isinstance(ep_info[key], torch.Tensor): - ep_info[key] = torch.Tensor([ep_info[key]]) - if len(ep_info[key].shape) == 0: - ep_info[key] = ep_info[key].unsqueeze(0) - infotensor = torch.cat((infotensor, ep_info[key].to(self.device))) - value = torch.mean(infotensor) - # log to logger and terminal - if "/" in key: - self.writer.add_scalar(key, value, locs["it"]) - ep_string += f"""{f'{key}:':>{pad}} {value:.4f}\n""" - else: - self.writer.add_scalar("Episode/" + key, value, locs["it"]) - ep_string += f"""{f'Mean episode {key}:':>{pad}} {value:.4f}\n""" - - mean_std = self.alg.policy.action_std.mean() - fps = int(collection_size / (locs["collection_time"] + locs["learn_time"])) - - # -- Losses - for key, value in locs["loss_dict"].items(): - self.writer.add_scalar(f"Loss/{key}", value, locs["it"]) - self.writer.add_scalar("Loss/learning_rate", self.alg.learning_rate, locs["it"]) - - # -- Policy - self.writer.add_scalar("Policy/mean_noise_std", mean_std.item(), locs["it"]) - - # -- Performance - self.writer.add_scalar("Perf/total_fps", fps, locs["it"]) - self.writer.add_scalar("Perf/collection time", locs["collection_time"], locs["it"]) - self.writer.add_scalar("Perf/learning_time", locs["learn_time"], locs["it"]) - - # -- Training - if len(locs["rewbuffer"]) > 0: - # separate logging for intrinsic and extrinsic rewards - if hasattr(self.alg, "rnd") and self.alg.rnd: - self.writer.add_scalar("Rnd/mean_extrinsic_reward", statistics.mean(locs["erewbuffer"]), locs["it"]) - self.writer.add_scalar("Rnd/mean_intrinsic_reward", statistics.mean(locs["irewbuffer"]), locs["it"]) - self.writer.add_scalar("Rnd/weight", self.alg.rnd.weight, locs["it"]) - # everything else - self.writer.add_scalar("Train/mean_reward", statistics.mean(locs["rewbuffer"]), locs["it"]) - self.writer.add_scalar("Train/mean_episode_length", statistics.mean(locs["lenbuffer"]), locs["it"]) - if self.logger_type != "wandb": # wandb does not support non-integer x-axis logging - self.writer.add_scalar("Train/mean_reward/time", statistics.mean(locs["rewbuffer"]), self.tot_time) - self.writer.add_scalar( - "Train/mean_episode_length/time", statistics.mean(locs["lenbuffer"]), self.tot_time - ) - - str = f" \033[1m Learning iteration {locs['it']}/{locs['tot_iter']} \033[0m " - - if len(locs["rewbuffer"]) > 0: - log_string = ( - f"""{'#' * width}\n""" - f"""{str.center(width, ' ')}\n\n""" - f"""{'Computation:':>{pad}} {fps:.0f} steps/s (collection: {locs[ - 'collection_time']:.3f}s, learning {locs['learn_time']:.3f}s)\n""" - f"""{'Mean action noise std:':>{pad}} {mean_std.item():.2f}\n""" - ) - # -- Losses - for key, value in locs["loss_dict"].items(): - log_string += f"""{f'Mean {key} loss:':>{pad}} {value:.4f}\n""" - # -- Rewards - if hasattr(self.alg, "rnd") and self.alg.rnd: - log_string += ( - f"""{'Mean extrinsic reward:':>{pad}} {statistics.mean(locs['erewbuffer']):.2f}\n""" - f"""{'Mean intrinsic reward:':>{pad}} {statistics.mean(locs['irewbuffer']):.2f}\n""" - ) - log_string += f"""{'Mean reward:':>{pad}} {statistics.mean(locs['rewbuffer']):.2f}\n""" - # -- episode info - log_string += f"""{'Mean episode length:':>{pad}} {statistics.mean(locs['lenbuffer']):.2f}\n""" - else: - log_string = ( - f"""{'#' * width}\n""" - f"""{str.center(width, ' ')}\n\n""" - f"""{'Computation:':>{pad}} {fps:.0f} steps/s (collection: {locs[ - 'collection_time']:.3f}s, learning {locs['learn_time']:.3f}s)\n""" - f"""{'Mean action noise std:':>{pad}} {mean_std.item():.2f}\n""" - ) - for key, value in locs["loss_dict"].items(): - log_string += f"""{f'{key}:':>{pad}} {value:.4f}\n""" - - log_string += ep_string - log_string += ( - f"""{'-' * width}\n""" - f"""{'Total timesteps:':>{pad}} {self.tot_timesteps}\n""" - f"""{'Iteration time:':>{pad}} {iteration_time:.2f}s\n""" - f"""{'Time elapsed:':>{pad}} {time.strftime("%H:%M:%S", time.gmtime(self.tot_time))}\n""" - f"""{'ETA:':>{pad}} {time.strftime( - "%H:%M:%S", - time.gmtime( - self.tot_time / (locs['it'] - locs['start_iter'] + 1) - * (locs['start_iter'] + locs['num_learning_iterations'] - locs['it']) - ) - )}\n""" - ) - print(log_string) - - def save(self, path: str, infos=None): - # -- Save model - saved_dict = { - "model_state_dict": self.alg.policy.state_dict(), - "optimizer_state_dict": self.alg.optimizer.state_dict(), - "iter": self.current_learning_iteration, - "infos": infos, - } - # -- Save RND model if used - if hasattr(self.alg, "rnd") and self.alg.rnd: - saved_dict["rnd_state_dict"] = self.alg.rnd.state_dict() - saved_dict["rnd_optimizer_state_dict"] = self.alg.rnd_optimizer.state_dict() - torch.save(saved_dict, path) - - # upload model to external logging service - if self.logger_type in ["neptune", "wandb"] and not self.disable_logs: - self.writer.save_model(path, self.current_learning_iteration) - - def load(self, path: str, load_optimizer: bool = True, map_location: str | None = None): - loaded_dict = torch.load(path, weights_only=False, map_location=map_location) - # -- Load model - resumed_training = self.alg.policy.load_state_dict(loaded_dict["model_state_dict"]) - # -- Load RND model if used - if hasattr(self.alg, "rnd") and self.alg.rnd: - self.alg.rnd.load_state_dict(loaded_dict["rnd_state_dict"]) - # -- load optimizer if used - if load_optimizer and resumed_training: - # -- algorithm optimizer - self.alg.optimizer.load_state_dict(loaded_dict["optimizer_state_dict"]) - # -- RND optimizer if used - if hasattr(self.alg, "rnd") and self.alg.rnd: - self.alg.rnd_optimizer.load_state_dict(loaded_dict["rnd_optimizer_state_dict"]) - # -- load current learning iteration - if resumed_training: - self.current_learning_iteration = loaded_dict["iter"] - return loaded_dict["infos"] - - def get_inference_policy(self, device=None): - self.eval_mode() # switch to evaluation mode (dropout for example) - if device is not None: - self.alg.policy.to(device) - return self.alg.policy.act_inference - - def train_mode(self): - # -- PPO - self.alg.policy.train() - # -- RND - if hasattr(self.alg, "rnd") and self.alg.rnd: - self.alg.rnd.train() - - def eval_mode(self): - # -- PPO - self.alg.policy.eval() - # -- RND - if hasattr(self.alg, "rnd") and self.alg.rnd: - self.alg.rnd.eval() - - def add_git_repo_to_log(self, repo_file_path): - self.git_status_repos.append(repo_file_path) - - """ - Helper functions. - """ - - def _configure_multi_gpu(self): - """Configure multi-gpu training.""" - # check if distributed training is enabled - self.gpu_world_size = int(os.getenv("WORLD_SIZE", "1")) - self.is_distributed = self.gpu_world_size > 1 - - # if not distributed training, set local and global rank to 0 and return - if not self.is_distributed: - self.gpu_local_rank = 0 - self.gpu_global_rank = 0 - self.multi_gpu_cfg = None - return - - # get rank and world size - self.gpu_local_rank = int(os.getenv("LOCAL_RANK", "0")) - self.gpu_global_rank = int(os.getenv("RANK", "0")) - - # make a configuration dictionary - self.multi_gpu_cfg = { - "global_rank": self.gpu_global_rank, # rank of the main process - "local_rank": self.gpu_local_rank, # rank of the current process - "world_size": self.gpu_world_size, # total number of processes - } - - # check if user has device specified for local rank - if self.device != f"cuda:{self.gpu_local_rank}": - raise ValueError( - f"Device '{self.device}' does not match expected device for local rank '{self.gpu_local_rank}'." - ) - # validate multi-gpu configuration - if self.gpu_local_rank >= self.gpu_world_size: - raise ValueError( - f"Local rank '{self.gpu_local_rank}' is greater than or equal to world size '{self.gpu_world_size}'." - ) - if self.gpu_global_rank >= self.gpu_world_size: - raise ValueError( - f"Global rank '{self.gpu_global_rank}' is greater than or equal to world size '{self.gpu_world_size}'." - ) - - # initialize torch distributed - torch.distributed.init_process_group(backend="nccl", rank=self.gpu_global_rank, world_size=self.gpu_world_size) - # set device to the local rank - torch.cuda.set_device(self.gpu_local_rank) - - def _construct_algorithm(self, obs) -> PPO: - """Construct the actor-critic algorithm.""" - # resolve RND config - self.alg_cfg = resolve_rnd_config(self.alg_cfg, obs, self.cfg["obs_groups"], self.env) - - # resolve symmetry config - self.alg_cfg = resolve_symmetry_config(self.alg_cfg, self.env) - - # resolve deprecated normalization config - if self.cfg.get("empirical_normalization") is not None: - warnings.warn( - "The `empirical_normalization` parameter is deprecated. Please set `actor_obs_normalization` and " - "`critic_obs_normalization` as part of the `policy` configuration instead.", - DeprecationWarning, - ) - if self.policy_cfg.get("actor_obs_normalization") is None: - self.policy_cfg["actor_obs_normalization"] = self.cfg["empirical_normalization"] - if self.policy_cfg.get("critic_obs_normalization") is None: - self.policy_cfg["critic_obs_normalization"] = self.cfg["empirical_normalization"] - - # initialize the actor-critic - actor_critic_class = eval(self.policy_cfg.pop("class_name")) - actor_critic: ActorCritic | ActorCriticRecurrent = actor_critic_class( - obs, self.cfg["obs_groups"], self.env.num_actions, **self.policy_cfg - ).to(self.device) - - # initialize the algorithm - alg_class = eval(self.alg_cfg.pop("class_name")) - alg: PPO = alg_class(actor_critic, device=self.device, **self.alg_cfg, multi_gpu_cfg=self.multi_gpu_cfg) - - # initialize the storage - alg.init_storage( - "rl", - self.env.num_envs, - self.num_steps_per_env, - obs, - [self.env.num_actions], - ) - - return alg - - def _prepare_logging_writer(self): - """Prepares the logging writers.""" - if self.log_dir is not None and self.writer is None and not self.disable_logs: - # Launch either Tensorboard or Neptune & Tensorboard summary writer(s), default: Tensorboard. - self.logger_type = self.cfg.get("logger", "tensorboard") - self.logger_type = self.logger_type.lower() - - if self.logger_type == "neptune": - from rsl_rl.utils.neptune_utils import NeptuneSummaryWriter - - self.writer = NeptuneSummaryWriter(log_dir=self.log_dir, flush_secs=10, cfg=self.cfg) - self.writer.log_config(self.env.cfg, self.cfg, self.alg_cfg, self.policy_cfg) - elif self.logger_type == "wandb": - from rsl_rl.utils.wandb_utils import WandbSummaryWriter - - self.writer = WandbSummaryWriter(log_dir=self.log_dir, flush_secs=10, cfg=self.cfg) - self.writer.log_config(self.env.cfg, self.cfg, self.alg_cfg, self.policy_cfg) - elif self.logger_type == "tensorboard": - from torch.utils.tensorboard import SummaryWriter - - self.writer = SummaryWriter(log_dir=self.log_dir, flush_secs=10) - else: - raise ValueError("Logger type not found. Please choose 'neptune', 'wandb' or 'tensorboard'.") diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/storage/__init__.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/storage/__init__.py deleted file mode 100644 index 1624330..0000000 --- a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/storage/__init__.py +++ /dev/null @@ -1,10 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -"""Implementation of transitions storage for RL-agent.""" - -from .rollout_storage import RolloutStorage -from .replay_buffer_multi import ReplayBufferMulti -__all__ = ["RolloutStorage", "ReplayBufferMulti"] diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/storage/replay_buffer_multi.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/storage/replay_buffer_multi.py deleted file mode 100644 index 6462b8e..0000000 --- a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/storage/replay_buffer_multi.py +++ /dev/null @@ -1,38 +0,0 @@ -import torch -import numpy as np - - -class ReplayBufferMulti: - """Fixed-size buffer to store experience tuples.""" - - def __init__(self, obs_dim, buffer_size, num_amp_frames, device): - """Initialize a ReplayBuffer object. - Arguments: - buffer_size (int): maximum size of buffer - """ - self.states = torch.zeros(buffer_size, num_amp_frames, obs_dim).to(device) - self.num_amp_frames = num_amp_frames - self.buffer_size = buffer_size - self.device = device - - self.step = 0 - self.num_samples = 0 - - def insert(self, states): - """Add new states to memory.""" - num_states = states.shape[0] - start_idx = self.step - end_idx = self.step + num_states - if end_idx > self.buffer_size: - self.states[self.step:self.buffer_size] = states[:self.buffer_size - self.step] - self.states[:end_idx - self.buffer_size] = states[self.buffer_size - self.step:] - else: - self.states[start_idx:end_idx] = states - - self.num_samples = min(self.buffer_size, max(end_idx, self.num_samples)) - self.step = (self.step + num_states) % self.buffer_size - - def feed_forward_generator(self, num_mini_batch, mini_batch_size): - for _ in range(num_mini_batch): - sample_idxs = np.random.choice(self.num_samples, size=mini_batch_size) - yield (self.states[sample_idxs].to(self.device)) diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/storage/rollout_storage.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/storage/rollout_storage.py deleted file mode 100644 index e9309b3..0000000 --- a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/storage/rollout_storage.py +++ /dev/null @@ -1,260 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -from __future__ import annotations - -import torch -from tensordict import TensorDict - -from rsl_rl.utils import split_and_pad_trajectories - - -class RolloutStorage: - class Transition: - def __init__(self): - self.observations = None - self.actions = None - self.privileged_actions = None - self.rewards = None - self.dones = None - self.values = None - self.actions_log_prob = None - self.action_mean = None - self.action_sigma = None - self.hidden_states = None - - def clear(self): - self.__init__() - - def __init__( - self, - training_type, - num_envs, - num_transitions_per_env, - obs, - actions_shape, - device="cpu", - ): - # store inputs - self.training_type = training_type - self.device = device - self.num_transitions_per_env = num_transitions_per_env - self.num_envs = num_envs - self.actions_shape = actions_shape - - # Core - self.observations = TensorDict( - {key: torch.zeros(num_transitions_per_env, *value.shape, device=device) for key, value in obs.items()}, - batch_size=[num_transitions_per_env, num_envs], - device=self.device, - ) - self.rewards = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device) - self.actions = torch.zeros(num_transitions_per_env, num_envs, *actions_shape, device=self.device) - self.dones = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device).byte() - - # for distillation - if training_type == "distillation": - self.privileged_actions = torch.zeros(num_transitions_per_env, num_envs, *actions_shape, device=self.device) - - # for reinforcement learning - if training_type == "rl": - self.values = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device) - self.actions_log_prob = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device) - self.mu = torch.zeros(num_transitions_per_env, num_envs, *actions_shape, device=self.device) - self.sigma = torch.zeros(num_transitions_per_env, num_envs, *actions_shape, device=self.device) - self.returns = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device) - self.advantages = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device) - - # For RNN networks - self.saved_hidden_states_a = None - self.saved_hidden_states_c = None - - # counter for the number of transitions stored - self.step = 0 - - def add_transitions(self, transition: Transition): - # check if the transition is valid - if self.step >= self.num_transitions_per_env: - raise OverflowError("Rollout buffer overflow! You should call clear() before adding new transitions.") - - # Core - self.observations[self.step].copy_(transition.observations) - self.actions[self.step].copy_(transition.actions) - self.rewards[self.step].copy_(transition.rewards.view(-1, 1)) - self.dones[self.step].copy_(transition.dones.view(-1, 1)) - - # for distillation - if self.training_type == "distillation": - self.privileged_actions[self.step].copy_(transition.privileged_actions) - - # for reinforcement learning - if self.training_type == "rl": - self.values[self.step].copy_(transition.values) - self.actions_log_prob[self.step].copy_(transition.actions_log_prob.view(-1, 1)) - self.mu[self.step].copy_(transition.action_mean) - self.sigma[self.step].copy_(transition.action_sigma) - - # For RNN networks - self._save_hidden_states(transition.hidden_states) - - # increment the counter - self.step += 1 - - def _save_hidden_states(self, hidden_states): - if hidden_states is None or hidden_states == (None, None): - return - # make a tuple out of GRU hidden state sto match the LSTM format - hid_a = hidden_states[0] if isinstance(hidden_states[0], tuple) else (hidden_states[0],) - hid_c = hidden_states[1] if isinstance(hidden_states[1], tuple) else (hidden_states[1],) - # initialize if needed - if self.saved_hidden_states_a is None: - self.saved_hidden_states_a = [ - torch.zeros(self.observations.shape[0], *hid_a[i].shape, device=self.device) for i in range(len(hid_a)) - ] - self.saved_hidden_states_c = [ - torch.zeros(self.observations.shape[0], *hid_c[i].shape, device=self.device) for i in range(len(hid_c)) - ] - # copy the states - for i in range(len(hid_a)): - self.saved_hidden_states_a[i][self.step].copy_(hid_a[i]) - self.saved_hidden_states_c[i][self.step].copy_(hid_c[i]) - - def clear(self): - self.step = 0 - - def compute_returns(self, last_values, gamma, lam, normalize_advantage: bool = True): - advantage = 0 - for step in reversed(range(self.num_transitions_per_env)): - # if we are at the last step, bootstrap the return value - if step == self.num_transitions_per_env - 1: - next_values = last_values - else: - next_values = self.values[step + 1] - # 1 if we are not in a terminal state, 0 otherwise - next_is_not_terminal = 1.0 - self.dones[step].float() - # TD error: r_t + gamma * V(s_{t+1}) - V(s_t) - delta = self.rewards[step] + next_is_not_terminal * gamma * next_values - self.values[step] - # Advantage: A(s_t, a_t) = delta_t + gamma * lambda * A(s_{t+1}, a_{t+1}) - advantage = delta + next_is_not_terminal * gamma * lam * advantage - # Return: R_t = A(s_t, a_t) + V(s_t) - self.returns[step] = advantage + self.values[step] - - # Compute the advantages - self.advantages = self.returns - self.values - # Normalize the advantages if flag is set - # This is to prevent double normalization (i.e. if per minibatch normalization is used) - if normalize_advantage: - self.advantages = (self.advantages - self.advantages.mean()) / (self.advantages.std() + 1e-8) - - # for distillation - def generator(self): - if self.training_type != "distillation": - raise ValueError("This function is only available for distillation training.") - - for i in range(self.num_transitions_per_env): - yield self.observations[i], self.actions[i], self.privileged_actions[i], self.dones[i] - - # for reinforcement learning with feedforward networks - def mini_batch_generator(self, num_mini_batches, num_epochs=8): - if self.training_type != "rl": - raise ValueError("This function is only available for reinforcement learning training.") - batch_size = self.num_envs * self.num_transitions_per_env - mini_batch_size = batch_size // num_mini_batches - indices = torch.randperm(num_mini_batches * mini_batch_size, requires_grad=False, device=self.device) - - # Core - observations = self.observations.flatten(0, 1) - actions = self.actions.flatten(0, 1) - values = self.values.flatten(0, 1) - returns = self.returns.flatten(0, 1) - - # For PPO - old_actions_log_prob = self.actions_log_prob.flatten(0, 1) - advantages = self.advantages.flatten(0, 1) - old_mu = self.mu.flatten(0, 1) - old_sigma = self.sigma.flatten(0, 1) - - for epoch in range(num_epochs): - for i in range(num_mini_batches): - # Select the indices for the mini-batch - start = i * mini_batch_size - end = (i + 1) * mini_batch_size - batch_idx = indices[start:end] - - # Create the mini-batch - # -- Core - obs_batch = observations[batch_idx] - actions_batch = actions[batch_idx] - - # -- For PPO - target_values_batch = values[batch_idx] - returns_batch = returns[batch_idx] - old_actions_log_prob_batch = old_actions_log_prob[batch_idx] - advantages_batch = advantages[batch_idx] - old_mu_batch = old_mu[batch_idx] - old_sigma_batch = old_sigma[batch_idx] - - # yield the mini-batch - yield obs_batch, actions_batch, target_values_batch, advantages_batch, returns_batch, old_actions_log_prob_batch, old_mu_batch, old_sigma_batch, ( - None, - None, - ), None - - # for reinfrocement learning with recurrent networks - def recurrent_mini_batch_generator(self, num_mini_batches, num_epochs=8): - if self.training_type != "rl": - raise ValueError("This function is only available for reinforcement learning training.") - padded_obs_trajectories, trajectory_masks = split_and_pad_trajectories(self.observations, self.dones) - - mini_batch_size = self.num_envs // num_mini_batches - for ep in range(num_epochs): - first_traj = 0 - for i in range(num_mini_batches): - start = i * mini_batch_size - stop = (i + 1) * mini_batch_size - - dones = self.dones.squeeze(-1) - last_was_done = torch.zeros_like(dones, dtype=torch.bool) - last_was_done[1:] = dones[:-1] - last_was_done[0] = True - trajectories_batch_size = torch.sum(last_was_done[:, start:stop]) - last_traj = first_traj + trajectories_batch_size - - masks_batch = trajectory_masks[:, first_traj:last_traj] - obs_batch = padded_obs_trajectories[:, first_traj:last_traj] - actions_batch = self.actions[:, start:stop] - old_mu_batch = self.mu[:, start:stop] - old_sigma_batch = self.sigma[:, start:stop] - returns_batch = self.returns[:, start:stop] - advantages_batch = self.advantages[:, start:stop] - values_batch = self.values[:, start:stop] - old_actions_log_prob_batch = self.actions_log_prob[:, start:stop] - - # reshape to [num_envs, time, num layers, hidden dim] (original shape: [time, num_layers, num_envs, hidden_dim]) - # then take only time steps after dones (flattens num envs and time dimensions), - # take a batch of trajectories and finally reshape back to [num_layers, batch, hidden_dim] - last_was_done = last_was_done.permute(1, 0) - hid_a_batch = [ - saved_hidden_states.permute(2, 0, 1, 3)[last_was_done][first_traj:last_traj] - .transpose(1, 0) - .contiguous() - for saved_hidden_states in self.saved_hidden_states_a - ] - hid_c_batch = [ - saved_hidden_states.permute(2, 0, 1, 3)[last_was_done][first_traj:last_traj] - .transpose(1, 0) - .contiguous() - for saved_hidden_states in self.saved_hidden_states_c - ] - # remove the tuple for GRU - hid_a_batch = hid_a_batch[0] if len(hid_a_batch) == 1 else hid_a_batch - hid_c_batch = hid_c_batch[0] if len(hid_c_batch) == 1 else hid_c_batch - - yield obs_batch, actions_batch, values_batch, advantages_batch, returns_batch, old_actions_log_prob_batch, old_mu_batch, old_sigma_batch, ( - hid_a_batch, - hid_c_batch, - ), masks_batch - - first_traj = last_traj diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/__init__.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/__init__.py deleted file mode 100644 index f5781f1..0000000 --- a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/__init__.py +++ /dev/null @@ -1,13 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -"""Helper functions.""" - -from .utils import * -from .motion_loader_g1 import G1_AMPLoader - -__all__ = [ - "G1_AMPLoader", -] \ No newline at end of file diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/motion_loader_g1.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/motion_loader_g1.py deleted file mode 100644 index 677c630..0000000 --- a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/motion_loader_g1.py +++ /dev/null @@ -1,388 +0,0 @@ -import os -from os.path import join as pjoin -import glob -import json -import logging - -import torch -import numpy as np -from pybullet_utils import transformations - -from rsl_rl.utils import motion_util - -_EPS = np.finfo(float).eps * 4.0 -def quaternion_slerp(q0, q1, fraction, spin=0, shortestpath=True): - """Batch quaternion spherical linear interpolation.""" - - out = torch.zeros_like(q0) - - zero_mask = torch.isclose(fraction, torch.zeros_like(fraction)).squeeze() - ones_mask = torch.isclose(fraction, torch.ones_like(fraction)).squeeze() - out[zero_mask] = q0[zero_mask] - out[ones_mask] = q1[ones_mask] - - d = torch.sum(q0 * q1, dim=-1, keepdim=True) - dist_mask = (torch.abs(torch.abs(d) - 1.0) < _EPS).squeeze() - out[dist_mask] = q0[dist_mask] - - if shortestpath: - d_old = torch.clone(d) - d = torch.where(d_old < 0, -d, d) - q1 = torch.where(d_old < 0, -q1, q1) - - angle = torch.acos(d) + spin * torch.pi - angle_mask = (torch.abs(angle) < _EPS).squeeze() - out[angle_mask] = q0[angle_mask] - - final_mask = torch.logical_or(zero_mask, ones_mask) - final_mask = torch.logical_or(final_mask, dist_mask) - final_mask = torch.logical_or(final_mask, angle_mask) - final_mask = torch.logical_not(final_mask) - - isin = 1.0 / angle - q0 *= torch.sin((1.0 - fraction) * angle) * isin - q1 *= torch.sin(fraction * angle) * isin - q0 += q1 - out[final_mask] = q0[final_mask] - return out - - -class G1_AMPLoader: - - def __init__( - self, - device, - time_between_frames, - motion_files, - preload_transitions=False, - num_preload_transitions=1000000, - num_frames=5, - ): - """Expert dataset provides AMP observations from Dog mocap dataset. - - time_between_frames: Amount of time in seconds between transition. - """ - self.device = device - self.time_between_frames = time_between_frames - self.num_frames = num_frames - - # Values to store for each trajectory. - self.trajectories = [] - self.trajectories_full = [] - self.trajectory_names = [] - self.trajectory_idxs = [] - self.trajectory_lens = [] # Traj length in seconds. - self.trajectory_weights = [] - self.trajectory_frame_durations = [] - self.trajectory_num_frames = [] - self.motion_dir = motion_files - # import ipdb; ipdb.set_trace() - for i, motion_file in enumerate(os.listdir(motion_files)): - self.trajectory_names.append(motion_file) - motion_path = pjoin(motion_files, motion_file) - motion_data = np.load(motion_path, allow_pickle=True) - motion_data_processed = np.zeros((motion_data.shape[0],36)) - - for f_i in range(motion_data.shape[0]): - motion_data_processed[f_i, :3] = motion_data[f_i, :3] # base pos - motion_data_processed[f_i, 3:7] = motion_data[f_i, 3:7] # base quat (wxyz) - motion_data_processed[f_i, 7:35] = motion_data[f_i, 7:35] # base vel - ''' - NOTE The order of motion_data_processed is - base pos 0:3, - base quat 3:7, wxyz - dof pos 7:36, (mujoco joint order) - ''' - self.trajectories.append(torch.tensor( - motion_data_processed[:, 7:], - dtype=torch.float32, - device=self.device - )) - - self.trajectories_full.append(torch.tensor( - motion_data_processed, - dtype=torch.float32, - device=self.device - )) - - self.trajectory_idxs.append(i) - self.trajectory_weights.append(1 / len(os.listdir(motion_files))) - frame_duration = 1 / 50 - - self.trajectory_frame_durations.append(frame_duration) - traj_len = (motion_data_processed.shape[0] - 1) * frame_duration # seconds - self.trajectory_lens.append(traj_len) - self.trajectory_num_frames.append(float(motion_data_processed.shape[0])) - print(f"Loaded {traj_len}s. motion from {motion_file}.") - - # Trajectory weights are used to sample some trajectories more than others. - self.trajectory_weights = np.array(self.trajectory_weights) / np.sum(self.trajectory_weights) - self.trajectory_frame_durations = np.array(self.trajectory_frame_durations) - self.trajectory_lens = np.array(self.trajectory_lens) - self.trajectory_num_frames = np.array(self.trajectory_num_frames) - - # Preload transitions. - self.preload_transitions = preload_transitions - if self.preload_transitions: - print(f'Preloading {num_preload_transitions} transitions') - - traj_idxs = self.weighted_traj_idx_sample_batch(num_preload_transitions) - times = self.traj_time_sample_batch(traj_idxs) - self.preloaded_s_prior = self.get_full_frame_at_time_batch(traj_idxs, times - self.time_between_frames) - self.preloaded_s = self.get_full_frame_at_time_batch(traj_idxs, times) - self.preloaded_s_next = self.get_full_frame_at_time_batch(traj_idxs, times + self.time_between_frames) - print(f'Finished preloading') - - # 预加载多帧数据 - self.preloaded_frames = [] - for i in range(self.num_frames): - frame_time = times + (i - (self.num_frames - 2)) * self.time_between_frames - full_frame = self.get_full_frame_at_time_batch(traj_idxs, frame_time) - # 预处理:提前提取并连接需要的列(7:26 和 29:33),避免每次生成时重复切片 - processed_frame = torch.cat([ - full_frame[:, 7:26], - full_frame[:, 29:33] - ], dim=-1) - self.preloaded_frames.append(processed_frame) - print(f'Finished preloading multiple frames') - - self.all_trajectories_full = torch.vstack(self.trajectories_full) - - def weighted_traj_idx_sample(self): - """Get traj idx via weighted sampling.""" - return np.random.choice( - self.trajectory_idxs, p=self.trajectory_weights) - - def weighted_traj_idx_sample_batch(self, size): - """Batch sample traj idxs.""" - return np.random.choice( - self.trajectory_idxs, size=size, p=self.trajectory_weights, - replace=True) - - def traj_time_sample(self, traj_idx): - """Sample random time for traj.""" - subst = self.time_between_frames + self.trajectory_frame_durations[traj_idx] - return max( - 0, (self.trajectory_lens[traj_idx] * np.random.uniform() - subst)) - - def traj_time_sample_batch(self, traj_idxs): - """Sample random time for multiple trajectories.""" - subst = self.time_between_frames + self.trajectory_frame_durations[traj_idxs] - time_samples = self.trajectory_lens[traj_idxs] * np.random.uniform(size=len(traj_idxs)) - subst - return np.maximum(np.zeros_like(time_samples), time_samples) - - def slerp(self, val0, val1, blend): - return (1.0 - blend) * val0 + blend * val1 - - def get_trajectory(self, traj_idx): - """Returns trajectory of AMP observations.""" - return self.trajectories_full[traj_idx] - - def get_frame_at_time(self, traj_idx, time): - """Returns frame for the given trajectory at the specified time.""" - p = float(time) / self.trajectory_lens[traj_idx] - n = self.trajectories[traj_idx].shape[0] - idx_low, idx_high = int(np.floor(p * n)), int(np.ceil(p * n)) - frame_start = self.trajectories[traj_idx][idx_low] - frame_end = self.trajectories[traj_idx][idx_high] - blend = p * n - idx_low - return self.slerp(frame_start, frame_end, blend) - - def get_frame_at_time_batch(self, traj_idxs, times): - """Returns frame for the given trajectory at the specified time.""" - p = times / self.trajectory_lens[traj_idxs] - n = self.trajectory_num_frames[traj_idxs] - idx_low, idx_high = np.floor(p * n).astype(np.int32), np.ceil(p * n).astype(np.int32) - all_frame_starts = torch.zeros(len(traj_idxs), self.observation_dim, device=self.device) - all_frame_ends = torch.zeros(len(traj_idxs), self.observation_dim, device=self.device) - for traj_idx in set(traj_idxs): - trajectory = self.trajectories[traj_idx] - traj_mask = traj_idxs == traj_idx - all_frame_starts[traj_mask] = trajectory[idx_low[traj_mask]] - all_frame_ends[traj_mask] = trajectory[idx_high[traj_mask]] - blend = torch.tensor(p * n - idx_low, device=self.device, dtype=torch.float32).unsqueeze(-1) - return self.slerp(all_frame_starts, all_frame_ends, blend) - - def get_full_frame_at_time(self, traj_idx, time): - """Returns full frame for the given trajectory at the specified time.""" - p = float(time) / self.trajectory_lens[traj_idx] - n = self.trajectories_full[traj_idx].shape[0] - idx_low, idx_high = int(np.floor(p * n)), int(np.ceil(p * n)) - frame_start = self.trajectories_full[traj_idx][idx_low] - frame_end = self.trajectories_full[traj_idx][idx_high] - blend = p * n - idx_low - print(idx_low, idx_high) - return self.blend_frame_pose(frame_start, frame_end, blend) - - def get_full_frame_at_time_batch(self, traj_idxs, times): - p = times / self.trajectory_lens[traj_idxs] - n = self.trajectory_num_frames[traj_idxs] - idx_low, idx_high = np.floor(p * n).astype(np.int32), np.ceil(p * n).astype(np.int32) - all_frame_pos_starts = torch.zeros(len(traj_idxs), 3, device=self.device) - all_frame_pos_ends = torch.zeros(len(traj_idxs), 3, device=self.device) - all_frame_rot_starts = torch.zeros(len(traj_idxs), 4, device=self.device) - all_frame_rot_ends = torch.zeros(len(traj_idxs), 4, device=self.device) - all_frame_amp_starts = torch.zeros(len(traj_idxs), 29, device=self.device) - all_frame_amp_ends = torch.zeros(len(traj_idxs), 29, device=self.device) - for traj_idx in set(traj_idxs): - trajectory = self.trajectories_full[traj_idx] - traj_mask = traj_idxs == traj_idx - all_frame_pos_starts[traj_mask] = G1_AMPLoader.get_root_pos_batch(trajectory[idx_low[traj_mask]]) - all_frame_pos_ends[traj_mask] = G1_AMPLoader.get_root_pos_batch(trajectory[idx_high[traj_mask]]) - all_frame_rot_starts[traj_mask] = G1_AMPLoader.get_root_rot_batch(trajectory[idx_low[traj_mask]]) - all_frame_rot_ends[traj_mask] = G1_AMPLoader.get_root_rot_batch(trajectory[idx_high[traj_mask]]) - all_frame_amp_starts[traj_mask] = trajectory[idx_low[traj_mask]][:, 7:36] # base vel3+ang3, dof vel23+ang23 - all_frame_amp_ends[traj_mask] = trajectory[idx_high[traj_mask]][:, 7:36] # base vel3+ang3, dof vel23+ang23 - blend = torch.tensor(p * n - idx_low, device=self.device, dtype=torch.float32).unsqueeze(-1) - pos_blend = self.slerp(all_frame_pos_starts, all_frame_pos_ends, blend) - rot_blend = quaternion_slerp(all_frame_rot_starts, all_frame_rot_ends, blend) - amp_blend = self.slerp(all_frame_amp_starts, all_frame_amp_ends, blend) - return torch.cat([pos_blend, rot_blend, amp_blend], dim=-1) - - def get_frame(self): - """Returns random frame.""" - traj_idx = self.weighted_traj_idx_sample() - sampled_time = self.traj_time_sample(traj_idx) - return self.get_frame_at_time(traj_idx, sampled_time) - - def get_full_frame(self): - """Returns random full frame.""" - traj_idx = self.weighted_traj_idx_sample() - sampled_time = self.traj_time_sample(traj_idx) - return self.get_full_frame_at_time(traj_idx, sampled_time) - - def get_full_frame_batch(self, num_frames): - if self.preload_transitions: - idxs = np.random.choice( - self.preloaded_s.shape[0], size=num_frames) - return self.preloaded_s[idxs] - else: - traj_idxs = self.weighted_traj_idx_sample_batch(num_frames) - times = self.traj_time_sample_batch(traj_idxs) - return self.get_full_frame_at_time_batch(traj_idxs, times) - - def blend_frame_pose(self, frame0, frame1, blend): - """Linearly interpolate between two frames, including orientation. - - Args: - frame0: First frame to be blended corresponds to (blend = 0). - frame1: Second frame to be blended corresponds to (blend = 1). - blend: Float between [0, 1], specifying the interpolation between - the two frames. - Returns: - An interpolation of the two frames. - """ - root_pos0, root_pos1 = G1_AMPLoader.get_root_pos(frame0), G1_AMPLoader.get_root_pos(frame1) - root_rot0, root_rot1 = G1_AMPLoader.get_root_rot(frame0), G1_AMPLoader.get_root_rot(frame1) - joints0, joints1 = G1_AMPLoader.get_joint_pose(frame0), G1_AMPLoader.get_joint_pose(frame1) - # tar_toe_pos_0, tar_toe_pos_1 = G1_AMPLoader.get_tar_toe_pos_local(frame0), G1_AMPLoader.get_tar_toe_pos_local(frame1) - linear_vel_0, linear_vel_1 = G1_AMPLoader.get_linear_vel(frame0), G1_AMPLoader.get_linear_vel(frame1) - angular_vel_0, angular_vel_1 = G1_AMPLoader.get_angular_vel(frame0), G1_AMPLoader.get_angular_vel(frame1) - joint_vel_0, joint_vel_1 = G1_AMPLoader.get_joint_vel(frame0), G1_AMPLoader.get_joint_vel(frame1) - - blend_root_pos = self.slerp(root_pos0, root_pos1, blend) - blend_root_rot = transformations.quaternion_slerp(root_rot0.cpu().numpy(), root_rot1.cpu().numpy(), blend) - blend_root_rot = torch.tensor(motion_util.standardize_quaternion(blend_root_rot),dtype=torch.float32, device=self.device) - blend_joints = self.slerp(joints0, joints1, blend) - # blend_tar_toe_pos = self.slerp(tar_toe_pos_0, tar_toe_pos_1, blend) - blend_linear_vel = self.slerp(linear_vel_0, linear_vel_1, blend) - blend_angular_vel = self.slerp(angular_vel_0, angular_vel_1, blend) - blend_joints_vel = self.slerp(joint_vel_0, joint_vel_1, blend) - - # return - # torch.cat([ - # blend_root_pos, blend_root_rot, blend_linear_vel, blend_angular_vel, blend_joints, blend_joints_vel]) - return torch.cat([blend_root_pos, blend_root_rot, blend_linear_vel, blend_angular_vel, blend_joints]) - - def feed_forward_generator_23dof_multi(self, num_mini_batch, mini_batch_size): - """Generates a batch of AMP transitions.""" - # import ipdb; ipdb.set_trace() - for _ in range(num_mini_batch): - if self.preload_transitions: - idxs = np.random.choice(self.preloaded_s.shape[0], size=mini_batch_size) - - frames = [] - for i in range(self.num_frames): - # 数据已在预加载时预处理,直接索引即可 - s = self.preloaded_frames[i][idxs] - frames.append(s) - else: - NotImplementedError('preload transition') - yield torch.stack(frames, dim=1) # [batch, num_frames, 16] - - - - - def quaternion_to_euler_array(self, quat): - # Ensure quaternion is in the correct format [x, y, z, w] - x, y, z, w =quat - - # Roll (x-axis rotation) - t0 = +2.0 * (w * x + y * z) - t1 = +1.0 - 2.0 * (x * x + y * y) - roll_x = np.arctan2(t0, t1) - - # Pitch (y-axis rotation) - t2 = +2.0 * (w * y - z * x) - t2 = np.clip(t2, -1.0, 1.0) - pitch_y = np.arcsin(t2) - - # Yaw (z-axis rotation) - t3 = +2.0 * (w * z + x * y) - t4 = +1.0 - 2.0 * (y * y + z * z) - yaw_z = np.arctan2(t3, t4) - - # Returns roll, pitch, yaw in a NumPy array in radians - return np.array([roll_x, pitch_y, yaw_z]) - - def euler_to_quaternion(self, root_rot): - roll, pitch, yaw = root_rot[0], root_rot[1], root_rot[2] - cy = np.cos(yaw * 0.5) - sy = np.sin(yaw * 0.5) - cp = np.cos(pitch * 0.5) - sp = np.sin(pitch * 0.5) - cr = np.cos(roll * 0.5) - sr = np.sin(roll * 0.5) - - qw = cy * cp * cr + sy * sp * sr - qx = cy * cp * sr - sy * sp * cr - qy = sy * cp * sr + cy * sp * cr - qz = sy * cp * cr - cy * sp * sr - - return np.array([qx, qy, qz, qw]) - - @property - def observation_dim(self): - """Size of AMP observations.""" - return self.trajectories[0].shape[1] + 1 - - @property - def num_motions(self): - return len(self.trajectory_names) - @staticmethod - def get_root_pos(pose): - return pose[0:3] - - @staticmethod - def get_root_pos_batch(poses): - return poses[:, 0:3] - - @staticmethod - def get_root_rot(pose): - return pose[3:7] - - @staticmethod - def get_root_rot_batch(poses): - return poses[:, 3:7] - - @staticmethod - def get_joint_pose_batch_12dof(poses): - return poses[:, 13:25] - - @staticmethod - def get_tar_toe_pos_local(pose): - return pose[G1_AMPLoader.TAR_TOE_POS_LOCAL_START_IDX:G1_AMPLoader.TAR_TOE_POS_LOCAL_END_IDX] - - @staticmethod - def get_tar_toe_pos_local_batch(poses): - return poses[:, G1_AMPLoader.TAR_TOE_POS_LOCAL_START_IDX:G1_AMPLoader.TAR_TOE_POS_LOCAL_END_IDX] diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/motion_util.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/motion_util.py deleted file mode 100644 index 3d49bfa..0000000 --- a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/motion_util.py +++ /dev/null @@ -1,97 +0,0 @@ -# coding=utf-8 -# Copyright 2020 The Google Research Authors. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Utility functions for processing motion clips.""" - -import os -import inspect -currentdir = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe()))) -parentdir = os.path.dirname(os.path.dirname(currentdir)) -os.sys.path.insert(0, parentdir) - -import numpy as np - -from rsl_rl.utils import pose3d -# from pybullet_utils import transformations - - -def standardize_quaternion(q): - """Returns a quaternion where q.w >= 0 to remove redundancy due to q = -q. - - Args: - q: A quaternion to be standardized. - - Returns: - A quaternion with q.w >= 0. - - """ - if q[-1] < 0: - q = -q - return q - - -def normalize_rotation_angle(theta): - """Returns a rotation angle normalized between [-pi, pi]. - - Args: - theta: angle of rotation (radians). - - Returns: - An angle of rotation normalized between [-pi, pi]. - - """ - norm_theta = theta - if np.abs(norm_theta) > np.pi: - norm_theta = np.fmod(norm_theta, 2 * np.pi) - if norm_theta >= 0: - norm_theta += -2 * np.pi - else: - norm_theta += 2 * np.pi - - return norm_theta - - -def calc_heading(q): - """Returns the heading of a rotation q, specified as a quaternion. - - The heading represents the rotational component of q along the vertical - axis (z axis). - - Args: - q: A quaternion that the heading is to be computed from. - - Returns: - An angle representing the rotation about the z axis. - - """ - ref_dir = np.array([1, 0, 0]) - rot_dir = pose3d.QuaternionRotatePoint(ref_dir, q) - heading = np.arctan2(rot_dir[1], rot_dir[0]) - return heading - - -# def calc_heading_rot(q): -# """Return a quaternion representing the heading rotation of q along the vertical axis (z axis). - -# Args: -# q: A quaternion that the heading is to be computed from. - -# Returns: -# A quaternion representing the rotation about the z axis. - -# """ -# heading = calc_heading(q) -# q_heading = transformations.quaternion_about_axis(heading, [0, 0, 1]) -# return q_heading diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/neptune_utils.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/neptune_utils.py deleted file mode 100644 index 3796ec8..0000000 --- a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/neptune_utils.py +++ /dev/null @@ -1,94 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -from __future__ import annotations - -import os -from dataclasses import asdict -from torch.utils.tensorboard import SummaryWriter - -try: - import neptune -except ModuleNotFoundError: - raise ModuleNotFoundError("neptune-client is required to log to Neptune.") - - -class NeptuneLogger: - def __init__(self, project, token): - self.run = neptune.init_run(project=project, api_token=token) - - def store_config(self, env_cfg, runner_cfg, alg_cfg, policy_cfg): - self.run["runner_cfg"] = runner_cfg - self.run["policy_cfg"] = policy_cfg - self.run["alg_cfg"] = alg_cfg - self.run["env_cfg"] = asdict(env_cfg) - - -class NeptuneSummaryWriter(SummaryWriter): - """Summary writer for Neptune.""" - - def __init__(self, log_dir: str, flush_secs: int, cfg): - super().__init__(log_dir, flush_secs) - - try: - project = cfg["neptune_project"] - except KeyError: - raise KeyError("Please specify neptune_project in the runner config, e.g. legged_gym.") - - try: - token = os.environ["NEPTUNE_API_TOKEN"] - except KeyError: - raise KeyError( - "Neptune api token not found. Please run or add to ~/.bashrc: export NEPTUNE_API_TOKEN=YOUR_API_TOKEN" - ) - - try: - entity = os.environ["NEPTUNE_USERNAME"] - except KeyError: - raise KeyError( - "Neptune username not found. Please run or add to ~/.bashrc: export NEPTUNE_USERNAME=YOUR_USERNAME" - ) - - neptune_project = entity + "/" + project - - self.neptune_logger = NeptuneLogger(neptune_project, token) - - self.name_map = { - "Train/mean_reward/time": "Train/mean_reward_time", - "Train/mean_episode_length/time": "Train/mean_episode_length_time", - } - - run_name = os.path.split(log_dir)[-1] - - self.neptune_logger.run["log_dir"].log(run_name) - - def _map_path(self, path): - if path in self.name_map: - return self.name_map[path] - else: - return path - - def add_scalar(self, tag, scalar_value, global_step=None, walltime=None, new_style=False): - super().add_scalar( - tag, - scalar_value, - global_step=global_step, - walltime=walltime, - new_style=new_style, - ) - self.neptune_logger.run[self._map_path(tag)].log(scalar_value, step=global_step) - - def stop(self): - self.neptune_logger.run.stop() - - def log_config(self, env_cfg, runner_cfg, alg_cfg, policy_cfg): - self.neptune_logger.store_config(env_cfg, runner_cfg, alg_cfg, policy_cfg) - - def save_model(self, model_path, iter): - self.neptune_logger.run["model/saved_model_" + str(iter)].upload(model_path) - - def save_file(self, path, iter=None): - name = path.rsplit("/", 1)[-1].split(".")[0] - self.neptune_logger.run["git_diff/" + name].upload(path) diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/pose3d.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/pose3d.py deleted file mode 100644 index a4cba1e..0000000 --- a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/pose3d.py +++ /dev/null @@ -1,283 +0,0 @@ -# coding=utf-8 -# Copyright 2020 The Google Research Authors. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Utilities for 3D pose conversion.""" -import math -import numpy as np - -# from pybullet_utils import transformations - -VECTOR3_0 = np.zeros(3, dtype=np.float64) -VECTOR3_1 = np.ones(3, dtype=np.float64) -VECTOR3_X = np.array([1, 0, 0], dtype=np.float64) -VECTOR3_Y = np.array([0, 1, 0], dtype=np.float64) -VECTOR3_Z = np.array([0, 0, 1], dtype=np.float64) - -# QUATERNION_IDENTITY is the multiplicative identity 1.0 + 0i + 0j + 0k. -# When interpreted as a rotation, it is the identity rotation. -QUATERNION_IDENTITY = np.array([0.0, 0.0, 0.0, 1.0], dtype=np.float64) - - -def Vector3RandomNormal(sigma, mu=VECTOR3_0): - """Returns a random 3D vector from a normal distribution. - - Each component is selected independently from a normal distribution. - - Args: - sigma: Scale (or stddev) of distribution for all variables. - mu: Mean of distribution for each variable. - - Returns: - A 3D vector in a numpy array. - """ - - random_v3 = np.random.normal(scale=sigma, size=3) + mu - return random_v3 - - -def Vector3RandomUniform(low=VECTOR3_0, high=VECTOR3_1): - """Returns a 3D vector selected uniformly from the input box. - - Args: - low: The min-value corner of the box. - high: The max-value corner of the box. - - Returns: - A 3D vector in a numpy array. - """ - - random_x = np.random.uniform(low=low[0], high=high[0]) - random_y = np.random.uniform(low=low[1], high=high[1]) - random_z = np.random.uniform(low=low[2], high=high[2]) - return np.array([random_x, random_y, random_z]) - - -def Vector3RandomUnit(): - """Returns a random 3D vector with unit length. - - Generates a 3D vector selected uniformly from the unit sphere. - - Returns: - A normalized 3D vector in a numpy array. - """ - longitude = np.random.uniform(low=-math.pi, high=math.pi) - sin_latitude = np.random.uniform(low=-1.0, high=1.0) - cos_latitude = math.sqrt(1.0 - sin_latitude * sin_latitude) - x = math.cos(longitude) * cos_latitude - y = math.sin(longitude) * cos_latitude - z = sin_latitude - return np.array([x, y, z], dtype=np.float64) - - -def QuaternionNormalize(q): - """Normalizes the quaternion to length 1. - - Divides the quaternion by its magnitude. If the magnitude is too - small, returns the quaternion identity value (1.0). - - Args: - q: A quaternion to be normalized. - - Raises: - ValueError: If input quaternion has length near zero. - - Returns: - A quaternion with magnitude 1 in a numpy array [x, y, z, w]. - - """ - q_norm = np.linalg.norm(q) - if np.isclose(q_norm, 0.0): - raise ValueError( - 'Quaternion may not be zero in QuaternionNormalize: |q| = %f, q = %s' % - (q_norm, q)) - return q / q_norm - - -def QuaternionFromAxisAngle(axis, angle): - """Returns a quaternion that generates the given axis-angle rotation. - - Returns the quaternion: sin(angle/2) * axis + cos(angle/2). - - Args: - axis: Axis of rotation, a 3D vector in a numpy array. - angle: The angle of rotation (radians). - - Raises: - ValueError: If input axis is not a normalizable 3D vector. - - Returns: - A unit quaternion in a numpy array. - - """ - if len(axis) != 3: - raise ValueError('Axis vector should have three components: %s' % axis) - axis_norm = np.linalg.norm(axis) - if np.isclose(axis_norm, 0.0): - raise ValueError('Axis vector may not have zero length: |v| = %f, v = %s' % - (axis_norm, axis)) - half_angle = angle * 0.5 - q = np.zeros(4, dtype=np.float64) - q[0:3] = axis - q[0:3] *= math.sin(half_angle) / axis_norm - q[3] = math.cos(half_angle) - return q - - -def QuaternionToAxisAngle(quat, default_axis=VECTOR3_Z, direction_axis=None): - """Calculates axis and angle of rotation performed by a quaternion. - - Calculates the axis and angle of the rotation performed by the quaternion. - The quaternion should have four values and be normalized. - - Args: - quat: Unit quaternion in a numpy array. - default_axis: 3D vector axis used if the rotation is near to zero. Without - this default, small rotations would result in an exception. It is - reasonable to use a default axis for tiny rotations, because zero angle - rotations about any axis are equivalent. - direction_axis: Used to disambiguate rotation directions. If the - direction_axis is specified, the axis of the rotation will be chosen such - that its inner product with the direction_axis is non-negative. - - Raises: - ValueError: If quat is not a normalized quaternion. - - Returns: - axis: Axis of rotation. - angle: Angle in radians. - """ - if len(quat) != 4: - raise ValueError( - 'Quaternion should have four components [x, y, z, w]: %s' % quat) - if not np.isclose(1.0, np.linalg.norm(quat)): - raise ValueError('Quaternion should have unit length: |q| = %f, q = %s' % - (np.linalg.norm(quat), quat)) - axis = quat[:3].copy() - axis_norm = np.linalg.norm(axis) - min_axis_norm = 1e-8 - if axis_norm < min_axis_norm: - axis = default_axis - if len(default_axis) != 3: - raise ValueError('Axis vector should have three components: %s' % axis) - if not np.isclose(np.linalg.norm(axis), 1.0): - raise ValueError('Axis vector should have unit length: |v| = %f, v = %s' % - (np.linalg.norm(axis), axis)) - else: - axis /= axis_norm - sin_half_angle = axis_norm - if direction_axis is not None and np.inner(axis, direction_axis) < 0: - sin_half_angle = -sin_half_angle - axis = -axis - cos_half_angle = quat[3] - half_angle = math.atan2(sin_half_angle, cos_half_angle) - angle = half_angle * 2 - return axis, angle - - -def QuaternionRandomRotation(max_angle=math.pi): - """Creates a random small rotation around a random axis. - - Generates a small rotation with the axis vector selected uniformly - from the unit sphere and an angle selected from a uniform - distribution over [0, max_angle]. - - If the max_angle is not specified, the rotation should be selected - uniformly over all possible rotation angles. - - Args: - max_angle: The maximum angle of rotation (radians). - - Returns: - A unit quaternion in a numpy array. - - """ - - angle = np.random.uniform(low=0, high=max_angle) - axis = Vector3RandomUnit() - return QuaternionFromAxisAngle(axis, angle) - - -# def QuaternionRotatePoint(point, quat): -# """Performs a rotation by quaternion. - -# Rotate the point by the quaternion using quaternion multiplication, -# (q * p * q^-1), without constructing the rotation matrix. - -# Args: -# point: The point to be rotated. -# quat: The rotation represented as a quaternion [x, y, z, w]. - -# Returns: -# A 3D vector in a numpy array. -# """ - -# q_point = np.array([point[0], point[1], point[2], 0.0]) -# quat_inverse = transformations.quaternion_inverse(quat) -# q_point_rotated = transformations.quaternion_multiply( -# transformations.quaternion_multiply(quat, q_point), quat_inverse) -# return q_point_rotated[:3] - - -def IsRotationMatrix(m): - """Returns true if the 3x3 submatrix represents a rotation. - - Args: - m: A transformation matrix. - - Raises: - ValueError: If input is not a matrix of size at least 3x3. - - Returns: - True if the 3x3 submatrix is a rotation (orthogonal). - """ - if len(m.shape) != 2 or m.shape[0] < 3 or m.shape[1] < 3: - raise ValueError('Matrix should be 3x3 or 4x4: %s\n %s' % (m.shape, m)) - rot = m[:3, :3] - eye = np.matmul(rot, np.transpose(rot)) - return np.isclose(eye, np.identity(3), atol=1e-4).all() - -# def ZAxisAlignedRobotPoseTool(robot_pose_tool): -# """Returns the current gripper pose rotated for alignment with the z-axis. - -# Args: -# robot_pose_tool: a pose3d.Pose3d() instance. - -# Returns: -# An instance of pose.Transform representing the current gripper pose -# rotated for alignment with the z-axis. -# """ -# # Align the current pose to the z-axis. -# robot_pose_tool.quaternion = transformations.quaternion_multiply( -# RotationBetween( -# robot_pose_tool.matrix4x4[0:3, 0:3].dot(np.array([0, 0, 1])), -# np.array([0.0, 0.0, -1.0])), robot_pose_tool.quaternion) -# return robot_pose_tool - -# def RotationBetween(a_translation_b, a_translation_c): -# """Computes the rotation from one vector to another. - -# The computed rotation has the property that: - -# a_translation_c = a_rotation_b_to_c * a_translation_b - -# Args: -# a_translation_b: vec3, vector to rotate from -# a_translation_c: vec3, vector to rotate to - -# Returns: -# a_rotation_b_to_c: new Orientation -# """ -# rotation = rotation3.Rotation3.rotation_between( -# a_translation_b, a_translation_c, err_msg='RotationBetween') -# return rotation.quaternion.xyzw diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/utils.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/utils.py deleted file mode 100644 index 3605622..0000000 --- a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/utils.py +++ /dev/null @@ -1,360 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -from __future__ import annotations - -import git -import importlib -import os -import pathlib -import torch -import warnings -from tensordict import TensorDict -from typing import Callable -import numpy as np -class RunningMeanStd: - def __init__(self, epsilon: float = 1e-4, shape: Tuple[int, ...] = ()): - """ - Calculates the running mean and std of a data stream - https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Parallel_algorithm - :param epsilon: helps with arithmetic issues - :param shape: the shape of the data stream's output - """ - self.mean = np.zeros(shape, np.float64) - self.var = np.ones(shape, np.float64) - self.count = epsilon - - def update(self, arr: np.ndarray) -> None: - batch_mean = np.mean(arr, axis=0) - batch_var = np.var(arr, axis=0) - batch_count = arr.shape[0] - self.update_from_moments(batch_mean, batch_var, batch_count) - - def update_from_moments(self, batch_mean: np.ndarray, batch_var: np.ndarray, batch_count: int) -> None: - delta = batch_mean - self.mean - tot_count = self.count + batch_count - - new_mean = self.mean + delta * batch_count / tot_count - m_a = self.var * self.count - m_b = batch_var * batch_count - m_2 = m_a + m_b + np.square(delta) * self.count * batch_count / (self.count + batch_count) - new_var = m_2 / (self.count + batch_count) - - new_count = batch_count + self.count - - self.mean = new_mean - self.var = new_var - self.count = new_count - - -class Normalizer(RunningMeanStd): - def __init__(self, input_dim, epsilon=1e-4, clip_obs=10.0): - super().__init__(shape=input_dim) - self.epsilon = epsilon - self.clip_obs = clip_obs - - def normalize(self, input): - return np.clip((input - self.mean) / np.sqrt(self.var + self.epsilon), -self.clip_obs, self.clip_obs) - - def normalize_torch(self, input, device): - mean_torch = torch.tensor(self.mean, device=device, dtype=torch.float32) - std_torch = torch.sqrt(torch.tensor(self.var + self.epsilon, device=device, dtype=torch.float32)) - return torch.clamp((input - mean_torch) / std_torch, -self.clip_obs, self.clip_obs) - - def update_normalizer(self, rollouts, expert_loader): - policy_data_generator = rollouts.feed_forward_generator_amp(None, mini_batch_size=expert_loader.batch_size) - expert_data_generator = expert_loader.dataset.feed_forward_generator_amp(expert_loader.batch_size) - - for expert_batch, policy_batch in zip(expert_data_generator, policy_data_generator): - self.update(torch.vstack(tuple(policy_batch) + tuple(expert_batch)).cpu().numpy()) - - -def resolve_nn_activation(act_name: str) -> torch.nn.Module: - """Resolves the activation function from the name. - - Args: - act_name: The name of the activation function. - - Returns: - The activation function. - - Raises: - ValueError: If the activation function is not found. - """ - act_dict = { - "elu": torch.nn.ELU(), - "selu": torch.nn.SELU(), - "relu": torch.nn.ReLU(), - "crelu": torch.nn.CELU(), - "lrelu": torch.nn.LeakyReLU(), - "tanh": torch.nn.Tanh(), - "sigmoid": torch.nn.Sigmoid(), - "softplus": torch.nn.Softplus(), - "gelu": torch.nn.GELU(), - "swish": torch.nn.SiLU(), - "mish": torch.nn.Mish(), - "identity": torch.nn.Identity(), - } - - act_name = act_name.lower() - if act_name in act_dict: - return act_dict[act_name] - else: - raise ValueError(f"Invalid activation function '{act_name}'. Valid activations are: {list(act_dict.keys())}") - - -def resolve_optimizer(optimizer_name: str) -> torch.optim.Optimizer: - """Resolves the optimizer from the name. - - Args: - optimizer_name: The name of the optimizer. - - Returns: - The optimizer. - - Raises: - ValueError: If the optimizer is not found. - """ - optimizer_dict = { - "adam": torch.optim.Adam, - "adamw": torch.optim.AdamW, - "sgd": torch.optim.SGD, - "rmsprop": torch.optim.RMSprop, - } - - optimizer_name = optimizer_name.lower() - if optimizer_name in optimizer_dict: - return optimizer_dict[optimizer_name] - else: - raise ValueError(f"Invalid optimizer '{optimizer_name}'. Valid optimizers are: {list(optimizer_dict.keys())}") - - -def split_and_pad_trajectories( - tensor: torch.Tensor | TensorDict, dones: torch.Tensor -) -> tuple[torch.Tensor | TensorDict, torch.Tensor]: - """Splits trajectories at done indices. Then concatenates them and pads with zeros up to the length of the longest - trajectory. Returns masks corresponding to valid parts of the trajectories. - - Example: - Input: [[a1, a2, a3, a4 | a5, a6], - [b1, b2 | b3, b4, b5 | b6]] - - Output:[[a1, a2, a3, a4], | [[True, True, True, True], - [a5, a6, 0, 0], | [True, True, False, False], - [b1, b2, 0, 0], | [True, True, False, False], - [b3, b4, b5, 0], | [True, True, True, False], - [b6, 0, 0, 0]] | [True, False, False, False]] - - Assumes that the input has the following order of dimensions: [time, number of envs, additional dimensions] - """ - - dones = dones.clone() - dones[-1] = 1 - # Permute the buffers to have order (num_envs, num_transitions_per_env, ...), for correct reshaping - flat_dones = dones.transpose(1, 0).reshape(-1, 1) - # Get length of trajectory by counting the number of successive not done elements - done_indices = torch.cat((flat_dones.new_tensor([-1], dtype=torch.int64), flat_dones.nonzero()[:, 0])) - trajectory_lengths = done_indices[1:] - done_indices[:-1] - trajectory_lengths_list = trajectory_lengths.tolist() - # Extract the individual trajectories - if isinstance(tensor, TensorDict): - padded_trajectories = {} - for k, v in tensor.items(): - # split the tensor into trajectories - trajectories = torch.split(v.transpose(1, 0).flatten(0, 1), trajectory_lengths_list) - # add at least one full length trajectory - trajectories = trajectories + (torch.zeros(v.shape[0], *v.shape[2:], device=v.device),) - # pad the trajectories to the length of the longest trajectory - padded_trajectories[k] = torch.nn.utils.rnn.pad_sequence(trajectories) - # remove the added tensor - padded_trajectories[k] = padded_trajectories[k][:, :-1] - padded_trajectories = TensorDict( - padded_trajectories, batch_size=[tensor.batch_size[0], len(trajectory_lengths_list)] - ) - else: - # split the tensor into trajectories - trajectories = torch.split(tensor.transpose(1, 0).flatten(0, 1), trajectory_lengths_list) - # add at least one full length trajectory - trajectories = trajectories + (torch.zeros(tensor.shape[0], *tensor.shape[2:], device=tensor.device),) - # pad the trajectories to the length of the longest trajectory - padded_trajectories = torch.nn.utils.rnn.pad_sequence(trajectories) - # remove the added tensor - padded_trajectories = padded_trajectories[:, :-1] - # create masks for the valid parts of the trajectories - trajectory_masks = trajectory_lengths > torch.arange(0, tensor.shape[0], device=tensor.device).unsqueeze(1) - return padded_trajectories, trajectory_masks - - -def unpad_trajectories(trajectories, masks): - """Does the inverse operation of split_and_pad_trajectories()""" - # Need to transpose before and after the masking to have proper reshaping - return ( - trajectories.transpose(1, 0)[masks.transpose(1, 0)] - .view(-1, trajectories.shape[0], trajectories.shape[-1]) - .transpose(1, 0) - ) - - -def store_code_state(logdir, repositories) -> list: - git_log_dir = os.path.join(logdir, "git") - os.makedirs(git_log_dir, exist_ok=True) - file_paths = [] - for repository_file_path in repositories: - try: - repo = git.Repo(repository_file_path, search_parent_directories=True) - t = repo.head.commit.tree - except Exception: - print(f"Could not find git repository in {repository_file_path}. Skipping.") - # skip if not a git repository - continue - # get the name of the repository - repo_name = pathlib.Path(repo.working_dir).name - diff_file_name = os.path.join(git_log_dir, f"{repo_name}.diff") - # check if the diff file already exists - if os.path.isfile(diff_file_name): - continue - # write the diff file - print(f"Storing git diff for '{repo_name}' in: {diff_file_name}") - with open(diff_file_name, "x", encoding="utf-8") as f: - content = f"--- git status ---\n{repo.git.status()} \n\n\n--- git diff ---\n{repo.git.diff(t)}" - f.write(content) - # add the file path to the list of files to be uploaded - file_paths.append(diff_file_name) - return file_paths - - -def string_to_callable(name: str) -> Callable: - """Resolves the module and function names to return the function. - - Args: - name: The function name. The format should be 'module:attribute_name'. - - Raises: - ValueError: When the resolved attribute is not a function. - ValueError: When unable to resolve the attribute. - - Returns: - The function loaded from the module. - """ - try: - mod_name, attr_name = name.split(":") - mod = importlib.import_module(mod_name) - callable_object = getattr(mod, attr_name) - # check if attribute is callable - if callable(callable_object): - return callable_object - else: - raise ValueError(f"The imported object is not callable: '{name}'") - except AttributeError as e: - msg = ( - "We could not interpret the entry as a callable object. The format of input should be" - f" 'module:attribute_name'\nWhile processing input '{name}', received the error:\n {e}." - ) - raise ValueError(msg) - - -def resolve_obs_groups( - obs: TensorDict, obs_groups: dict[str, list[str]], default_sets: list[str] -) -> dict[str, list[str]]: - """Validates the observation configuration and defaults missing observation sets. - - The input is an observation dictionary `obs` containing observation groups and a configuration dictionary - `obs_groups` where the keys are the observation sets and the values are lists of observation groups. - - The configuration dictionary could for example look like: - { - "policy": ["group_1", "group_2"], - "critic": ["group_1", "group_3"] - } - - This means that the 'policy' observation set will contain the observations "group_1" and "group_2" and the - 'critic' observation set will contain the observations "group_1" and "group_3". This function will check that all - the observations in the 'policy' and 'critic' observation sets are present in the observation dictionary from the - environment. - - Additionally, if one of the `default_sets`, e.g. "critic", is not present in the configuration dictionary, - this function will: - - 1. Check if a group with the same name exists in the observations and assign this group to the observation set. - 2. If 1. fails, it will assign the observations from the 'policy' observation set to the default observation set. - - Args: - obs: Observations from the environment in the form of a dictionary. - obs_groups: Observation sets configuration. - default_sets: Reserved observation set names used by the algorithm (besides 'policy'). - If not provided in 'obs_groups', a default behavior gets triggered. - - Returns: - The resolved observation groups. - - Raises: - ValueError: If any observation set is an empty list. - ValueError: If any observation set contains an observation term that is not present in the observations. - """ - # check if policy observation set exists - if "policy" not in obs_groups.keys(): - if "policy" in obs: - obs_groups["policy"] = ["policy"] - warnings.warn( - "The observation configuration dictionary 'obs_groups' must contain the 'policy' key." - " As an observation group with the name 'policy' was found, this is assumed to be the observation set." - " Consider adding the 'policy' key to the 'obs_groups' dictionary for clarity." - " This behavior will be removed in a future version." - ) - else: - raise ValueError( - "The observation configuration dictionary 'obs_groups' must contain the 'policy' key." - f" Found keys: {list(obs_groups.keys())}" - ) - - # check all observation sets for valid observation groups - for set_name, groups in obs_groups.items(): - # check if the list is empty - if len(groups) == 0: - msg = f"The '{set_name}' key in the 'obs_groups' dictionary can not be an empty list." - if set_name in default_sets: - if set_name not in obs: - msg += " Consider removing the key to default to the observations used for the 'policy' set." - else: - msg += ( - f" Consider removing the key to default to the observation '{set_name}' from the environment." - ) - raise ValueError(msg) - # check groups exist inside the observations from the environment - for group in groups: - if group not in obs: - raise ValueError( - f"Observation '{group}' in observation set '{set_name}' not found in the observations from the" - f" environment. Available observations from the environment: {list(obs.keys())}" - ) - - # fill missing observation sets - for default_set_name in default_sets: - if default_set_name not in obs_groups.keys(): - if default_set_name in obs: - obs_groups[default_set_name] = [default_set_name] - warnings.warn( - f"The observation configuration dictionary 'obs_groups' must contain the '{default_set_name}' key." - f" As an observation group with the name '{default_set_name}' was found, this is assumed to be the" - f" observation set. Consider adding the '{default_set_name}' key to the 'obs_groups' dictionary for" - " clarity. This behavior will be removed in a future version." - ) - else: - obs_groups[default_set_name] = obs_groups["policy"].copy() - warnings.warn( - f"The observation configuration dictionary 'obs_groups' must contain the '{default_set_name}' key." - f" As the configuration for '{default_set_name}' is missing, the observations from the 'policy' set" - f" are used. Consider adding the '{default_set_name}' key to the 'obs_groups' dictionary for" - " clarity. This behavior will be removed in a future version." - ) - - # print the final parsed observation sets - print("-" * 80) - print("Resolved observation sets: ") - for set_name, groups in obs_groups.items(): - print("\t", set_name, ": ", groups) - print("-" * 80) - - return obs_groups diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/wandb_utils.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/wandb_utils.py deleted file mode 100644 index 243e82d..0000000 --- a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/wandb_utils.py +++ /dev/null @@ -1,87 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -from __future__ import annotations - -import os -from dataclasses import asdict -from torch.utils.tensorboard import SummaryWriter - -try: - import wandb -except ModuleNotFoundError: - raise ModuleNotFoundError("Wandb is required to log to Weights and Biases.") - - -class WandbSummaryWriter(SummaryWriter): - """Summary writer for Weights and Biases.""" - - def __init__(self, log_dir: str, flush_secs: int, cfg): - super().__init__(log_dir, flush_secs) - - # Get the run name - run_name = os.path.split(log_dir)[-1] - - try: - project = cfg["wandb_project"] - except KeyError: - raise KeyError("Please specify wandb_project in the runner config, e.g. legged_gym.") - - try: - entity = os.environ["WANDB_USERNAME"] - except KeyError: - entity = None - - # Initialize wandb - wandb.init(project=project, entity=entity, name=run_name) - - # Add log directory to wandb - wandb.config.update({"log_dir": log_dir}) - - self.name_map = { - "Train/mean_reward/time": "Train/mean_reward_time", - "Train/mean_episode_length/time": "Train/mean_episode_length_time", - } - - def store_config(self, env_cfg, runner_cfg, alg_cfg, policy_cfg): - wandb.config.update({"runner_cfg": runner_cfg}) - wandb.config.update({"policy_cfg": policy_cfg}) - wandb.config.update({"alg_cfg": alg_cfg}) - try: - wandb.config.update({"env_cfg": env_cfg.to_dict()}) - except Exception: - wandb.config.update({"env_cfg": asdict(env_cfg)}) - - def add_scalar(self, tag, scalar_value, global_step=None, walltime=None, new_style=False): - super().add_scalar( - tag, - scalar_value, - global_step=global_step, - walltime=walltime, - new_style=new_style, - ) - wandb.log({self._map_path(tag): scalar_value}, step=global_step) - - def stop(self): - wandb.finish() - - def log_config(self, env_cfg, runner_cfg, alg_cfg, policy_cfg): - self.store_config(env_cfg, runner_cfg, alg_cfg, policy_cfg) - - def save_model(self, model_path, iter): - wandb.save(model_path, base_path=os.path.dirname(model_path)) - - def save_file(self, path, iter=None): - wandb.save(path, base_path=os.path.dirname(path)) - - """ - Private methods. - """ - - def _map_path(self, path): - if path in self.name_map: - return self.name_map[path] - else: - return path diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/env/__init__.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/env/__init__.py deleted file mode 100644 index ab7c056..0000000 --- a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/env/__init__.py +++ /dev/null @@ -1,10 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -"""Submodule defining the environment definitions.""" - -from .vec_env import VecEnv - -__all__ = ["VecEnv"] diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/env/vec_env.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/env/vec_env.py deleted file mode 100644 index ae471cf..0000000 --- a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/env/vec_env.py +++ /dev/null @@ -1,113 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -from __future__ import annotations - -import torch -from abc import ABC, abstractmethod -from tensordict import TensorDict - - -class VecEnv(ABC): - """Abstract class for a vectorized environment. - - The vectorized environment is a collection of environments that are synchronized. This means that - the same type of action is applied to all environments and the same type of observation is returned from all - environments. - """ - - num_envs: int - """Number of environments.""" - - num_actions: int - """Number of actions.""" - - max_episode_length: int | torch.Tensor - - max_episode_length_s: float - """Maximum episode length. - - The maximum episode length can be a scalar or a tensor. If it is a scalar, it is the same for all environments. - If it is a tensor, it is the maximum episode length for each environment. This is useful for dynamic episode - lengths. - """ - - episode_length_buf: torch.Tensor - """Buffer for current episode lengths.""" - - device: torch.device | str - """Device to use.""" - - cfg: dict | object - """Configuration object.""" - - reset_env_ids: torch.Tensor | None = None - - contact_phase: torch.Tensor | None = None - """ - Operations. - """ - - @abstractmethod - def get_observations(self) -> TensorDict: - """Return the current observations. - - Returns: - observations (TensorDict): Observations from the environment. - """ - raise NotImplementedError - - @abstractmethod - def get_amp_observations(self) -> TensorDict: - """Return the current AMP observations. - - Returns: - observations (TensorDict): Observations from the environment. - """ - raise NotImplementedError - - @abstractmethod - def step(self, actions: torch.Tensor) -> tuple[TensorDict, torch.Tensor, torch.Tensor, dict]: - """Apply input action to the environment. - - Args: - actions (torch.Tensor): Input actions to apply. Shape: (num_envs, num_actions) - - Returns: - observations (TensorDict): Observations from the environment. - rewards (torch.Tensor): Rewards from the environment. Shape: (num_envs,) - dones (torch.Tensor): Done flags from the environment. Shape: (num_envs,) - extras (dict): Extra information from the environment. - - Observations: - - The observations TensorDict usually contains multiple observation groups. The `obs_groups` - dictionary of the runner configuration specifies which observation groups are used for which - purpose, i.e., it maps the available observation groups to observation sets. The observation sets - (keys of the `obs_groups` dictionary) currently used by rsl_rl are: - - - "policy": Specified observation groups are used as input to the actor/student network. - - "critic": Specified observation groups are used as input to the critic network. - - "teacher": Specified observation groups are used as input to the teacher network. - - "rnd_state": Specified observation groups are used as input to the RND network. - - Incomplete or incorrect configurations are handled in the `resolve_obs_groups()` function in - `rsl_rl/utils/utils.py`. - - Extras: - - The extras dictionary includes metrics such as the episode reward, episode length, etc. The following - dictionary keys are used by rsl_rl: - - - "time_outs" (torch.Tensor): Timeouts for the environments. These correspond to terminations that - happen due to time limits and not due to the environment reaching a terminal state. This is useful - for environments that have a fixed episode length. - - - "log" (dict[str, float | torch.Tensor]): Additional information for logging and debugging purposes. - The key should be a string and start with "/" for namespacing. The value can be a scalar or a - tensor. If it is a tensor, the mean of the tensor is used for logging. - """ - raise NotImplementedError - \ No newline at end of file diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/__init__.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/__init__.py deleted file mode 100644 index 9afe1a2..0000000 --- a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/__init__.py +++ /dev/null @@ -1,21 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -"""Definitions for neural-network components for RL-agents.""" - -from .actor_critic import ActorCritic -from .actor_critic_recurrent import ActorCriticRecurrent -from .rnd import * -from .student_teacher import StudentTeacher -from .student_teacher_recurrent import StudentTeacherRecurrent -from .symmetry import * -from .discriminator_multi import DiscriminatorMulti -__all__ = [ - "ActorCritic", - "ActorCriticRecurrent", - "StudentTeacher", - "StudentTeacherRecurrent", - "DiscriminatorMulti", -] diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/actor_critic.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/actor_critic.py deleted file mode 100644 index 0efc36b..0000000 --- a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/actor_critic.py +++ /dev/null @@ -1,195 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -from __future__ import annotations - -import torch -import torch.nn as nn -from torch.distributions import Normal - -from rsl_rl.networks import MLP, EmpiricalNormalization - - -class ActorCritic(nn.Module): - is_recurrent = False - - def __init__( - self, - obs, - obs_groups, - num_actions, - actor_obs_normalization=False, - critic_obs_normalization=False, - actor_hidden_dims=[256, 256, 256], - critic_hidden_dims=[256, 256, 256], - activation="elu", - init_noise_std=1.0, - noise_std_type: str = "scalar", - state_dependent_std=False, - **kwargs, - ): - if kwargs: - print( - "ActorCritic.__init__ got unexpected arguments, which will be ignored: " - + str([key for key in kwargs.keys()]) - ) - super().__init__() - - # get the observation dimensions - self.obs_groups = obs_groups - num_actor_obs = 0 - for obs_group in obs_groups["policy"]: - assert len(obs[obs_group].shape) == 2, "The ActorCritic module only supports 1D observations." - num_actor_obs += obs[obs_group].shape[-1] - num_critic_obs = 0 - for obs_group in obs_groups["critic"]: - assert len(obs[obs_group].shape) == 2, "The ActorCritic module only supports 1D observations." - num_critic_obs += obs[obs_group].shape[-1] - - self.state_dependent_std = state_dependent_std - # actor - if self.state_dependent_std: - self.actor = MLP(num_actor_obs, [2, num_actions], actor_hidden_dims, activation) - else: - self.actor = MLP(num_actor_obs, num_actions, actor_hidden_dims, activation) - # actor observation normalization - self.actor_obs_normalization = actor_obs_normalization - if actor_obs_normalization: - self.actor_obs_normalizer = EmpiricalNormalization(num_actor_obs) - else: - self.actor_obs_normalizer = torch.nn.Identity() - print(f"Actor MLP: {self.actor}") - - # critic - self.critic = MLP(num_critic_obs, 1, critic_hidden_dims, activation) - # critic observation normalization - self.critic_obs_normalization = critic_obs_normalization - if critic_obs_normalization: - self.critic_obs_normalizer = EmpiricalNormalization(num_critic_obs) - else: - self.critic_obs_normalizer = torch.nn.Identity() - print(f"Critic MLP: {self.critic}") - - # Action noise - self.noise_std_type = noise_std_type - if self.state_dependent_std: - torch.nn.init.zeros_(self.actor[-2].weight[num_actions:]) - if self.noise_std_type == "scalar": - torch.nn.init.constant_(self.actor[-2].bias[num_actions:], init_noise_std) - elif self.noise_std_type == "log": - torch.nn.init.constant_( - self.actor[-2].bias[num_actions:], torch.log(torch.tensor(init_noise_std + 1e-7)) - ) - else: - raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") - else: - if self.noise_std_type == "scalar": - self.std = nn.Parameter(init_noise_std * torch.ones(num_actions)) - elif self.noise_std_type == "log": - self.log_std = nn.Parameter(torch.log(init_noise_std * torch.ones(num_actions))) - else: - raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") - - # Action distribution (populated in update_distribution) - self.distribution = None - # disable args validation for speedup - Normal.set_default_validate_args(False) - - def reset(self, dones=None): - pass - - def forward(self): - raise NotImplementedError - - @property - def action_mean(self): - return self.distribution.mean - - @property - def action_std(self): - return self.distribution.stddev - - @property - def entropy(self): - return self.distribution.entropy().sum(dim=-1) - - def update_distribution(self, obs): - if self.state_dependent_std: - # compute mean and standard deviation - mean_and_std = self.actor(obs) - if self.noise_std_type == "scalar": - mean, std = torch.unbind(mean_and_std, dim=-2) - elif self.noise_std_type == "log": - mean, log_std = torch.unbind(mean_and_std, dim=-2) - std = torch.exp(log_std) - else: - raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") - else: - # compute mean - mean = self.actor(obs) - # compute standard deviation - if self.noise_std_type == "scalar": - std = self.std.expand_as(mean) - elif self.noise_std_type == "log": - std = torch.exp(self.log_std).expand_as(mean) - else: - raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") - # create distribution - self.distribution = Normal(mean, std) - - def act(self, obs, **kwargs): - obs = self.get_actor_obs(obs) - obs = self.actor_obs_normalizer(obs) - self.update_distribution(obs) - return self.distribution.sample() - - def act_inference(self, obs): - obs = self.get_actor_obs(obs) - obs = self.actor_obs_normalizer(obs) - return self.actor(obs) - - def evaluate(self, obs, **kwargs): - obs = self.get_critic_obs(obs) - obs = self.critic_obs_normalizer(obs) - return self.critic(obs) - - def get_actor_obs(self, obs): - obs_list = [] - for obs_group in self.obs_groups["policy"]: - obs_list.append(obs[obs_group]) - return torch.cat(obs_list, dim=-1) - - def get_critic_obs(self, obs): - obs_list = [] - for obs_group in self.obs_groups["critic"]: - obs_list.append(obs[obs_group]) - return torch.cat(obs_list, dim=-1) - - def get_actions_log_prob(self, actions): - return self.distribution.log_prob(actions).sum(dim=-1) - - def update_normalization(self, obs): - if self.actor_obs_normalization: - actor_obs = self.get_actor_obs(obs) - self.actor_obs_normalizer.update(actor_obs) - if self.critic_obs_normalization: - critic_obs = self.get_critic_obs(obs) - self.critic_obs_normalizer.update(critic_obs) - - def load_state_dict(self, state_dict, strict=True): - """Load the parameters of the actor-critic model. - - Args: - state_dict (dict): State dictionary of the model. - strict (bool): Whether to strictly enforce that the keys in state_dict match the keys returned by this - module's state_dict() function. - - Returns: - bool: Whether this training resumes a previous training. This flag is used by the `load()` function of - `OnPolicyRunner` to determine how to load further parameters (relevant for, e.g., distillation). - """ - - super().load_state_dict(state_dict, strict=strict) - return True # training resumes diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/actor_critic_recurrent.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/actor_critic_recurrent.py deleted file mode 100644 index bba46ca..0000000 --- a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/actor_critic_recurrent.py +++ /dev/null @@ -1,218 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -from __future__ import annotations - -import torch -import torch.nn as nn -import warnings -from torch.distributions import Normal - -from rsl_rl.networks import MLP, EmpiricalNormalization, Memory - - -class ActorCriticRecurrent(nn.Module): - is_recurrent = True - - def __init__( - self, - obs, - obs_groups, - num_actions, - actor_obs_normalization=False, - critic_obs_normalization=False, - actor_hidden_dims=[256, 256, 256], - critic_hidden_dims=[256, 256, 256], - activation="elu", - init_noise_std=1.0, - noise_std_type: str = "scalar", - state_dependent_std=False, - rnn_type="lstm", - rnn_hidden_dim=256, - rnn_num_layers=1, - **kwargs, - ): - if "rnn_hidden_size" in kwargs: - warnings.warn( - "The argument `rnn_hidden_size` is deprecated and will be removed in a future version. " - "Please use `rnn_hidden_dim` instead.", - DeprecationWarning, - ) - if rnn_hidden_dim == 256: # Only override if the new argument is at its default - rnn_hidden_dim = kwargs.pop("rnn_hidden_size") - if kwargs: - print( - "ActorCriticRecurrent.__init__ got unexpected arguments, which will be ignored: " + str(kwargs.keys()), - ) - super().__init__() - - # get the observation dimensions - self.obs_groups = obs_groups - num_actor_obs = 0 - for obs_group in obs_groups["policy"]: - assert len(obs[obs_group].shape) == 2, "The ActorCriticRecurrent module only supports 1D observations." - num_actor_obs += obs[obs_group].shape[-1] - num_critic_obs = 0 - for obs_group in obs_groups["critic"]: - assert len(obs[obs_group].shape) == 2, "The ActorCriticRecurrent module only supports 1D observations." - num_critic_obs += obs[obs_group].shape[-1] - - self.state_dependent_std = state_dependent_std - # actor - self.memory_a = Memory(num_actor_obs, type=rnn_type, num_layers=rnn_num_layers, hidden_size=rnn_hidden_dim) - if self.state_dependent_std: - self.actor = MLP(rnn_hidden_dim, [2, num_actions], actor_hidden_dims, activation) - else: - self.actor = MLP(rnn_hidden_dim, num_actions, actor_hidden_dims, activation) - - # actor observation normalization - self.actor_obs_normalization = actor_obs_normalization - if actor_obs_normalization: - self.actor_obs_normalizer = EmpiricalNormalization(num_actor_obs) - else: - self.actor_obs_normalizer = torch.nn.Identity() - print(f"Actor RNN: {self.memory_a}") - print(f"Actor MLP: {self.actor}") - - # critic - self.memory_c = Memory(num_critic_obs, type=rnn_type, num_layers=rnn_num_layers, hidden_size=rnn_hidden_dim) - self.critic = MLP(rnn_hidden_dim, 1, critic_hidden_dims, activation) - # critic observation normalization - self.critic_obs_normalization = critic_obs_normalization - if critic_obs_normalization: - self.critic_obs_normalizer = EmpiricalNormalization(num_critic_obs) - else: - self.critic_obs_normalizer = torch.nn.Identity() - print(f"Critic RNN: {self.memory_c}") - print(f"Critic MLP: {self.critic}") - - # Action noise - self.noise_std_type = noise_std_type - if self.state_dependent_std: - torch.nn.init.zeros_(self.actor[-2].weight[num_actions:]) - if self.noise_std_type == "scalar": - torch.nn.init.constant_(self.actor[-2].bias[num_actions:], init_noise_std) - elif self.noise_std_type == "log": - torch.nn.init.constant_( - self.actor[-2].bias[num_actions:], torch.log(torch.tensor(init_noise_std + 1e-7)) - ) - else: - raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") - else: - if self.noise_std_type == "scalar": - self.std = nn.Parameter(init_noise_std * torch.ones(num_actions)) - elif self.noise_std_type == "log": - self.log_std = nn.Parameter(torch.log(init_noise_std * torch.ones(num_actions))) - else: - raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") - - # Action distribution (populated in update_distribution) - self.distribution = None - # disable args validation for speedup - Normal.set_default_validate_args(False) - - @property - def action_mean(self): - return self.distribution.mean - - @property - def action_std(self): - return self.distribution.stddev - - @property - def entropy(self): - return self.distribution.entropy().sum(dim=-1) - - def reset(self, dones=None): - self.memory_a.reset(dones) - self.memory_c.reset(dones) - - def forward(self): - raise NotImplementedError - - def update_distribution(self, obs): - if self.state_dependent_std: - # compute mean and standard deviation - mean_and_std = self.actor(obs) - if self.noise_std_type == "scalar": - mean, std = torch.unbind(mean_and_std, dim=-2) - elif self.noise_std_type == "log": - mean, log_std = torch.unbind(mean_and_std, dim=-2) - std = torch.exp(log_std) - else: - raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") - else: - # compute mean - mean = self.actor(obs) - # compute standard deviation - if self.noise_std_type == "scalar": - std = self.std.expand_as(mean) - elif self.noise_std_type == "log": - std = torch.exp(self.log_std).expand_as(mean) - else: - raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") - # create distribution - self.distribution = Normal(mean, std) - - def act(self, obs, masks=None, hidden_states=None): - obs = self.get_actor_obs(obs) - obs = self.actor_obs_normalizer(obs) - out_mem = self.memory_a(obs, masks, hidden_states).squeeze(0) - self.update_distribution(out_mem) - return self.distribution.sample() - - def act_inference(self, obs): - obs = self.get_actor_obs(obs) - obs = self.actor_obs_normalizer(obs) - out_mem = self.memory_a(obs).squeeze(0) - return self.actor(out_mem) - - def evaluate(self, obs, masks=None, hidden_states=None): - obs = self.get_critic_obs(obs) - obs = self.critic_obs_normalizer(obs) - out_mem = self.memory_c(obs, masks, hidden_states).squeeze(0) - return self.critic(out_mem) - - def get_actor_obs(self, obs): - obs_list = [] - for obs_group in self.obs_groups["policy"]: - obs_list.append(obs[obs_group]) - return torch.cat(obs_list, dim=-1) - - def get_critic_obs(self, obs): - obs_list = [] - for obs_group in self.obs_groups["critic"]: - obs_list.append(obs[obs_group]) - return torch.cat(obs_list, dim=-1) - - def get_actions_log_prob(self, actions): - return self.distribution.log_prob(actions).sum(dim=-1) - - def get_hidden_states(self): - return self.memory_a.hidden_states, self.memory_c.hidden_states - - def update_normalization(self, obs): - if self.actor_obs_normalization: - actor_obs = self.get_actor_obs(obs) - self.actor_obs_normalizer.update(actor_obs) - if self.critic_obs_normalization: - critic_obs = self.get_critic_obs(obs) - self.critic_obs_normalizer.update(critic_obs) - - def load_state_dict(self, state_dict, strict=True): - """Load the parameters of the actor-critic model. - - Args: - state_dict (dict): State dictionary of the model. - strict (bool): Whether to strictly enforce that the keys in state_dict match the keys returned by this - module's state_dict() function. - - Returns: - bool: Whether this training resumes a previous training. This flag is used by the `load()` function of - `OnPolicyRunner` to determine how to load further parameters (relevant for, e.g., distillation). - """ - - super().load_state_dict(state_dict, strict=strict) - return True diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/discriminator_multi.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/discriminator_multi.py deleted file mode 100644 index 7899d23..0000000 --- a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/discriminator_multi.py +++ /dev/null @@ -1,102 +0,0 @@ -import torch -import torch.nn as nn -from torch import autograd -import torch.nn.utils.spectral_norm as spectral_norm - - -class DiscriminatorMulti(nn.Module): - def __init__( - self, state_dim, amp_reward_coef, hidden_layer_sizes, device, - num_frames=2, task_reward_lerp=0.0, use_lerp=True): - super(DiscriminatorMulti, self).__init__() - - self.device = device - self.state_dim = state_dim - self.use_lerp = use_lerp - self.num_frames = num_frames # 存储帧数参数 - - self.amp_reward_coef = amp_reward_coef - amp_layers = [] - - curr_in_dim = state_dim * num_frames - for hidden_dim in hidden_layer_sizes: - amp_layers.append(spectral_norm(nn.Linear(curr_in_dim, hidden_dim))) - amp_layers.append(nn.ReLU()) - curr_in_dim = hidden_dim - self.trunk = nn.Sequential(*amp_layers).to(device) - self.amp_linear = spectral_norm(nn.Linear(hidden_layer_sizes[-1], 1)).to(device) - - self.trunk.train() - self.amp_linear.train() - - self.task_reward_lerp = task_reward_lerp - - def forward(self, x): - h = self.trunk(x) - d = self.amp_linear(h) - return d - - def compute_grad_pen(self, - expert_states, # 改为接收多帧状态列表 - lambda_=10): - # 将多帧状态沿最后一个维度拼接 - expert_data = expert_states.flatten(1) - expert_data.requires_grad = True - - disc = self.amp_linear(self.trunk(expert_data)) - ones = torch.ones(disc.size(), device=disc.device) - grad = autograd.grad( - outputs=disc, inputs=expert_data, - grad_outputs=ones, create_graph=True, - retain_graph=True, only_inputs=True)[0] - - # Enforce that the grad norm approaches 0. - grad_pen = lambda_ * (grad.norm(2, dim=1) - 0).pow(2).mean() - return grad_pen - - - def get_disc_weights(self): - weights = [] - for m in self.trunk.modules(): - if isinstance(m, nn.Linear): - weights.append(torch.flatten(m.weight)) - - weights.append(torch.flatten(self.amp_linear.weight)) - return weights - - def get_disc_logit_weights(self): - return torch.flatten(self.amp_linear.weight) - - def predict_amp_reward( - self, states, # 改为接收多帧状态列表 - task_reward, normalizer=None): - """ - states: torch.Tensor, shape=(num_envs, num_frames, state_dim) - task_reward: torch.Tensor, shape=(num_envs, 1) - """ - # import ipdb; ipdb.set_trace() - with torch.no_grad(): - self.eval() - if normalizer is not None: - # 对每一帧状态进行归一化 - states = normalizer.normalize_torch(states, self.device) - - # 拼接多帧状态 - state_cat = states.flatten(1) - d = self.amp_linear(self.trunk(state_cat)) - disc_reward = self.amp_reward_coef * torch.clamp(1 - (1/4) * torch.square(d - 1), min=0) - - if self.use_lerp: - if self.task_reward_lerp > 0: - reward = self._lerp_reward(disc_reward, task_reward.unsqueeze(-1)) - self.train() - return reward.squeeze(), d, disc_reward.squeeze() * (1.0 - self.task_reward_lerp) - else: - disc_reward *= 0.02 - reward = task_reward.unsqueeze(-1) + disc_reward - self.train() - return reward.squeeze(), d, disc_reward.squeeze() - - def _lerp_reward(self, disc_r, task_r): - r = (1.0 - self.task_reward_lerp) * disc_r + self.task_reward_lerp * task_r - return r \ No newline at end of file diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/rnd.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/rnd.py deleted file mode 100644 index 8e65c43..0000000 --- a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/rnd.py +++ /dev/null @@ -1,209 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -from __future__ import annotations - -import torch -import torch.nn as nn - -from rsl_rl.networks import MLP, EmpiricalDiscountedVariationNormalization, EmpiricalNormalization - - -class RandomNetworkDistillation(nn.Module): - """Implementation of Random Network Distillation (RND) [1] - - References: - .. [1] Burda, Yuri, et al. "Exploration by random network distillation." arXiv preprint arXiv:1810.12894 (2018). - """ - - def __init__( - self, - num_states: int, - obs_groups: dict, - num_outputs: int, - predictor_hidden_dims: list[int], - target_hidden_dims: list[int], - activation: str = "elu", - weight: float = 0.0, - state_normalization: bool = False, - reward_normalization: bool = False, - device: str = "cpu", - weight_schedule: dict | None = None, - ): - """Initialize the RND module. - - - If :attr:`state_normalization` is True, then the input state is normalized using an Empirical Normalization layer. - - If :attr:`reward_normalization` is True, then the intrinsic reward is normalized using an Empirical Discounted - Variation Normalization layer. - - .. note:: - If the hidden dimensions are -1 in the predictor and target networks configuration, then the number of states - is used as the hidden dimension. - - Args: - num_states: Number of states/inputs to the predictor and target networks. - num_outputs: Number of outputs (embedding size) of the predictor and target networks. - predictor_hidden_dims: List of hidden dimensions of the predictor network. - target_hidden_dims: List of hidden dimensions of the target network. - activation: Activation function. Defaults to "elu". - weight: Scaling factor of the intrinsic reward. Defaults to 0.0. - state_normalization: Whether to normalize the input state. Defaults to False. - reward_normalization: Whether to normalize the intrinsic reward. Defaults to False. - device: Device to use. Defaults to "cpu". - weight_schedule: The type of schedule to use for the RND weight parameter. - Defaults to None, in which case the weight parameter is constant. - It is a dictionary with the following keys: - - - "mode": The type of schedule to use for the RND weight parameter. - - "constant": Constant weight schedule. - - "step": Step weight schedule. - - "linear": Linear weight schedule. - - For the "step" weight schedule, the following parameters are required: - - - "final_step": The step at which the weight parameter is set to the final value. - - "final_value": The final value of the weight parameter. - - For the "linear" weight schedule, the following parameters are required: - - "initial_step": The step at which the weight parameter is set to the initial value. - - "final_step": The step at which the weight parameter is set to the final value. - - "final_value": The final value of the weight parameter. - """ - # initialize parent class - super().__init__() - - # Store parameters - self.num_states = num_states - self.obs_groups = obs_groups - self.num_outputs = num_outputs - self.initial_weight = weight - self.device = device - self.state_normalization = state_normalization - self.reward_normalization = reward_normalization - - # Normalization of input gates - if state_normalization: - self.state_normalizer = EmpiricalNormalization(shape=[self.num_states], until=1.0e8).to(self.device) - else: - self.state_normalizer = torch.nn.Identity() - # Normalization of intrinsic reward - if reward_normalization: - self.reward_normalizer = EmpiricalDiscountedVariationNormalization(shape=[], until=1.0e8).to(self.device) - else: - self.reward_normalizer = torch.nn.Identity() - - # counter for the number of updates - self.update_counter = 0 - - # resolve weight schedule - if weight_schedule is not None: - self.weight_scheduler_params = weight_schedule - self.weight_scheduler = getattr(self, f"_{weight_schedule['mode']}_weight_schedule") - else: - self.weight_scheduler = None - # Create network architecture - self.predictor = MLP(num_states, num_outputs, predictor_hidden_dims, activation).to(self.device) - self.target = MLP(num_states, num_outputs, target_hidden_dims, activation).to(self.device) - - # make target network not trainable - self.target.eval() - - def get_intrinsic_reward(self, obs) -> torch.Tensor: - # Note: the counter is updated number of env steps per learning iteration - self.update_counter += 1 - # Extract the rnd state from the observation - rnd_state = self.get_rnd_state(obs) - rnd_state = self.state_normalizer(rnd_state) - # Obtain the embedding of the rnd state from the target and predictor networks - target_embedding = self.target(rnd_state).detach() - predictor_embedding = self.predictor(rnd_state).detach() - # Compute the intrinsic reward as the distance between the embeddings - intrinsic_reward = torch.linalg.norm(target_embedding - predictor_embedding, dim=1) - # Normalize intrinsic reward - intrinsic_reward = self.reward_normalizer(intrinsic_reward) - - # Check the weight schedule - if self.weight_scheduler is not None: - self.weight = self.weight_scheduler(step=self.update_counter, **self.weight_scheduler_params) - else: - self.weight = self.initial_weight - # Scale intrinsic reward - intrinsic_reward *= self.weight - - return intrinsic_reward - - def forward(self, *args, **kwargs): - raise RuntimeError("Forward method is not implemented. Use get_intrinsic_reward instead.") - - def train(self, mode: bool = True): - # sets module into training mode - self.predictor.train(mode) - if self.state_normalization: - self.state_normalizer.train(mode) - if self.reward_normalization: - self.reward_normalizer.train(mode) - return self - - def eval(self): - return self.train(False) - - def get_rnd_state(self, obs): - obs_list = [] - for obs_group in self.obs_groups["rnd_state"]: - obs_list.append(obs[obs_group]) - return torch.cat(obs_list, dim=-1) - - def update_normalization(self, obs): - # Normalize the state - if self.state_normalization: - rnd_state = self.get_rnd_state(obs) - self.state_normalizer.update(rnd_state) - - """ - Different weight schedules. - """ - - def _constant_weight_schedule(self, step: int, **kwargs): - return self.initial_weight - - def _step_weight_schedule(self, step: int, final_step: int, final_value: float, **kwargs): - return self.initial_weight if step < final_step else final_value - - def _linear_weight_schedule(self, step: int, initial_step: int, final_step: int, final_value: float, **kwargs): - if step < initial_step: - return self.initial_weight - elif step > final_step: - return final_value - else: - return self.initial_weight + (final_value - self.initial_weight) * (step - initial_step) / ( - final_step - initial_step - ) - - -def resolve_rnd_config(alg_cfg, obs, obs_groups, env): - """Resolve the RND configuration. - - Args: - alg_cfg: The algorithm configuration dictionary. - obs: The observation dictionary. - obs_groups: The observation groups dictionary. - env: The environment. - - Returns: - The resolved algorithm configuration dictionary. - """ - # resolve dimension of rnd gated state - if "rnd_cfg" in alg_cfg and alg_cfg["rnd_cfg"] is not None: - # get dimension of rnd gated state - num_rnd_state = 0 - for obs_group in obs_groups["rnd_state"]: - assert len(obs[obs_group].shape) == 2, "The RND module only supports 1D observations." - num_rnd_state += obs[obs_group].shape[-1] - # add rnd gated state to config - alg_cfg["rnd_cfg"]["num_states"] = num_rnd_state - alg_cfg["rnd_cfg"]["obs_groups"] = obs_groups - # scale down the rnd weight with timestep - alg_cfg["rnd_cfg"]["weight"] *= env.unwrapped.step_dt - return alg_cfg diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/student_teacher.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/student_teacher.py deleted file mode 100644 index 6bf1380..0000000 --- a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/student_teacher.py +++ /dev/null @@ -1,206 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -from __future__ import annotations - -import torch -import torch.nn as nn -from torch.distributions import Normal - -from rsl_rl.networks import MLP, EmpiricalNormalization - - -class StudentTeacher(nn.Module): - is_recurrent = False - - def __init__( - self, - obs, - obs_groups, - num_actions, - student_obs_normalization=False, - teacher_obs_normalization=False, - student_hidden_dims=[256, 256, 256], - teacher_hidden_dims=[256, 256, 256], - activation="elu", - init_noise_std=0.1, - noise_std_type: str = "scalar", - **kwargs, - ): - if kwargs: - print( - "StudentTeacher.__init__ got unexpected arguments, which will be ignored: " - + str([key for key in kwargs.keys()]) - ) - super().__init__() - - self.loaded_teacher = False # indicates if teacher has been loaded - - # get the observation dimensions - self.obs_groups = obs_groups - num_student_obs = 0 - for obs_group in obs_groups["policy"]: - assert len(obs[obs_group].shape) == 2, "The StudentTeacher module only supports 1D observations." - num_student_obs += obs[obs_group].shape[-1] - num_teacher_obs = 0 - for obs_group in obs_groups["teacher"]: - assert len(obs[obs_group].shape) == 2, "The StudentTeacher module only supports 1D observations." - num_teacher_obs += obs[obs_group].shape[-1] - - # student - self.student = MLP(num_student_obs, num_actions, student_hidden_dims, activation) - - # student observation normalization - self.student_obs_normalization = student_obs_normalization - if student_obs_normalization: - self.student_obs_normalizer = EmpiricalNormalization(num_student_obs) - else: - self.student_obs_normalizer = torch.nn.Identity() - - print(f"Student MLP: {self.student}") - - # teacher - self.teacher = MLP(num_teacher_obs, num_actions, teacher_hidden_dims, activation) - self.teacher.eval() - - # teacher observation normalization - self.teacher_obs_normalization = teacher_obs_normalization - if teacher_obs_normalization: - self.teacher_obs_normalizer = EmpiricalNormalization(num_teacher_obs) - else: - self.teacher_obs_normalizer = torch.nn.Identity() - - print(f"Teacher MLP: {self.teacher}") - - # action noise - self.noise_std_type = noise_std_type - if self.noise_std_type == "scalar": - self.std = nn.Parameter(init_noise_std * torch.ones(num_actions)) - elif self.noise_std_type == "log": - self.log_std = nn.Parameter(torch.log(init_noise_std * torch.ones(num_actions))) - else: - raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") - - # action distribution (populated in update_distribution) - self.distribution = None - # disable args validation for speedup - Normal.set_default_validate_args(False) - - def reset(self, dones=None, hidden_states=None): - pass - - def forward(self): - raise NotImplementedError - - @property - def action_mean(self): - return self.distribution.mean - - @property - def action_std(self): - return self.distribution.stddev - - @property - def entropy(self): - return self.distribution.entropy().sum(dim=-1) - - def update_distribution(self, obs): - # compute mean - mean = self.student(obs) - # compute standard deviation - if self.noise_std_type == "scalar": - std = self.std.expand_as(mean) - elif self.noise_std_type == "log": - std = torch.exp(self.log_std).expand_as(mean) - else: - raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") - # create distribution - self.distribution = Normal(mean, std) - - def act(self, obs): - obs = self.get_student_obs(obs) - obs = self.student_obs_normalizer(obs) - self.update_distribution(obs) - return self.distribution.sample() - - def act_inference(self, obs): - obs = self.get_student_obs(obs) - obs = self.student_obs_normalizer(obs) - return self.student(obs) - - def evaluate(self, obs): - obs = self.get_teacher_obs(obs) - obs = self.teacher_obs_normalizer(obs) - with torch.no_grad(): - return self.teacher(obs) - - def get_student_obs(self, obs): - obs_list = [] - for obs_group in self.obs_groups["policy"]: - obs_list.append(obs[obs_group]) - return torch.cat(obs_list, dim=-1) - - def get_teacher_obs(self, obs): - obs_list = [] - for obs_group in self.obs_groups["teacher"]: - obs_list.append(obs[obs_group]) - return torch.cat(obs_list, dim=-1) - - def get_hidden_states(self): - return None - - def detach_hidden_states(self, dones=None): - pass - - def train(self, mode=True): - super().train(mode) - # make sure teacher is in eval mode - self.teacher.eval() - self.teacher_obs_normalizer.eval() - - def update_normalization(self, obs): - if self.student_obs_normalization: - student_obs = self.get_student_obs(obs) - self.student_obs_normalizer.update(student_obs) - - def load_state_dict(self, state_dict, strict=True): - """Load the parameters of the student and teacher networks. - - Args: - state_dict (dict): State dictionary of the model. - strict (bool): Whether to strictly enforce that the keys in state_dict match the keys returned by this - module's state_dict() function. - - Returns: - bool: Whether this training resumes a previous training. This flag is used by the `load()` function of - `OnPolicyRunner` to determine how to load further parameters. - """ - - # check if state_dict contains teacher and student or just teacher parameters - if any("actor" in key for key in state_dict.keys()): # loading parameters from rl training - # rename keys to match teacher and remove critic parameters - teacher_state_dict = {} - teacher_obs_normalizer_state_dict = {} - for key, value in state_dict.items(): - if "actor." in key: - teacher_state_dict[key.replace("actor.", "")] = value - if "actor_obs_normalizer." in key: - teacher_obs_normalizer_state_dict[key.replace("actor_obs_normalizer.", "")] = value - self.teacher.load_state_dict(teacher_state_dict, strict=strict) - self.teacher_obs_normalizer.load_state_dict(teacher_obs_normalizer_state_dict, strict=strict) - # set flag for successfully loading the parameters - self.loaded_teacher = True - self.teacher.eval() - self.teacher_obs_normalizer.eval() - return False # training does not resume - elif any("student" in key for key in state_dict.keys()): # loading parameters from distillation training - super().load_state_dict(state_dict, strict=strict) - # set flag for successfully loading the parameters - self.loaded_teacher = True - self.teacher.eval() - self.teacher_obs_normalizer.eval() - return True # training resumes - else: - raise ValueError("state_dict does not contain student or teacher parameters") diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/student_teacher_recurrent.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/student_teacher_recurrent.py deleted file mode 100644 index 964a2dc..0000000 --- a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/student_teacher_recurrent.py +++ /dev/null @@ -1,249 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -from __future__ import annotations - -import torch -import torch.nn as nn -import warnings -from torch.distributions import Normal - -from rsl_rl.networks import MLP, EmpiricalNormalization, Memory - - -class StudentTeacherRecurrent(nn.Module): - is_recurrent = True - - def __init__( - self, - obs, - obs_groups, - num_actions, - student_obs_normalization=False, - teacher_obs_normalization=False, - student_hidden_dims=[256, 256, 256], - teacher_hidden_dims=[256, 256, 256], - activation="elu", - init_noise_std=0.1, - noise_std_type: str = "scalar", - rnn_type="lstm", - rnn_hidden_dim=256, - rnn_num_layers=1, - teacher_recurrent=False, - **kwargs, - ): - if "rnn_hidden_size" in kwargs: - warnings.warn( - "The argument `rnn_hidden_size` is deprecated and will be removed in a future version. " - "Please use `rnn_hidden_dim` instead.", - DeprecationWarning, - ) - if rnn_hidden_dim == 256: # Only override if the new argument is at its default - rnn_hidden_dim = kwargs.pop("rnn_hidden_size") - if kwargs: - print( - "StudentTeacherRecurrent.__init__ got unexpected arguments, which will be ignored: " - + str(kwargs.keys()), - ) - super().__init__() - - self.loaded_teacher = False # indicates if teacher has been loaded - self.teacher_recurrent = teacher_recurrent # indicates if teacher is recurrent too - - # get the observation dimensions - self.obs_groups = obs_groups - num_student_obs = 0 - for obs_group in obs_groups["policy"]: - assert len(obs[obs_group].shape) == 2, "The StudentTeacher module only supports 1D observations." - num_student_obs += obs[obs_group].shape[-1] - num_teacher_obs = 0 - for obs_group in obs_groups["teacher"]: - assert len(obs[obs_group].shape) == 2, "The StudentTeacher module only supports 1D observations." - num_teacher_obs += obs[obs_group].shape[-1] - - # student - self.memory_s = Memory(num_student_obs, type=rnn_type, num_layers=rnn_num_layers, hidden_size=rnn_hidden_dim) - self.student = MLP(rnn_hidden_dim, num_actions, student_hidden_dims, activation) - - # student observation normalization - self.student_obs_normalization = student_obs_normalization - if student_obs_normalization: - self.student_obs_normalizer = EmpiricalNormalization(num_student_obs) - else: - self.student_obs_normalizer = torch.nn.Identity() - - print(f"Student RNN: {self.memory_s}") - print(f"Student MLP: {self.student}") - - # teacher - if self.teacher_recurrent: - self.memory_t = Memory( - num_teacher_obs, type=rnn_type, num_layers=rnn_num_layers, hidden_size=rnn_hidden_dim - ) - num_teacher_obs = rnn_hidden_dim - self.teacher = MLP(num_teacher_obs, num_actions, teacher_hidden_dims, activation) - - # teacher observation normalization - self.teacher_obs_normalization = teacher_obs_normalization - if teacher_obs_normalization: - self.teacher_obs_normalizer = EmpiricalNormalization(num_teacher_obs) - else: - self.teacher_obs_normalizer = torch.nn.Identity() - - if self.teacher_recurrent: - print(f"Teacher RNN: {self.memory_t}") - print(f"Teacher MLP: {self.teacher}") - - # action noise - self.noise_std_type = noise_std_type - if self.noise_std_type == "scalar": - self.std = nn.Parameter(init_noise_std * torch.ones(num_actions)) - elif self.noise_std_type == "log": - self.log_std = nn.Parameter(torch.log(init_noise_std * torch.ones(num_actions))) - else: - raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") - - # action distribution (populated in update_distribution) - self.distribution = None - # disable args validation for speedup - Normal.set_default_validate_args(False) - - def reset(self, dones=None, hidden_states=None): - if hidden_states is None: - hidden_states = (None, None) - self.memory_s.reset(dones, hidden_states[0]) - if self.teacher_recurrent: - self.memory_t.reset(dones, hidden_states[1]) - - def forward(self): - raise NotImplementedError - - @property - def action_mean(self): - return self.distribution.mean - - @property - def action_std(self): - return self.distribution.stddev - - @property - def entropy(self): - return self.distribution.entropy().sum(dim=-1) - - def update_distribution(self, obs): - # compute mean - mean = self.student(obs) - # compute standard deviation - if self.noise_std_type == "scalar": - std = self.std.expand_as(mean) - elif self.noise_std_type == "log": - std = torch.exp(self.log_std).expand_as(mean) - else: - raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") - # create distribution - self.distribution = Normal(mean, std) - - def act(self, obs): - obs = self.get_student_obs(obs) - obs = self.student_obs_normalizer(obs) - out_mem = self.memory_s(obs).squeeze(0) - self.update_distribution(out_mem) - return self.distribution.sample() - - def act_inference(self, obs): - obs = self.get_student_obs(obs) - obs = self.student_obs_normalizer(obs) - out_mem = self.memory_s(obs).squeeze(0) - return self.student(out_mem) - - def evaluate(self, obs): - obs = self.get_teacher_obs(obs) - obs = self.teacher_obs_normalizer(obs) - with torch.no_grad(): - if self.teacher_recurrent: - self.memory_t.eval() - obs = self.memory_t(obs).squeeze(0) - return self.teacher(obs) - - def get_student_obs(self, obs): - obs_list = [] - for obs_group in self.obs_groups["policy"]: - obs_list.append(obs[obs_group]) - return torch.cat(obs_list, dim=-1) - - def get_teacher_obs(self, obs): - obs_list = [] - for obs_group in self.obs_groups["teacher"]: - obs_list.append(obs[obs_group]) - return torch.cat(obs_list, dim=-1) - - def get_hidden_states(self): - if self.teacher_recurrent: - return self.memory_s.hidden_states, self.memory_t.hidden_states - else: - return self.memory_s.hidden_states, None - - def detach_hidden_states(self, dones=None): - self.memory_s.detach_hidden_states(dones) - if self.teacher_recurrent: - self.memory_t.detach_hidden_states(dones) - - def train(self, mode=True): - super().train(mode) - # make sure teacher is in eval mode - self.teacher.eval() - self.teacher_obs_normalizer.eval() - - def update_normalization(self, obs): - if self.student_obs_normalization: - student_obs = self.get_student_obs(obs) - self.student_obs_normalizer.update(student_obs) - - def load_state_dict(self, state_dict, strict=True): - """Load the parameters of the student and teacher networks. - - Args: - state_dict (dict): State dictionary of the model. - strict (bool): Whether to strictly enforce that the keys in state_dict match the keys returned by this - module's state_dict() function. - - Returns: - bool: Whether this training resumes a previous training. This flag is used by the `load()` function of - `OnPolicyRunner` to determine how to load further parameters. - """ - - # check if state_dict contains teacher and student or just teacher parameters - if any("actor" in key for key in state_dict.keys()): # loading parameters from rl training - # rename keys to match teacher and remove critic parameters - teacher_state_dict = {} - teacher_obs_normalizer_state_dict = {} - for key, value in state_dict.items(): - if "actor." in key: - teacher_state_dict[key.replace("actor.", "")] = value - if "actor_obs_normalizer." in key: - teacher_obs_normalizer_state_dict[key.replace("actor_obs_normalizer.", "")] = value - self.teacher.load_state_dict(teacher_state_dict, strict=strict) - self.teacher_obs_normalizer.load_state_dict(teacher_obs_normalizer_state_dict, strict=strict) - # also load recurrent memory if teacher is recurrent - if self.teacher_recurrent: - memory_t_state_dict = {} - for key, value in state_dict.items(): - if "memory_a." in key: - memory_t_state_dict[key.replace("memory_a.", "")] = value - self.memory_t.load_state_dict(memory_t_state_dict, strict=strict) - # set flag for successfully loading the parameters - self.loaded_teacher = True - self.teacher.eval() - self.teacher_obs_normalizer.eval() - return False # training does not resume - elif any("student" in key for key in state_dict.keys()): # loading parameters from distillation training - super().load_state_dict(state_dict, strict=strict) - # set flag for successfully loading the parameters - self.loaded_teacher = True - self.teacher.eval() - self.teacher_obs_normalizer.eval() - return True # training resumes - else: - raise ValueError("state_dict does not contain student or teacher parameters") diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/symmetry.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/symmetry.py deleted file mode 100644 index b017515..0000000 --- a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/symmetry.py +++ /dev/null @@ -1,24 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -from __future__ import annotations - - -def resolve_symmetry_config(alg_cfg, env): - """Resolve the symmetry configuration. - - Args: - alg_cfg: The algorithm configuration dictionary. - env: The environment. - - Returns: - The resolved algorithm configuration dictionary. - """ - - # if using symmetry then pass the environment config object - if "symmetry_cfg" in alg_cfg and alg_cfg["symmetry_cfg"] is not None: - # this is used by the symmetry function for handling different observation terms - alg_cfg["symmetry_cfg"]["_env"] = env - return alg_cfg diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/networks/__init__.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/networks/__init__.py deleted file mode 100644 index c18f487..0000000 --- a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/networks/__init__.py +++ /dev/null @@ -1,10 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -"""Definitions for components of modules.""" - -from .memory import Memory -from .mlp import MLP -from .normalization import EmpiricalDiscountedVariationNormalization, EmpiricalNormalization diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/networks/memory.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/networks/memory.py deleted file mode 100644 index 7577357..0000000 --- a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/networks/memory.py +++ /dev/null @@ -1,70 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -from __future__ import annotations - -import torch.nn as nn - -from rsl_rl.utils import unpad_trajectories - - -class Memory(nn.Module): - """Memory module for recurrent networks. - - This module is used to store the hidden states of the policy. - Currently only supports GRU and LSTM. - """ - - def __init__(self, input_size, type="lstm", num_layers=1, hidden_size=256): - super().__init__() - # RNN - rnn_cls = nn.GRU if type.lower() == "gru" else nn.LSTM - self.rnn = rnn_cls(input_size=input_size, hidden_size=hidden_size, num_layers=num_layers) - self.hidden_states = None - - def forward(self, input, masks=None, hidden_states=None): - batch_mode = masks is not None - if batch_mode: - # batch mode: needs saved hidden states - if hidden_states is None: - raise ValueError("Hidden states not passed to memory module during policy update") - out, _ = self.rnn(input, hidden_states) - out = unpad_trajectories(out, masks) - else: - # inference/distillation mode: uses hidden states of last step - out, self.hidden_states = self.rnn(input.unsqueeze(0), self.hidden_states) - return out - - def reset(self, dones=None, hidden_states=None): - if dones is None: # reset all hidden states - if hidden_states is None: - self.hidden_states = None - else: - self.hidden_states = hidden_states - elif self.hidden_states is not None: # reset hidden states of done environments - if hidden_states is None: - if isinstance(self.hidden_states, tuple): # tuple in case of LSTM - for hidden_state in self.hidden_states: - hidden_state[..., dones == 1, :] = 0.0 - else: - self.hidden_states[..., dones == 1, :] = 0.0 - else: - NotImplementedError( - "Resetting hidden states of done environments with custom hidden states is not implemented" - ) - - def detach_hidden_states(self, dones=None): - if self.hidden_states is not None: - if dones is None: # detach all hidden states - if isinstance(self.hidden_states, tuple): # tuple in case of LSTM - self.hidden_states = tuple(hidden_state.detach() for hidden_state in self.hidden_states) - else: - self.hidden_states = self.hidden_states.detach() - else: # detach hidden states of done environments - if isinstance(self.hidden_states, tuple): # tuple in case of LSTM - for hidden_state in self.hidden_states: - hidden_state[..., dones == 1, :] = hidden_state[..., dones == 1, :].detach() - else: - self.hidden_states[..., dones == 1, :] = self.hidden_states[..., dones == 1, :].detach() diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/networks/mlp.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/networks/mlp.py deleted file mode 100644 index e91574e..0000000 --- a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/networks/mlp.py +++ /dev/null @@ -1,120 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -from __future__ import annotations - -import torch -import torch.nn as nn -from functools import reduce - -from rsl_rl.utils import resolve_nn_activation - - -class MLP(nn.Sequential): - """Multi-layer perceptron. - - The MLP network is a sequence of linear layers and activation functions. The - last layer is a linear layer that outputs the desired dimension unless the - last activation function is specified. - - It provides additional conveniences: - - - If the hidden dimensions have a value of ``-1``, the dimension is inferred - from the input dimension. - - If the output dimension is a tuple, the output is reshaped to the desired - shape. - - """ - - def __init__( - self, - input_dim: int, - output_dim: int | tuple[int] | list[int], - hidden_dims: tuple[int] | list[int], - activation: str = "elu", - last_activation: str | None = None, - ): - """Initialize the MLP. - - Args: - input_dim: Dimension of the input. - output_dim: Dimension of the output. - hidden_dims: Dimensions of the hidden layers. A value of ``-1`` indicates - that the dimension should be inferred from the input dimension. - activation: Activation function. Defaults to "elu". - last_activation: Activation function of the last layer. Defaults to None, - in which case the last layer is linear. - """ - super().__init__() - - # resolve activation functions - activation_mod = resolve_nn_activation(activation) - last_activation_mod = resolve_nn_activation(last_activation) if last_activation is not None else None - # resolve number of hidden dims if they are -1 - hidden_dims_processed = [input_dim if dim == -1 else dim for dim in hidden_dims] - - # create layers sequentially - layers = [] - layers.append(nn.Linear(input_dim, hidden_dims_processed[0])) - layers.append(activation_mod) - - for layer_index in range(len(hidden_dims_processed) - 1): - layers.append(nn.Linear(hidden_dims_processed[layer_index], hidden_dims_processed[layer_index + 1])) - layers.append(activation_mod) - - # add last layer - if isinstance(output_dim, int): - layers.append(nn.Linear(hidden_dims_processed[-1], output_dim)) - else: - # compute the total output dimension - total_out_dim = reduce(lambda x, y: x * y, output_dim) - # add a layer to reshape the output to the desired shape - layers.append(nn.Linear(hidden_dims_processed[-1], total_out_dim)) - layers.append(nn.Unflatten(dim=-1, unflattened_size=output_dim)) - - # add last activation function if specified - if last_activation_mod is not None: - layers.append(last_activation_mod) - - # register the layers - for idx, layer in enumerate(layers): - self.add_module(f"{idx}", layer) - - def init_weights(self, scales: float | tuple[float]): - """Initialize the weights of the MLP. - - Args: - scales: Scale factor for the weights. - """ - - def get_scale(idx) -> float: - """Get the scale factor for the weights of the MLP. - - Args: - idx: Index of the layer. - """ - return scales[idx] if isinstance(scales, (list, tuple)) else scales - - # initialize the weights - for idx, module in enumerate(self): - if isinstance(module, nn.Linear): - nn.init.orthogonal_(module.weight, gain=get_scale(idx)) - nn.init.zeros_(module.bias) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - """Forward pass of the MLP. - - Args: - x: Input tensor. - """ - for layer in self: - x = layer(x) - return x - - def reset(self, dones=None, hidden_states=None): - pass - - def detach_hidden_states(self, dones=None): - pass diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/networks/normalization.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/networks/normalization.py deleted file mode 100644 index 5fd9692..0000000 --- a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/networks/normalization.py +++ /dev/null @@ -1,130 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -# Copyright (c) 2020 Preferred Networks, Inc. - -from __future__ import annotations - -import torch -from torch import nn - - -class EmpiricalNormalization(nn.Module): - """Normalize mean and variance of values based on empirical values.""" - - def __init__(self, shape, eps=1e-2, until=None): - """Initialize EmpiricalNormalization module. - - Args: - shape (int or tuple of int): Shape of input values except batch axis. - eps (float): Small value for stability. - until (int or None): If this arg is specified, the module learns input values until the sum of batch sizes - exceeds it. - - Note: The normalization parameters are computed over the whole batch, not for each environment separately. - """ - super().__init__() - self.eps = eps - self.until = until - self.register_buffer("_mean", torch.zeros(shape).unsqueeze(0)) - self.register_buffer("_var", torch.ones(shape).unsqueeze(0)) - self.register_buffer("_std", torch.ones(shape).unsqueeze(0)) - self.register_buffer("count", torch.tensor(0, dtype=torch.long)) - - @property - def mean(self): - return self._mean.squeeze(0).clone() - - @property - def std(self): - return self._std.squeeze(0).clone() - - def forward(self, x): - """Normalize mean and variance of values based on empirical values.""" - - return (x - self._mean) / (self._std + self.eps) - - @torch.jit.unused - def update(self, x): - """Learn input values without computing the output values of them""" - - if not self.training: - return - if self.until is not None and self.count >= self.until: - return - - count_x = x.shape[0] - self.count += count_x - rate = count_x / self.count - var_x = torch.var(x, dim=0, unbiased=False, keepdim=True) - mean_x = torch.mean(x, dim=0, keepdim=True) - delta_mean = mean_x - self._mean - self._mean += rate * delta_mean - self._var += rate * (var_x - self._var + delta_mean * (mean_x - self._mean)) - self._std = torch.sqrt(self._var) - - @torch.jit.unused - def inverse(self, y): - """De-normalize values based on empirical values.""" - - return y * (self._std + self.eps) + self._mean - - -class EmpiricalDiscountedVariationNormalization(nn.Module): - """Reward normalization from Pathak's large scale study on PPO. - - Reward normalization. Since the reward function is non-stationary, it is useful to normalize - the scale of the rewards so that the value function can learn quickly. We did this by dividing - the rewards by a running estimate of the standard deviation of the sum of discounted rewards. - """ - - def __init__(self, shape, eps=1e-2, gamma=0.99, until=None): - super().__init__() - - self.emp_norm = EmpiricalNormalization(shape, eps, until) - self.disc_avg = _DiscountedAverage(gamma) - - def forward(self, rew): - if self.training: - # update discounted rewards - avg = self.disc_avg.update(rew) - # update moments from discounted rewards - self.emp_norm.update(avg) - - # normalize rewards with the empirical std - if self.emp_norm._std > 0: - return rew / self.emp_norm._std - else: - return rew - - -""" -Helper class. -""" - - -class _DiscountedAverage: - r"""Discounted average of rewards. - - The discounted average is defined as: - - .. math:: - - \bar{R}_t = \gamma \bar{R}_{t-1} + r_t - - Args: - gamma (float): Discount factor. - """ - - def __init__(self, gamma): - self.avg = None - self.gamma = gamma - - def update(self, rew: torch.Tensor) -> torch.Tensor: - if self.avg is None: - self.avg = rew - else: - self.avg = self.avg * self.gamma + rew - return self.avg diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/runners/__init__.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/runners/__init__.py deleted file mode 100644 index 61f1682..0000000 --- a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/runners/__init__.py +++ /dev/null @@ -1,12 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -"""Implementation of runners for environment-agent interaction.""" - -from .on_policy_runner import OnPolicyRunner # isort:skip -from .distillation_runner import DistillationRunner -from .amp_on_policy_runner import AMPOnPolicyRunner - -__all__ = ["OnPolicyRunner", "DistillationRunner", "AMPOnPolicyRunner"] diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/runners/amp_on_policy_runner.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/runners/amp_on_policy_runner.py deleted file mode 100644 index c0b9b9e..0000000 --- a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/runners/amp_on_policy_runner.py +++ /dev/null @@ -1,521 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -from __future__ import annotations - -import os -import statistics -import time -import torch -import warnings -from collections import deque - -import rsl_rl -from rsl_rl.algorithms import AMP_PPO -from rsl_rl.env import VecEnv -from rsl_rl.modules import ActorCritic, ActorCriticRecurrent,DiscriminatorMulti, resolve_rnd_config, resolve_symmetry_config -from rsl_rl.utils import resolve_obs_groups, store_code_state, Normalizer, G1_AMPLoader - - -class AMPOnPolicyRunner: - """On-policy runner for training and evaluation of actor-critic methods.""" - - def __init__(self, env: VecEnv, train_cfg: dict, log_dir: str | None = None, device="cpu"): - self.cfg = train_cfg - self.alg_cfg = train_cfg["algorithm"] - self.policy_cfg = train_cfg["policy"] - self.device = device - self.env = env - - # check if multi-gpu is enabled - self._configure_multi_gpu() - - # store training configuration - self.num_steps_per_env = self.cfg["num_steps_per_env"] - self.save_interval = self.cfg["save_interval"] - - # query observations from environment for algorithm construction - obs = self.env.get_observations() - default_sets = ["critic"] - if "rnd_cfg" in self.alg_cfg and self.alg_cfg["rnd_cfg"] is not None: - default_sets.append("rnd_state") - self.cfg["obs_groups"] = resolve_obs_groups(obs, self.cfg["obs_groups"], default_sets) - - self.amp_data = G1_AMPLoader( - device, - time_between_frames=1/50.0, - preload_transitions=True, - num_preload_transitions=train_cfg["amp_num_preload_transitions"], - motion_files=train_cfg["amp_motion_files"], - num_frames=train_cfg['amp_num_frames'] - ) - - self.amp_observation_dim = self.amp_data.observation_dim if self.cfg["amp_num_obs"] == 0 else self.cfg["amp_num_obs"] - self.amp_num_frames = 0 if self.cfg["amp_num_frames"] == 0 else self.cfg["amp_num_frames"] - self.amp_normalizer = Normalizer(self.amp_observation_dim) - self.discriminator = DiscriminatorMulti( - self.amp_observation_dim, - train_cfg["amp_reward_coef"], - train_cfg["amp_discr_hidden_dims"], - device, - train_cfg["amp_num_frames"], - train_cfg["amp_task_reward_lerp"], - train_cfg['use_lerp'], - ).to(self.device) - - - # create the algorithm - self.alg = self._construct_algorithm(obs) - - # Decide whether to disable logging - # We only log from the process with rank 0 (main process) - self.disable_logs = self.is_distributed and self.gpu_global_rank != 0 - - # Logging - self.log_dir = log_dir - self.writer = None - self.tot_timesteps = 0 - self.tot_time = 0 - self.current_learning_iteration = 0 - self.git_status_repos = [rsl_rl.__file__] - - def learn(self, num_learning_iterations: int, init_at_random_ep_len: bool = False): # noqa: C901 - # initialize writer - self._prepare_logging_writer() - - # randomize initial episode lengths (for exploration) - if init_at_random_ep_len: - self.env.episode_length_buf = torch.randint_like( - self.env.episode_length_buf, high=int(self.env.max_episode_length) - ) - - # start learning - obs = self.env.get_observations().to(self.device) - self.train_mode() # switch to train mode (for dropout for example) - amp_obs = self.env.get_amp_observations() - amp_obs = amp_obs.to(self.device) - if self.amp_num_frames != 0: - self.amp_obs_frames = torch.zeros(size=(self.env.num_envs, self.amp_num_frames, self.amp_observation_dim), device=self.device) - self.amp_obs_frames = torch.concat((self.amp_obs_frames[:, 1:], amp_obs.unsqueeze(1)), dim=1) - - # Book keeping - ep_infos = [] - rewbuffer = deque(maxlen=100) - lenbuffer = deque(maxlen=100) - step_discrewbuffer = deque(maxlen=100) - - cur_reward_sum = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) - cur_episode_length = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) - cur_single_step_disc_rew = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) - # create buffers for logging extrinsic and intrinsic rewards - if self.alg.rnd: - erewbuffer = deque(maxlen=100) - irewbuffer = deque(maxlen=100) - cur_ereward_sum = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) - cur_ireward_sum = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) - - # Ensure all parameters are in-synced - if self.is_distributed: - print(f"Synchronizing parameters for rank {self.gpu_global_rank}...") - self.alg.broadcast_parameters() - - # Start training - start_iter = self.current_learning_iteration - tot_iter = start_iter + num_learning_iterations - for it in range(start_iter, tot_iter): - start = time.time() - # Rollout - with torch.inference_mode(): - for _ in range(self.num_steps_per_env): - # Sample actions - actions = self.alg.act(obs,amp_obs) - # Step the environment - obs, rewards, dones, extras = self.env.step(actions.to(self.env.device)) - # Move to device - obs, rewards, dones = (obs.to(self.device), rewards.to(self.device), dones.to(self.device)) - - next_amp_obs = self.env.get_amp_observations() - next_amp_obs = next_amp_obs.to(self.device) - next_amp_obs_with_term = torch.clone(next_amp_obs) - - reset_env_ids = self.env.reset_env_ids - terminal_amp_states = self.env.get_amp_observations()[reset_env_ids] - next_amp_obs_with_term[reset_env_ids] = terminal_amp_states - self.amp_obs_frames = torch.concat((self.amp_obs_frames[:, 1:], next_amp_obs_with_term.unsqueeze(1)), dim=1) - - amp_reward = torch.zeros(self.env.num_envs, device=obs.device) - - mask = self.env.contact_phase[:, 0] == 1.0 - if mask.any(): - rewards[mask], logit, disc_reward = self.alg.discriminator.predict_amp_reward( - self.amp_obs_frames[mask], rewards[mask], normalizer=self.alg.amp_normalizer - ) - amp_reward[mask] += disc_reward - - # process the step - self.alg.process_env_step(obs, rewards, dones, extras, next_amp_obs_with_term, self.amp_obs_frames) - self.amp_obs_frames[reset_env_ids] = 0 - - amp_obs = torch.clone(next_amp_obs) - # Extract intrinsic rewards (only for logging) - intrinsic_rewards = self.alg.intrinsic_rewards if self.alg.rnd else None - # book keeping - if self.log_dir is not None: - if "episode" in extras: - ep_infos.append(extras["episode"]) - elif "log" in extras: - ep_infos.append(extras["log"]) - # Update rewards - if self.alg.rnd: - cur_ereward_sum += rewards - cur_ireward_sum += intrinsic_rewards # type: ignore - cur_reward_sum += rewards + intrinsic_rewards - else: - cur_reward_sum += rewards - # Update episode length - cur_episode_length += 1 - cur_single_step_disc_rew += amp_reward - # Clear data for completed episodes - # -- common - new_ids = (dones > 0).nonzero(as_tuple=False) - rewbuffer.extend(cur_reward_sum[new_ids][:, 0].cpu().numpy().tolist()) - lenbuffer.extend(cur_episode_length[new_ids][:, 0].cpu().numpy().tolist()) - cur_reward_sum[new_ids] = 0 - cur_episode_length[new_ids] = 0 - to_extend_disc = (cur_single_step_disc_rew[new_ids] / self.env.max_episode_length_s)[:, 0].cpu().numpy() - step_discrewbuffer.extend(to_extend_disc.tolist()) - cur_single_step_disc_rew[new_ids] = 0 - # -- intrinsic and extrinsic rewards - if self.alg.rnd: - erewbuffer.extend(cur_ereward_sum[new_ids][:, 0].cpu().numpy().tolist()) - irewbuffer.extend(cur_ireward_sum[new_ids][:, 0].cpu().numpy().tolist()) - cur_ereward_sum[new_ids] = 0 - cur_ireward_sum[new_ids] = 0 - - stop = time.time() - collection_time = stop - start - start = stop - - # compute returns - self.alg.compute_returns(obs) - - # update policy - loss_dict = self.alg.update() - - stop = time.time() - learn_time = stop - start - self.current_learning_iteration = it - # log info - if self.log_dir is not None and not self.disable_logs: - # Log information - self.log(locals()) - # Save model - if it % self.save_interval == 0: - self.save(os.path.join(self.log_dir, f"model_{it}.pt")) - - # Clear episode infos - ep_infos.clear() - # Save code state - if it == start_iter and not self.disable_logs: - # obtain all the diff files - git_file_paths = store_code_state(self.log_dir, self.git_status_repos) - # if possible store them to wandb - if self.logger_type in ["wandb", "neptune"] and git_file_paths: - for path in git_file_paths: - self.writer.save_file(path) - - # Save the final model after training - if self.log_dir is not None and not self.disable_logs: - self.save(os.path.join(self.log_dir, f"model_{self.current_learning_iteration}.pt")) - - def log(self, locs: dict, width: int = 80, pad: int = 35): - # Compute the collection size - collection_size = self.num_steps_per_env * self.env.num_envs * self.gpu_world_size - # Update total time-steps and time - self.tot_timesteps += collection_size - self.tot_time += locs["collection_time"] + locs["learn_time"] - iteration_time = locs["collection_time"] + locs["learn_time"] - - # -- Episode info - ep_string = "" - if locs["ep_infos"]: - for key in locs["ep_infos"][0]: - infotensor = torch.tensor([], device=self.device) - for ep_info in locs["ep_infos"]: - # handle scalar and zero dimensional tensor infos - if key not in ep_info: - continue - if not isinstance(ep_info[key], torch.Tensor): - ep_info[key] = torch.Tensor([ep_info[key]]) - if len(ep_info[key].shape) == 0: - ep_info[key] = ep_info[key].unsqueeze(0) - infotensor = torch.cat((infotensor, ep_info[key].to(self.device))) - value = torch.mean(infotensor) - # log to logger and terminal - if "/" in key: - self.writer.add_scalar(key, value, locs["it"]) - ep_string += f"""{f'{key}:':>{pad}} {value:.4f}\n""" - else: - self.writer.add_scalar("Episode/" + key, value, locs["it"]) - ep_string += f"""{f'Mean episode {key}:':>{pad}} {value:.4f}\n""" - - mean_std = self.alg.policy.action_std.mean() - fps = int(collection_size / (locs["collection_time"] + locs["learn_time"])) - - # -- Losses - for key, value in locs["loss_dict"].items(): - self.writer.add_scalar(f"Loss/{key}", value, locs["it"]) - self.writer.add_scalar("Loss/learning_rate", self.alg.learning_rate, locs["it"]) - - # -- Policy - self.writer.add_scalar("Policy/mean_noise_std", mean_std.item(), locs["it"]) - - # -- Performance - self.writer.add_scalar("Perf/total_fps", fps, locs["it"]) - self.writer.add_scalar("Perf/collection time", locs["collection_time"], locs["it"]) - self.writer.add_scalar("Perf/learning_time", locs["learn_time"], locs["it"]) - - # -- Training - if len(locs["rewbuffer"]) > 0: - # separate logging for intrinsic and extrinsic rewards - if hasattr(self.alg, "rnd") and self.alg.rnd: - self.writer.add_scalar("Rnd/mean_extrinsic_reward", statistics.mean(locs["erewbuffer"]), locs["it"]) - self.writer.add_scalar("Rnd/mean_intrinsic_reward", statistics.mean(locs["irewbuffer"]), locs["it"]) - self.writer.add_scalar("Rnd/weight", self.alg.rnd.weight, locs["it"]) - # everything else - self.writer.add_scalar("Train/mean_reward", statistics.mean(locs["rewbuffer"]), locs["it"]) - self.writer.add_scalar("Train/mean_episode_length", statistics.mean(locs["lenbuffer"]), locs["it"]) - self.writer.add_scalar('Train/mean_step_disc_reward', statistics.mean(locs['step_discrewbuffer']), locs['it']) - if self.logger_type != "wandb": # wandb does not support non-integer x-axis logging - self.writer.add_scalar("Train/mean_reward/time", statistics.mean(locs["rewbuffer"]), self.tot_time) - self.writer.add_scalar( - "Train/mean_episode_length/time", statistics.mean(locs["lenbuffer"]), self.tot_time - ) - - str = f" \033[1m Learning iteration {locs['it']}/{locs['tot_iter']} \033[0m " - - if len(locs["rewbuffer"]) > 0: - log_string = ( - f"""{'#' * width}\n""" - f"""{str.center(width, ' ')}\n\n""" - f"""{'Computation:':>{pad}} {fps:.0f} steps/s (collection: {locs[ - 'collection_time']:.3f}s, learning {locs['learn_time']:.3f}s)\n""" - f"""{'Mean action noise std:':>{pad}} {mean_std.item():.2f}\n""" - f"""{'Step disc reward:':>{pad}} {statistics.mean(locs['step_discrewbuffer']):.2f}\n""" - ) - # -- Losses - for key, value in locs["loss_dict"].items(): - log_string += f"""{f'Mean {key} loss:':>{pad}} {value:.4f}\n""" - # -- Rewards - if hasattr(self.alg, "rnd") and self.alg.rnd: - log_string += ( - f"""{'Mean extrinsic reward:':>{pad}} {statistics.mean(locs['erewbuffer']):.2f}\n""" - f"""{'Mean intrinsic reward:':>{pad}} {statistics.mean(locs['irewbuffer']):.2f}\n""" - ) - log_string += f"""{'Mean reward:':>{pad}} {statistics.mean(locs['rewbuffer']):.2f}\n""" - # -- episode info - log_string += f"""{'Mean episode length:':>{pad}} {statistics.mean(locs['lenbuffer']):.2f}\n""" - else: - log_string = ( - f"""{'#' * width}\n""" - f"""{str.center(width, ' ')}\n\n""" - f"""{'Computation:':>{pad}} {fps:.0f} steps/s (collection: {locs[ - 'collection_time']:.3f}s, learning {locs['learn_time']:.3f}s)\n""" - f"""{'Mean action noise std:':>{pad}} {mean_std.item():.2f}\n""" - ) - for key, value in locs["loss_dict"].items(): - log_string += f"""{f'{key}:':>{pad}} {value:.4f}\n""" - - log_string += ep_string - log_string += ( - f"""{'-' * width}\n""" - f"""{'Total timesteps:':>{pad}} {self.tot_timesteps}\n""" - f"""{'Iteration time:':>{pad}} {iteration_time:.2f}s\n""" - f"""{'Time elapsed:':>{pad}} {time.strftime("%H:%M:%S", time.gmtime(self.tot_time))}\n""" - f"""{'ETA:':>{pad}} {time.strftime( - "%H:%M:%S", - time.gmtime( - self.tot_time / (locs['it'] - locs['start_iter'] + 1) - * (locs['start_iter'] + locs['num_learning_iterations'] - locs['it']) - ) - )}\n""" - ) - print(log_string) - - def save(self, path: str, infos=None): - # -- Save model - saved_dict = { - "model_state_dict": self.alg.policy.state_dict(), - "optimizer_state_dict": self.alg.optimizer.state_dict(), - "iter": self.current_learning_iteration, - "infos": infos, - } - # -- Save RND model if used - if hasattr(self.alg, "rnd") and self.alg.rnd: - saved_dict["rnd_state_dict"] = self.alg.rnd.state_dict() - saved_dict["rnd_optimizer_state_dict"] = self.alg.rnd_optimizer.state_dict() - torch.save(saved_dict, path) - - # upload model to external logging service - if self.logger_type in ["neptune", "wandb"] and not self.disable_logs: - self.writer.save_model(path, self.current_learning_iteration) - - def load(self, path: str, load_optimizer: bool = True, map_location: str | None = None): - loaded_dict = torch.load(path, weights_only=False, map_location=map_location) - # -- Load model - resumed_training = self.alg.policy.load_state_dict(loaded_dict["model_state_dict"]) - # -- Load RND model if used - if hasattr(self.alg, "rnd") and self.alg.rnd: - self.alg.rnd.load_state_dict(loaded_dict["rnd_state_dict"]) - # -- load optimizer if used - if load_optimizer and resumed_training: - # -- algorithm optimizer - self.alg.optimizer.load_state_dict(loaded_dict["optimizer_state_dict"]) - # -- RND optimizer if used - if hasattr(self.alg, "rnd") and self.alg.rnd: - self.alg.rnd_optimizer.load_state_dict(loaded_dict["rnd_optimizer_state_dict"]) - # -- load current learning iteration - if resumed_training: - self.current_learning_iteration = loaded_dict["iter"] - return loaded_dict["infos"] - - def get_inference_policy(self, device=None): - self.eval_mode() # switch to evaluation mode (dropout for example) - if device is not None: - self.alg.policy.to(device) - return self.alg.policy.act_inference - - def train_mode(self): - # -- PPO - self.alg.policy.train() - self.alg.discriminator.train() - # -- RND - if hasattr(self.alg, "rnd") and self.alg.rnd: - self.alg.rnd.train() - - def eval_mode(self): - # -- PPO - self.alg.policy.eval() - self.alg.discriminator.eval() - # -- RND - if hasattr(self.alg, "rnd") and self.alg.rnd: - self.alg.rnd.eval() - - def add_git_repo_to_log(self, repo_file_path): - self.git_status_repos.append(repo_file_path) - - """ - Helper functions. - """ - - def _configure_multi_gpu(self): - """Configure multi-gpu training.""" - # check if distributed training is enabled - self.gpu_world_size = int(os.getenv("WORLD_SIZE", "1")) - self.is_distributed = self.gpu_world_size > 1 - - # if not distributed training, set local and global rank to 0 and return - if not self.is_distributed: - self.gpu_local_rank = 0 - self.gpu_global_rank = 0 - self.multi_gpu_cfg = None - return - - # get rank and world size - self.gpu_local_rank = int(os.getenv("LOCAL_RANK", "0")) - self.gpu_global_rank = int(os.getenv("RANK", "0")) - - # make a configuration dictionary - self.multi_gpu_cfg = { - "global_rank": self.gpu_global_rank, # rank of the main process - "local_rank": self.gpu_local_rank, # rank of the current process - "world_size": self.gpu_world_size, # total number of processes - } - - # check if user has device specified for local rank - if self.device != f"cuda:{self.gpu_local_rank}": - raise ValueError( - f"Device '{self.device}' does not match expected device for local rank '{self.gpu_local_rank}'." - ) - # validate multi-gpu configuration - if self.gpu_local_rank >= self.gpu_world_size: - raise ValueError( - f"Local rank '{self.gpu_local_rank}' is greater than or equal to world size '{self.gpu_world_size}'." - ) - if self.gpu_global_rank >= self.gpu_world_size: - raise ValueError( - f"Global rank '{self.gpu_global_rank}' is greater than or equal to world size '{self.gpu_world_size}'." - ) - - # initialize torch distributed - torch.distributed.init_process_group(backend="nccl", rank=self.gpu_global_rank, world_size=self.gpu_world_size) - # set device to the local rank - torch.cuda.set_device(self.gpu_local_rank) - - def _construct_algorithm(self, obs) -> AMP_PPO: - """Construct the actor-critic algorithm.""" - # resolve RND config - self.alg_cfg = resolve_rnd_config(self.alg_cfg, obs, self.cfg["obs_groups"], self.env) - - # resolve symmetry config - self.alg_cfg = resolve_symmetry_config(self.alg_cfg, self.env) - - # resolve deprecated normalization config - if self.cfg.get("empirical_normalization") is not None: - warnings.warn( - "The `empirical_normalization` parameter is deprecated. Please set `actor_obs_normalization` and " - "`critic_obs_normalization` as part of the `policy` configuration instead.", - DeprecationWarning, - ) - if self.policy_cfg.get("actor_obs_normalization") is None: - self.policy_cfg["actor_obs_normalization"] = self.cfg["empirical_normalization"] - if self.policy_cfg.get("critic_obs_normalization") is None: - self.policy_cfg["critic_obs_normalization"] = self.cfg["empirical_normalization"] - - # initialize the actor-critic - actor_critic_class = eval(self.policy_cfg.pop("class_name")) - actor_critic: ActorCritic | ActorCriticRecurrent = actor_critic_class( - obs, self.cfg["obs_groups"], self.env.num_actions, **self.policy_cfg - ).to(self.device) - - # initialize the algorithm - alg_class = eval(self.alg_cfg.pop("class_name")) - - alg: AMP_PPO = alg_class(actor_critic, self.discriminator, self.amp_data, self.amp_normalizer, self.amp_num_frames, device=self.device, **self.alg_cfg, multi_gpu_cfg=self.multi_gpu_cfg) - - # initialize the storage - alg.init_storage( - "rl", - self.env.num_envs, - self.num_steps_per_env, - obs, - [self.env.num_actions], - ) - - return alg - - def _prepare_logging_writer(self): - """Prepares the logging writers.""" - if self.log_dir is not None and self.writer is None and not self.disable_logs: - # Launch either Tensorboard or Neptune & Tensorboard summary writer(s), default: Tensorboard. - self.logger_type = self.cfg.get("logger", "tensorboard") - self.logger_type = self.logger_type.lower() - - if self.logger_type == "neptune": - from rsl_rl.utils.neptune_utils import NeptuneSummaryWriter - - self.writer = NeptuneSummaryWriter(log_dir=self.log_dir, flush_secs=10, cfg=self.cfg) - self.writer.log_config(self.env.cfg, self.cfg, self.alg_cfg, self.policy_cfg) - elif self.logger_type == "wandb": - from rsl_rl.utils.wandb_utils import WandbSummaryWriter - - self.writer = WandbSummaryWriter(log_dir=self.log_dir, flush_secs=10, cfg=self.cfg) - self.writer.log_config(self.env.cfg, self.cfg, self.alg_cfg, self.policy_cfg) - elif self.logger_type == "tensorboard": - from torch.utils.tensorboard import SummaryWriter - - self.writer = SummaryWriter(log_dir=self.log_dir, flush_secs=10) - else: - raise ValueError("Logger type not found. Please choose 'neptune', 'wandb' or 'tensorboard'.") diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/runners/distillation_runner.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/runners/distillation_runner.py deleted file mode 100644 index 9cc6a8b..0000000 --- a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/runners/distillation_runner.py +++ /dev/null @@ -1,179 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -from __future__ import annotations - -import os -import time -import torch -from collections import deque - -import rsl_rl -from rsl_rl.algorithms import Distillation -from rsl_rl.env import VecEnv -from rsl_rl.modules import StudentTeacher, StudentTeacherRecurrent -from rsl_rl.runners import OnPolicyRunner -from rsl_rl.utils import resolve_obs_groups, store_code_state - - -class DistillationRunner(OnPolicyRunner): - """On-policy runner for training and evaluation of teacher-student training.""" - - def __init__(self, env: VecEnv, train_cfg: dict, log_dir: str | None = None, device="cpu"): - self.cfg = train_cfg - self.alg_cfg = train_cfg["algorithm"] - self.policy_cfg = train_cfg["policy"] - self.device = device - self.env = env - - # check if multi-gpu is enabled - self._configure_multi_gpu() - - # store training configuration - self.num_steps_per_env = self.cfg["num_steps_per_env"] - self.save_interval = self.cfg["save_interval"] - - # query observations from environment for algorithm construction - obs = self.env.get_observations() - self.cfg["obs_groups"] = resolve_obs_groups(obs, self.cfg["obs_groups"], default_sets=["teacher"]) - - # create the algorithm - self.alg = self._construct_algorithm(obs) - - # Decide whether to disable logging - # We only log from the process with rank 0 (main process) - self.disable_logs = self.is_distributed and self.gpu_global_rank != 0 - - # Logging - self.log_dir = log_dir - self.writer = None - self.tot_timesteps = 0 - self.tot_time = 0 - self.current_learning_iteration = 0 - self.git_status_repos = [rsl_rl.__file__] - - def learn(self, num_learning_iterations: int, init_at_random_ep_len: bool = False): # noqa: C901 - # initialize writer - self._prepare_logging_writer() - # check if teacher is loaded - if not self.alg.policy.loaded_teacher: - raise ValueError("Teacher model parameters not loaded. Please load a teacher model to distill.") - - # randomize initial episode lengths (for exploration) - if init_at_random_ep_len: - self.env.episode_length_buf = torch.randint_like( - self.env.episode_length_buf, high=int(self.env.max_episode_length) - ) - - # start learning - obs = self.env.get_observations().to(self.device) - self.train_mode() # switch to train mode (for dropout for example) - - # Book keeping - ep_infos = [] - rewbuffer = deque(maxlen=100) - lenbuffer = deque(maxlen=100) - cur_reward_sum = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) - cur_episode_length = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) - - # Ensure all parameters are in-synced - if self.is_distributed: - print(f"Synchronizing parameters for rank {self.gpu_global_rank}...") - self.alg.broadcast_parameters() - - # Start training - start_iter = self.current_learning_iteration - tot_iter = start_iter + num_learning_iterations - for it in range(start_iter, tot_iter): - start = time.time() - # Rollout - with torch.inference_mode(): - for _ in range(self.num_steps_per_env): - # Sample actions - actions = self.alg.act(obs) - # Step the environment - obs, rewards, dones, extras = self.env.step(actions.to(self.env.device)) - # Move to device - obs, rewards, dones = (obs.to(self.device), rewards.to(self.device), dones.to(self.device)) - # process the step - self.alg.process_env_step(obs, rewards, dones, extras) - # book keeping - if self.log_dir is not None: - if "episode" in extras: - ep_infos.append(extras["episode"]) - elif "log" in extras: - ep_infos.append(extras["log"]) - # Update rewards - cur_reward_sum += rewards - # Update episode length - cur_episode_length += 1 - # Clear data for completed episodes - new_ids = (dones > 0).nonzero(as_tuple=False) - rewbuffer.extend(cur_reward_sum[new_ids][:, 0].cpu().numpy().tolist()) - lenbuffer.extend(cur_episode_length[new_ids][:, 0].cpu().numpy().tolist()) - cur_reward_sum[new_ids] = 0 - cur_episode_length[new_ids] = 0 - - stop = time.time() - collection_time = stop - start - start = stop - - # update policy - loss_dict = self.alg.update() - - stop = time.time() - learn_time = stop - start - self.current_learning_iteration = it - # log info - if self.log_dir is not None and not self.disable_logs: - # Log information - self.log(locals()) - # Save model - if it % self.save_interval == 0: - self.save(os.path.join(self.log_dir, f"model_{it}.pt")) - - # Clear episode infos - ep_infos.clear() - # Save code state - if it == start_iter and not self.disable_logs: - # obtain all the diff files - git_file_paths = store_code_state(self.log_dir, self.git_status_repos) - # if possible store them to wandb - if self.logger_type in ["wandb", "neptune"] and git_file_paths: - for path in git_file_paths: - self.writer.save_file(path) - - # Save the final model after training - if self.log_dir is not None and not self.disable_logs: - self.save(os.path.join(self.log_dir, f"model_{self.current_learning_iteration}.pt")) - - """ - Helper methods. - """ - - def _construct_algorithm(self, obs) -> Distillation: - """Construct the distillation algorithm.""" - # initialize the actor-critic - student_teacher_class = eval(self.policy_cfg.pop("class_name")) - student_teacher: StudentTeacher | StudentTeacherRecurrent = student_teacher_class( - obs, self.cfg["obs_groups"], self.env.num_actions, **self.policy_cfg - ).to(self.device) - - # initialize the algorithm - alg_class = eval(self.alg_cfg.pop("class_name")) - alg: Distillation = alg_class( - student_teacher, device=self.device, **self.alg_cfg, multi_gpu_cfg=self.multi_gpu_cfg - ) - - # initialize the storage - alg.init_storage( - "distillation", - self.env.num_envs, - self.num_steps_per_env, - obs, - [self.env.num_actions], - ) - - return alg diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/runners/on_policy_runner.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/runners/on_policy_runner.py deleted file mode 100644 index 36f11f3..0000000 --- a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/runners/on_policy_runner.py +++ /dev/null @@ -1,460 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -from __future__ import annotations - -import os -import statistics -import time -import torch -import warnings -from collections import deque - -import rsl_rl -from rsl_rl.algorithms import PPO -from rsl_rl.env import VecEnv -from rsl_rl.modules import ActorCritic, ActorCriticRecurrent, resolve_rnd_config, resolve_symmetry_config -from rsl_rl.utils import resolve_obs_groups, store_code_state - - -class OnPolicyRunner: - """On-policy runner for training and evaluation of actor-critic methods.""" - - def __init__(self, env: VecEnv, train_cfg: dict, log_dir: str | None = None, device="cpu"): - self.cfg = train_cfg - self.alg_cfg = train_cfg["algorithm"] - self.policy_cfg = train_cfg["policy"] - self.device = device - self.env = env - - # check if multi-gpu is enabled - self._configure_multi_gpu() - - # store training configuration - self.num_steps_per_env = self.cfg["num_steps_per_env"] - self.save_interval = self.cfg["save_interval"] - - # query observations from environment for algorithm construction - obs = self.env.get_observations() - default_sets = ["critic"] - if "rnd_cfg" in self.alg_cfg and self.alg_cfg["rnd_cfg"] is not None: - default_sets.append("rnd_state") - self.cfg["obs_groups"] = resolve_obs_groups(obs, self.cfg["obs_groups"], default_sets) - - # create the algorithm - self.alg = self._construct_algorithm(obs) - - # Decide whether to disable logging - # We only log from the process with rank 0 (main process) - self.disable_logs = self.is_distributed and self.gpu_global_rank != 0 - - # Logging - self.log_dir = log_dir - self.writer = None - self.tot_timesteps = 0 - self.tot_time = 0 - self.current_learning_iteration = 0 - self.git_status_repos = [rsl_rl.__file__] - - def learn(self, num_learning_iterations: int, init_at_random_ep_len: bool = False): # noqa: C901 - # initialize writer - self._prepare_logging_writer() - - # randomize initial episode lengths (for exploration) - if init_at_random_ep_len: - self.env.episode_length_buf = torch.randint_like( - self.env.episode_length_buf, high=int(self.env.max_episode_length) - ) - - # start learning - obs = self.env.get_observations().to(self.device) - self.train_mode() # switch to train mode (for dropout for example) - - # Book keeping - ep_infos = [] - rewbuffer = deque(maxlen=100) - lenbuffer = deque(maxlen=100) - cur_reward_sum = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) - cur_episode_length = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) - - # create buffers for logging extrinsic and intrinsic rewards - if self.alg.rnd: - erewbuffer = deque(maxlen=100) - irewbuffer = deque(maxlen=100) - cur_ereward_sum = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) - cur_ireward_sum = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) - - # Ensure all parameters are in-synced - if self.is_distributed: - print(f"Synchronizing parameters for rank {self.gpu_global_rank}...") - self.alg.broadcast_parameters() - - # Start training - start_iter = self.current_learning_iteration - tot_iter = start_iter + num_learning_iterations - for it in range(start_iter, tot_iter): - start = time.time() - # Rollout - with torch.inference_mode(): - for _ in range(self.num_steps_per_env): - # Sample actions - actions = self.alg.act(obs) - # Step the environment - obs, rewards, dones, extras = self.env.step(actions.to(self.env.device)) - # Move to device - obs, rewards, dones = (obs.to(self.device), rewards.to(self.device), dones.to(self.device)) - # process the step - self.alg.process_env_step(obs, rewards, dones, extras) - # Extract intrinsic rewards (only for logging) - intrinsic_rewards = self.alg.intrinsic_rewards if self.alg.rnd else None - # book keeping - if self.log_dir is not None: - if "episode" in extras: - ep_infos.append(extras["episode"]) - elif "log" in extras: - ep_infos.append(extras["log"]) - # Update rewards - if self.alg.rnd: - cur_ereward_sum += rewards - cur_ireward_sum += intrinsic_rewards # type: ignore - cur_reward_sum += rewards + intrinsic_rewards - else: - cur_reward_sum += rewards - # Update episode length - cur_episode_length += 1 - # Clear data for completed episodes - # -- common - new_ids = (dones > 0).nonzero(as_tuple=False) - rewbuffer.extend(cur_reward_sum[new_ids][:, 0].cpu().numpy().tolist()) - lenbuffer.extend(cur_episode_length[new_ids][:, 0].cpu().numpy().tolist()) - cur_reward_sum[new_ids] = 0 - cur_episode_length[new_ids] = 0 - # -- intrinsic and extrinsic rewards - if self.alg.rnd: - erewbuffer.extend(cur_ereward_sum[new_ids][:, 0].cpu().numpy().tolist()) - irewbuffer.extend(cur_ireward_sum[new_ids][:, 0].cpu().numpy().tolist()) - cur_ereward_sum[new_ids] = 0 - cur_ireward_sum[new_ids] = 0 - - stop = time.time() - collection_time = stop - start - start = stop - - # compute returns - self.alg.compute_returns(obs) - - # update policy - loss_dict = self.alg.update() - - stop = time.time() - learn_time = stop - start - self.current_learning_iteration = it - # log info - if self.log_dir is not None and not self.disable_logs: - # Log information - self.log(locals()) - # Save model - if it % self.save_interval == 0: - self.save(os.path.join(self.log_dir, f"model_{it}.pt")) - - # Clear episode infos - ep_infos.clear() - # Save code state - if it == start_iter and not self.disable_logs: - # obtain all the diff files - git_file_paths = store_code_state(self.log_dir, self.git_status_repos) - # if possible store them to wandb - if self.logger_type in ["wandb", "neptune"] and git_file_paths: - for path in git_file_paths: - self.writer.save_file(path) - - # Save the final model after training - if self.log_dir is not None and not self.disable_logs: - self.save(os.path.join(self.log_dir, f"model_{self.current_learning_iteration}.pt")) - - def log(self, locs: dict, width: int = 80, pad: int = 35): - # Compute the collection size - collection_size = self.num_steps_per_env * self.env.num_envs * self.gpu_world_size - # Update total time-steps and time - self.tot_timesteps += collection_size - self.tot_time += locs["collection_time"] + locs["learn_time"] - iteration_time = locs["collection_time"] + locs["learn_time"] - - # -- Episode info - ep_string = "" - if locs["ep_infos"]: - for key in locs["ep_infos"][0]: - infotensor = torch.tensor([], device=self.device) - for ep_info in locs["ep_infos"]: - # handle scalar and zero dimensional tensor infos - if key not in ep_info: - continue - if not isinstance(ep_info[key], torch.Tensor): - ep_info[key] = torch.Tensor([ep_info[key]]) - if len(ep_info[key].shape) == 0: - ep_info[key] = ep_info[key].unsqueeze(0) - infotensor = torch.cat((infotensor, ep_info[key].to(self.device))) - value = torch.mean(infotensor) - # log to logger and terminal - if "/" in key: - self.writer.add_scalar(key, value, locs["it"]) - ep_string += f"""{f'{key}:':>{pad}} {value:.4f}\n""" - else: - self.writer.add_scalar("Episode/" + key, value, locs["it"]) - ep_string += f"""{f'Mean episode {key}:':>{pad}} {value:.4f}\n""" - - mean_std = self.alg.policy.action_std.mean() - fps = int(collection_size / (locs["collection_time"] + locs["learn_time"])) - - # -- Losses - for key, value in locs["loss_dict"].items(): - self.writer.add_scalar(f"Loss/{key}", value, locs["it"]) - self.writer.add_scalar("Loss/learning_rate", self.alg.learning_rate, locs["it"]) - - # -- Policy - self.writer.add_scalar("Policy/mean_noise_std", mean_std.item(), locs["it"]) - - # -- Performance - self.writer.add_scalar("Perf/total_fps", fps, locs["it"]) - self.writer.add_scalar("Perf/collection time", locs["collection_time"], locs["it"]) - self.writer.add_scalar("Perf/learning_time", locs["learn_time"], locs["it"]) - - # -- Training - if len(locs["rewbuffer"]) > 0: - # separate logging for intrinsic and extrinsic rewards - if hasattr(self.alg, "rnd") and self.alg.rnd: - self.writer.add_scalar("Rnd/mean_extrinsic_reward", statistics.mean(locs["erewbuffer"]), locs["it"]) - self.writer.add_scalar("Rnd/mean_intrinsic_reward", statistics.mean(locs["irewbuffer"]), locs["it"]) - self.writer.add_scalar("Rnd/weight", self.alg.rnd.weight, locs["it"]) - # everything else - self.writer.add_scalar("Train/mean_reward", statistics.mean(locs["rewbuffer"]), locs["it"]) - self.writer.add_scalar("Train/mean_episode_length", statistics.mean(locs["lenbuffer"]), locs["it"]) - if self.logger_type != "wandb": # wandb does not support non-integer x-axis logging - self.writer.add_scalar("Train/mean_reward/time", statistics.mean(locs["rewbuffer"]), self.tot_time) - self.writer.add_scalar( - "Train/mean_episode_length/time", statistics.mean(locs["lenbuffer"]), self.tot_time - ) - - str = f" \033[1m Learning iteration {locs['it']}/{locs['tot_iter']} \033[0m " - - if len(locs["rewbuffer"]) > 0: - log_string = ( - f"""{'#' * width}\n""" - f"""{str.center(width, ' ')}\n\n""" - f"""{'Computation:':>{pad}} {fps:.0f} steps/s (collection: {locs[ - 'collection_time']:.3f}s, learning {locs['learn_time']:.3f}s)\n""" - f"""{'Mean action noise std:':>{pad}} {mean_std.item():.2f}\n""" - ) - # -- Losses - for key, value in locs["loss_dict"].items(): - log_string += f"""{f'Mean {key} loss:':>{pad}} {value:.4f}\n""" - # -- Rewards - if hasattr(self.alg, "rnd") and self.alg.rnd: - log_string += ( - f"""{'Mean extrinsic reward:':>{pad}} {statistics.mean(locs['erewbuffer']):.2f}\n""" - f"""{'Mean intrinsic reward:':>{pad}} {statistics.mean(locs['irewbuffer']):.2f}\n""" - ) - log_string += f"""{'Mean reward:':>{pad}} {statistics.mean(locs['rewbuffer']):.2f}\n""" - # -- episode info - log_string += f"""{'Mean episode length:':>{pad}} {statistics.mean(locs['lenbuffer']):.2f}\n""" - else: - log_string = ( - f"""{'#' * width}\n""" - f"""{str.center(width, ' ')}\n\n""" - f"""{'Computation:':>{pad}} {fps:.0f} steps/s (collection: {locs[ - 'collection_time']:.3f}s, learning {locs['learn_time']:.3f}s)\n""" - f"""{'Mean action noise std:':>{pad}} {mean_std.item():.2f}\n""" - ) - for key, value in locs["loss_dict"].items(): - log_string += f"""{f'{key}:':>{pad}} {value:.4f}\n""" - - log_string += ep_string - log_string += ( - f"""{'-' * width}\n""" - f"""{'Total timesteps:':>{pad}} {self.tot_timesteps}\n""" - f"""{'Iteration time:':>{pad}} {iteration_time:.2f}s\n""" - f"""{'Time elapsed:':>{pad}} {time.strftime("%H:%M:%S", time.gmtime(self.tot_time))}\n""" - f"""{'ETA:':>{pad}} {time.strftime( - "%H:%M:%S", - time.gmtime( - self.tot_time / (locs['it'] - locs['start_iter'] + 1) - * (locs['start_iter'] + locs['num_learning_iterations'] - locs['it']) - ) - )}\n""" - ) - print(log_string) - - def save(self, path: str, infos=None): - # -- Save model - saved_dict = { - "model_state_dict": self.alg.policy.state_dict(), - "optimizer_state_dict": self.alg.optimizer.state_dict(), - "iter": self.current_learning_iteration, - "infos": infos, - } - # -- Save RND model if used - if hasattr(self.alg, "rnd") and self.alg.rnd: - saved_dict["rnd_state_dict"] = self.alg.rnd.state_dict() - saved_dict["rnd_optimizer_state_dict"] = self.alg.rnd_optimizer.state_dict() - torch.save(saved_dict, path) - - # upload model to external logging service - if self.logger_type in ["neptune", "wandb"] and not self.disable_logs: - self.writer.save_model(path, self.current_learning_iteration) - - def load(self, path: str, load_optimizer: bool = True, map_location: str | None = None): - loaded_dict = torch.load(path, weights_only=False, map_location=map_location) - # -- Load model - resumed_training = self.alg.policy.load_state_dict(loaded_dict["model_state_dict"]) - # -- Load RND model if used - if hasattr(self.alg, "rnd") and self.alg.rnd: - self.alg.rnd.load_state_dict(loaded_dict["rnd_state_dict"]) - # -- load optimizer if used - if load_optimizer and resumed_training: - # -- algorithm optimizer - self.alg.optimizer.load_state_dict(loaded_dict["optimizer_state_dict"]) - # -- RND optimizer if used - if hasattr(self.alg, "rnd") and self.alg.rnd: - self.alg.rnd_optimizer.load_state_dict(loaded_dict["rnd_optimizer_state_dict"]) - # -- load current learning iteration - if resumed_training: - self.current_learning_iteration = loaded_dict["iter"] - return loaded_dict["infos"] - - def get_inference_policy(self, device=None): - self.eval_mode() # switch to evaluation mode (dropout for example) - if device is not None: - self.alg.policy.to(device) - return self.alg.policy.act_inference - - def train_mode(self): - # -- PPO - self.alg.policy.train() - # -- RND - if hasattr(self.alg, "rnd") and self.alg.rnd: - self.alg.rnd.train() - - def eval_mode(self): - # -- PPO - self.alg.policy.eval() - # -- RND - if hasattr(self.alg, "rnd") and self.alg.rnd: - self.alg.rnd.eval() - - def add_git_repo_to_log(self, repo_file_path): - self.git_status_repos.append(repo_file_path) - - """ - Helper functions. - """ - - def _configure_multi_gpu(self): - """Configure multi-gpu training.""" - # check if distributed training is enabled - self.gpu_world_size = int(os.getenv("WORLD_SIZE", "1")) - self.is_distributed = self.gpu_world_size > 1 - - # if not distributed training, set local and global rank to 0 and return - if not self.is_distributed: - self.gpu_local_rank = 0 - self.gpu_global_rank = 0 - self.multi_gpu_cfg = None - return - - # get rank and world size - self.gpu_local_rank = int(os.getenv("LOCAL_RANK", "0")) - self.gpu_global_rank = int(os.getenv("RANK", "0")) - - # make a configuration dictionary - self.multi_gpu_cfg = { - "global_rank": self.gpu_global_rank, # rank of the main process - "local_rank": self.gpu_local_rank, # rank of the current process - "world_size": self.gpu_world_size, # total number of processes - } - - # check if user has device specified for local rank - if self.device != f"cuda:{self.gpu_local_rank}": - raise ValueError( - f"Device '{self.device}' does not match expected device for local rank '{self.gpu_local_rank}'." - ) - # validate multi-gpu configuration - if self.gpu_local_rank >= self.gpu_world_size: - raise ValueError( - f"Local rank '{self.gpu_local_rank}' is greater than or equal to world size '{self.gpu_world_size}'." - ) - if self.gpu_global_rank >= self.gpu_world_size: - raise ValueError( - f"Global rank '{self.gpu_global_rank}' is greater than or equal to world size '{self.gpu_world_size}'." - ) - - # initialize torch distributed - torch.distributed.init_process_group(backend="nccl", rank=self.gpu_global_rank, world_size=self.gpu_world_size) - # set device to the local rank - torch.cuda.set_device(self.gpu_local_rank) - - def _construct_algorithm(self, obs) -> PPO: - """Construct the actor-critic algorithm.""" - # resolve RND config - self.alg_cfg = resolve_rnd_config(self.alg_cfg, obs, self.cfg["obs_groups"], self.env) - - # resolve symmetry config - self.alg_cfg = resolve_symmetry_config(self.alg_cfg, self.env) - - # resolve deprecated normalization config - if self.cfg.get("empirical_normalization") is not None: - warnings.warn( - "The `empirical_normalization` parameter is deprecated. Please set `actor_obs_normalization` and " - "`critic_obs_normalization` as part of the `policy` configuration instead.", - DeprecationWarning, - ) - if self.policy_cfg.get("actor_obs_normalization") is None: - self.policy_cfg["actor_obs_normalization"] = self.cfg["empirical_normalization"] - if self.policy_cfg.get("critic_obs_normalization") is None: - self.policy_cfg["critic_obs_normalization"] = self.cfg["empirical_normalization"] - - # initialize the actor-critic - actor_critic_class = eval(self.policy_cfg.pop("class_name")) - actor_critic: ActorCritic | ActorCriticRecurrent = actor_critic_class( - obs, self.cfg["obs_groups"], self.env.num_actions, **self.policy_cfg - ).to(self.device) - - # initialize the algorithm - alg_class = eval(self.alg_cfg.pop("class_name")) - alg: PPO = alg_class(actor_critic, device=self.device, **self.alg_cfg, multi_gpu_cfg=self.multi_gpu_cfg) - - # initialize the storage - alg.init_storage( - "rl", - self.env.num_envs, - self.num_steps_per_env, - obs, - [self.env.num_actions], - ) - - return alg - - def _prepare_logging_writer(self): - """Prepares the logging writers.""" - if self.log_dir is not None and self.writer is None and not self.disable_logs: - # Launch either Tensorboard or Neptune & Tensorboard summary writer(s), default: Tensorboard. - self.logger_type = self.cfg.get("logger", "tensorboard") - self.logger_type = self.logger_type.lower() - - if self.logger_type == "neptune": - from rsl_rl.utils.neptune_utils import NeptuneSummaryWriter - - self.writer = NeptuneSummaryWriter(log_dir=self.log_dir, flush_secs=10, cfg=self.cfg) - self.writer.log_config(self.env.cfg, self.cfg, self.alg_cfg, self.policy_cfg) - elif self.logger_type == "wandb": - from rsl_rl.utils.wandb_utils import WandbSummaryWriter - - self.writer = WandbSummaryWriter(log_dir=self.log_dir, flush_secs=10, cfg=self.cfg) - self.writer.log_config(self.env.cfg, self.cfg, self.alg_cfg, self.policy_cfg) - elif self.logger_type == "tensorboard": - from torch.utils.tensorboard import SummaryWriter - - self.writer = SummaryWriter(log_dir=self.log_dir, flush_secs=10) - else: - raise ValueError("Logger type not found. Please choose 'neptune', 'wandb' or 'tensorboard'.") diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/storage/__init__.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/storage/__init__.py deleted file mode 100644 index 1624330..0000000 --- a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/storage/__init__.py +++ /dev/null @@ -1,10 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -"""Implementation of transitions storage for RL-agent.""" - -from .rollout_storage import RolloutStorage -from .replay_buffer_multi import ReplayBufferMulti -__all__ = ["RolloutStorage", "ReplayBufferMulti"] diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/storage/replay_buffer_multi.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/storage/replay_buffer_multi.py deleted file mode 100644 index 6462b8e..0000000 --- a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/storage/replay_buffer_multi.py +++ /dev/null @@ -1,38 +0,0 @@ -import torch -import numpy as np - - -class ReplayBufferMulti: - """Fixed-size buffer to store experience tuples.""" - - def __init__(self, obs_dim, buffer_size, num_amp_frames, device): - """Initialize a ReplayBuffer object. - Arguments: - buffer_size (int): maximum size of buffer - """ - self.states = torch.zeros(buffer_size, num_amp_frames, obs_dim).to(device) - self.num_amp_frames = num_amp_frames - self.buffer_size = buffer_size - self.device = device - - self.step = 0 - self.num_samples = 0 - - def insert(self, states): - """Add new states to memory.""" - num_states = states.shape[0] - start_idx = self.step - end_idx = self.step + num_states - if end_idx > self.buffer_size: - self.states[self.step:self.buffer_size] = states[:self.buffer_size - self.step] - self.states[:end_idx - self.buffer_size] = states[self.buffer_size - self.step:] - else: - self.states[start_idx:end_idx] = states - - self.num_samples = min(self.buffer_size, max(end_idx, self.num_samples)) - self.step = (self.step + num_states) % self.buffer_size - - def feed_forward_generator(self, num_mini_batch, mini_batch_size): - for _ in range(num_mini_batch): - sample_idxs = np.random.choice(self.num_samples, size=mini_batch_size) - yield (self.states[sample_idxs].to(self.device)) diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/storage/rollout_storage.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/storage/rollout_storage.py deleted file mode 100644 index e9309b3..0000000 --- a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/storage/rollout_storage.py +++ /dev/null @@ -1,260 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -from __future__ import annotations - -import torch -from tensordict import TensorDict - -from rsl_rl.utils import split_and_pad_trajectories - - -class RolloutStorage: - class Transition: - def __init__(self): - self.observations = None - self.actions = None - self.privileged_actions = None - self.rewards = None - self.dones = None - self.values = None - self.actions_log_prob = None - self.action_mean = None - self.action_sigma = None - self.hidden_states = None - - def clear(self): - self.__init__() - - def __init__( - self, - training_type, - num_envs, - num_transitions_per_env, - obs, - actions_shape, - device="cpu", - ): - # store inputs - self.training_type = training_type - self.device = device - self.num_transitions_per_env = num_transitions_per_env - self.num_envs = num_envs - self.actions_shape = actions_shape - - # Core - self.observations = TensorDict( - {key: torch.zeros(num_transitions_per_env, *value.shape, device=device) for key, value in obs.items()}, - batch_size=[num_transitions_per_env, num_envs], - device=self.device, - ) - self.rewards = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device) - self.actions = torch.zeros(num_transitions_per_env, num_envs, *actions_shape, device=self.device) - self.dones = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device).byte() - - # for distillation - if training_type == "distillation": - self.privileged_actions = torch.zeros(num_transitions_per_env, num_envs, *actions_shape, device=self.device) - - # for reinforcement learning - if training_type == "rl": - self.values = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device) - self.actions_log_prob = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device) - self.mu = torch.zeros(num_transitions_per_env, num_envs, *actions_shape, device=self.device) - self.sigma = torch.zeros(num_transitions_per_env, num_envs, *actions_shape, device=self.device) - self.returns = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device) - self.advantages = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device) - - # For RNN networks - self.saved_hidden_states_a = None - self.saved_hidden_states_c = None - - # counter for the number of transitions stored - self.step = 0 - - def add_transitions(self, transition: Transition): - # check if the transition is valid - if self.step >= self.num_transitions_per_env: - raise OverflowError("Rollout buffer overflow! You should call clear() before adding new transitions.") - - # Core - self.observations[self.step].copy_(transition.observations) - self.actions[self.step].copy_(transition.actions) - self.rewards[self.step].copy_(transition.rewards.view(-1, 1)) - self.dones[self.step].copy_(transition.dones.view(-1, 1)) - - # for distillation - if self.training_type == "distillation": - self.privileged_actions[self.step].copy_(transition.privileged_actions) - - # for reinforcement learning - if self.training_type == "rl": - self.values[self.step].copy_(transition.values) - self.actions_log_prob[self.step].copy_(transition.actions_log_prob.view(-1, 1)) - self.mu[self.step].copy_(transition.action_mean) - self.sigma[self.step].copy_(transition.action_sigma) - - # For RNN networks - self._save_hidden_states(transition.hidden_states) - - # increment the counter - self.step += 1 - - def _save_hidden_states(self, hidden_states): - if hidden_states is None or hidden_states == (None, None): - return - # make a tuple out of GRU hidden state sto match the LSTM format - hid_a = hidden_states[0] if isinstance(hidden_states[0], tuple) else (hidden_states[0],) - hid_c = hidden_states[1] if isinstance(hidden_states[1], tuple) else (hidden_states[1],) - # initialize if needed - if self.saved_hidden_states_a is None: - self.saved_hidden_states_a = [ - torch.zeros(self.observations.shape[0], *hid_a[i].shape, device=self.device) for i in range(len(hid_a)) - ] - self.saved_hidden_states_c = [ - torch.zeros(self.observations.shape[0], *hid_c[i].shape, device=self.device) for i in range(len(hid_c)) - ] - # copy the states - for i in range(len(hid_a)): - self.saved_hidden_states_a[i][self.step].copy_(hid_a[i]) - self.saved_hidden_states_c[i][self.step].copy_(hid_c[i]) - - def clear(self): - self.step = 0 - - def compute_returns(self, last_values, gamma, lam, normalize_advantage: bool = True): - advantage = 0 - for step in reversed(range(self.num_transitions_per_env)): - # if we are at the last step, bootstrap the return value - if step == self.num_transitions_per_env - 1: - next_values = last_values - else: - next_values = self.values[step + 1] - # 1 if we are not in a terminal state, 0 otherwise - next_is_not_terminal = 1.0 - self.dones[step].float() - # TD error: r_t + gamma * V(s_{t+1}) - V(s_t) - delta = self.rewards[step] + next_is_not_terminal * gamma * next_values - self.values[step] - # Advantage: A(s_t, a_t) = delta_t + gamma * lambda * A(s_{t+1}, a_{t+1}) - advantage = delta + next_is_not_terminal * gamma * lam * advantage - # Return: R_t = A(s_t, a_t) + V(s_t) - self.returns[step] = advantage + self.values[step] - - # Compute the advantages - self.advantages = self.returns - self.values - # Normalize the advantages if flag is set - # This is to prevent double normalization (i.e. if per minibatch normalization is used) - if normalize_advantage: - self.advantages = (self.advantages - self.advantages.mean()) / (self.advantages.std() + 1e-8) - - # for distillation - def generator(self): - if self.training_type != "distillation": - raise ValueError("This function is only available for distillation training.") - - for i in range(self.num_transitions_per_env): - yield self.observations[i], self.actions[i], self.privileged_actions[i], self.dones[i] - - # for reinforcement learning with feedforward networks - def mini_batch_generator(self, num_mini_batches, num_epochs=8): - if self.training_type != "rl": - raise ValueError("This function is only available for reinforcement learning training.") - batch_size = self.num_envs * self.num_transitions_per_env - mini_batch_size = batch_size // num_mini_batches - indices = torch.randperm(num_mini_batches * mini_batch_size, requires_grad=False, device=self.device) - - # Core - observations = self.observations.flatten(0, 1) - actions = self.actions.flatten(0, 1) - values = self.values.flatten(0, 1) - returns = self.returns.flatten(0, 1) - - # For PPO - old_actions_log_prob = self.actions_log_prob.flatten(0, 1) - advantages = self.advantages.flatten(0, 1) - old_mu = self.mu.flatten(0, 1) - old_sigma = self.sigma.flatten(0, 1) - - for epoch in range(num_epochs): - for i in range(num_mini_batches): - # Select the indices for the mini-batch - start = i * mini_batch_size - end = (i + 1) * mini_batch_size - batch_idx = indices[start:end] - - # Create the mini-batch - # -- Core - obs_batch = observations[batch_idx] - actions_batch = actions[batch_idx] - - # -- For PPO - target_values_batch = values[batch_idx] - returns_batch = returns[batch_idx] - old_actions_log_prob_batch = old_actions_log_prob[batch_idx] - advantages_batch = advantages[batch_idx] - old_mu_batch = old_mu[batch_idx] - old_sigma_batch = old_sigma[batch_idx] - - # yield the mini-batch - yield obs_batch, actions_batch, target_values_batch, advantages_batch, returns_batch, old_actions_log_prob_batch, old_mu_batch, old_sigma_batch, ( - None, - None, - ), None - - # for reinfrocement learning with recurrent networks - def recurrent_mini_batch_generator(self, num_mini_batches, num_epochs=8): - if self.training_type != "rl": - raise ValueError("This function is only available for reinforcement learning training.") - padded_obs_trajectories, trajectory_masks = split_and_pad_trajectories(self.observations, self.dones) - - mini_batch_size = self.num_envs // num_mini_batches - for ep in range(num_epochs): - first_traj = 0 - for i in range(num_mini_batches): - start = i * mini_batch_size - stop = (i + 1) * mini_batch_size - - dones = self.dones.squeeze(-1) - last_was_done = torch.zeros_like(dones, dtype=torch.bool) - last_was_done[1:] = dones[:-1] - last_was_done[0] = True - trajectories_batch_size = torch.sum(last_was_done[:, start:stop]) - last_traj = first_traj + trajectories_batch_size - - masks_batch = trajectory_masks[:, first_traj:last_traj] - obs_batch = padded_obs_trajectories[:, first_traj:last_traj] - actions_batch = self.actions[:, start:stop] - old_mu_batch = self.mu[:, start:stop] - old_sigma_batch = self.sigma[:, start:stop] - returns_batch = self.returns[:, start:stop] - advantages_batch = self.advantages[:, start:stop] - values_batch = self.values[:, start:stop] - old_actions_log_prob_batch = self.actions_log_prob[:, start:stop] - - # reshape to [num_envs, time, num layers, hidden dim] (original shape: [time, num_layers, num_envs, hidden_dim]) - # then take only time steps after dones (flattens num envs and time dimensions), - # take a batch of trajectories and finally reshape back to [num_layers, batch, hidden_dim] - last_was_done = last_was_done.permute(1, 0) - hid_a_batch = [ - saved_hidden_states.permute(2, 0, 1, 3)[last_was_done][first_traj:last_traj] - .transpose(1, 0) - .contiguous() - for saved_hidden_states in self.saved_hidden_states_a - ] - hid_c_batch = [ - saved_hidden_states.permute(2, 0, 1, 3)[last_was_done][first_traj:last_traj] - .transpose(1, 0) - .contiguous() - for saved_hidden_states in self.saved_hidden_states_c - ] - # remove the tuple for GRU - hid_a_batch = hid_a_batch[0] if len(hid_a_batch) == 1 else hid_a_batch - hid_c_batch = hid_c_batch[0] if len(hid_c_batch) == 1 else hid_c_batch - - yield obs_batch, actions_batch, values_batch, advantages_batch, returns_batch, old_actions_log_prob_batch, old_mu_batch, old_sigma_batch, ( - hid_a_batch, - hid_c_batch, - ), masks_batch - - first_traj = last_traj diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/__init__.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/__init__.py deleted file mode 100644 index f5781f1..0000000 --- a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/__init__.py +++ /dev/null @@ -1,13 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -"""Helper functions.""" - -from .utils import * -from .motion_loader_g1 import G1_AMPLoader - -__all__ = [ - "G1_AMPLoader", -] \ No newline at end of file diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/motion_loader_g1.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/motion_loader_g1.py deleted file mode 100644 index 677c630..0000000 --- a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/motion_loader_g1.py +++ /dev/null @@ -1,388 +0,0 @@ -import os -from os.path import join as pjoin -import glob -import json -import logging - -import torch -import numpy as np -from pybullet_utils import transformations - -from rsl_rl.utils import motion_util - -_EPS = np.finfo(float).eps * 4.0 -def quaternion_slerp(q0, q1, fraction, spin=0, shortestpath=True): - """Batch quaternion spherical linear interpolation.""" - - out = torch.zeros_like(q0) - - zero_mask = torch.isclose(fraction, torch.zeros_like(fraction)).squeeze() - ones_mask = torch.isclose(fraction, torch.ones_like(fraction)).squeeze() - out[zero_mask] = q0[zero_mask] - out[ones_mask] = q1[ones_mask] - - d = torch.sum(q0 * q1, dim=-1, keepdim=True) - dist_mask = (torch.abs(torch.abs(d) - 1.0) < _EPS).squeeze() - out[dist_mask] = q0[dist_mask] - - if shortestpath: - d_old = torch.clone(d) - d = torch.where(d_old < 0, -d, d) - q1 = torch.where(d_old < 0, -q1, q1) - - angle = torch.acos(d) + spin * torch.pi - angle_mask = (torch.abs(angle) < _EPS).squeeze() - out[angle_mask] = q0[angle_mask] - - final_mask = torch.logical_or(zero_mask, ones_mask) - final_mask = torch.logical_or(final_mask, dist_mask) - final_mask = torch.logical_or(final_mask, angle_mask) - final_mask = torch.logical_not(final_mask) - - isin = 1.0 / angle - q0 *= torch.sin((1.0 - fraction) * angle) * isin - q1 *= torch.sin(fraction * angle) * isin - q0 += q1 - out[final_mask] = q0[final_mask] - return out - - -class G1_AMPLoader: - - def __init__( - self, - device, - time_between_frames, - motion_files, - preload_transitions=False, - num_preload_transitions=1000000, - num_frames=5, - ): - """Expert dataset provides AMP observations from Dog mocap dataset. - - time_between_frames: Amount of time in seconds between transition. - """ - self.device = device - self.time_between_frames = time_between_frames - self.num_frames = num_frames - - # Values to store for each trajectory. - self.trajectories = [] - self.trajectories_full = [] - self.trajectory_names = [] - self.trajectory_idxs = [] - self.trajectory_lens = [] # Traj length in seconds. - self.trajectory_weights = [] - self.trajectory_frame_durations = [] - self.trajectory_num_frames = [] - self.motion_dir = motion_files - # import ipdb; ipdb.set_trace() - for i, motion_file in enumerate(os.listdir(motion_files)): - self.trajectory_names.append(motion_file) - motion_path = pjoin(motion_files, motion_file) - motion_data = np.load(motion_path, allow_pickle=True) - motion_data_processed = np.zeros((motion_data.shape[0],36)) - - for f_i in range(motion_data.shape[0]): - motion_data_processed[f_i, :3] = motion_data[f_i, :3] # base pos - motion_data_processed[f_i, 3:7] = motion_data[f_i, 3:7] # base quat (wxyz) - motion_data_processed[f_i, 7:35] = motion_data[f_i, 7:35] # base vel - ''' - NOTE The order of motion_data_processed is - base pos 0:3, - base quat 3:7, wxyz - dof pos 7:36, (mujoco joint order) - ''' - self.trajectories.append(torch.tensor( - motion_data_processed[:, 7:], - dtype=torch.float32, - device=self.device - )) - - self.trajectories_full.append(torch.tensor( - motion_data_processed, - dtype=torch.float32, - device=self.device - )) - - self.trajectory_idxs.append(i) - self.trajectory_weights.append(1 / len(os.listdir(motion_files))) - frame_duration = 1 / 50 - - self.trajectory_frame_durations.append(frame_duration) - traj_len = (motion_data_processed.shape[0] - 1) * frame_duration # seconds - self.trajectory_lens.append(traj_len) - self.trajectory_num_frames.append(float(motion_data_processed.shape[0])) - print(f"Loaded {traj_len}s. motion from {motion_file}.") - - # Trajectory weights are used to sample some trajectories more than others. - self.trajectory_weights = np.array(self.trajectory_weights) / np.sum(self.trajectory_weights) - self.trajectory_frame_durations = np.array(self.trajectory_frame_durations) - self.trajectory_lens = np.array(self.trajectory_lens) - self.trajectory_num_frames = np.array(self.trajectory_num_frames) - - # Preload transitions. - self.preload_transitions = preload_transitions - if self.preload_transitions: - print(f'Preloading {num_preload_transitions} transitions') - - traj_idxs = self.weighted_traj_idx_sample_batch(num_preload_transitions) - times = self.traj_time_sample_batch(traj_idxs) - self.preloaded_s_prior = self.get_full_frame_at_time_batch(traj_idxs, times - self.time_between_frames) - self.preloaded_s = self.get_full_frame_at_time_batch(traj_idxs, times) - self.preloaded_s_next = self.get_full_frame_at_time_batch(traj_idxs, times + self.time_between_frames) - print(f'Finished preloading') - - # 预加载多帧数据 - self.preloaded_frames = [] - for i in range(self.num_frames): - frame_time = times + (i - (self.num_frames - 2)) * self.time_between_frames - full_frame = self.get_full_frame_at_time_batch(traj_idxs, frame_time) - # 预处理:提前提取并连接需要的列(7:26 和 29:33),避免每次生成时重复切片 - processed_frame = torch.cat([ - full_frame[:, 7:26], - full_frame[:, 29:33] - ], dim=-1) - self.preloaded_frames.append(processed_frame) - print(f'Finished preloading multiple frames') - - self.all_trajectories_full = torch.vstack(self.trajectories_full) - - def weighted_traj_idx_sample(self): - """Get traj idx via weighted sampling.""" - return np.random.choice( - self.trajectory_idxs, p=self.trajectory_weights) - - def weighted_traj_idx_sample_batch(self, size): - """Batch sample traj idxs.""" - return np.random.choice( - self.trajectory_idxs, size=size, p=self.trajectory_weights, - replace=True) - - def traj_time_sample(self, traj_idx): - """Sample random time for traj.""" - subst = self.time_between_frames + self.trajectory_frame_durations[traj_idx] - return max( - 0, (self.trajectory_lens[traj_idx] * np.random.uniform() - subst)) - - def traj_time_sample_batch(self, traj_idxs): - """Sample random time for multiple trajectories.""" - subst = self.time_between_frames + self.trajectory_frame_durations[traj_idxs] - time_samples = self.trajectory_lens[traj_idxs] * np.random.uniform(size=len(traj_idxs)) - subst - return np.maximum(np.zeros_like(time_samples), time_samples) - - def slerp(self, val0, val1, blend): - return (1.0 - blend) * val0 + blend * val1 - - def get_trajectory(self, traj_idx): - """Returns trajectory of AMP observations.""" - return self.trajectories_full[traj_idx] - - def get_frame_at_time(self, traj_idx, time): - """Returns frame for the given trajectory at the specified time.""" - p = float(time) / self.trajectory_lens[traj_idx] - n = self.trajectories[traj_idx].shape[0] - idx_low, idx_high = int(np.floor(p * n)), int(np.ceil(p * n)) - frame_start = self.trajectories[traj_idx][idx_low] - frame_end = self.trajectories[traj_idx][idx_high] - blend = p * n - idx_low - return self.slerp(frame_start, frame_end, blend) - - def get_frame_at_time_batch(self, traj_idxs, times): - """Returns frame for the given trajectory at the specified time.""" - p = times / self.trajectory_lens[traj_idxs] - n = self.trajectory_num_frames[traj_idxs] - idx_low, idx_high = np.floor(p * n).astype(np.int32), np.ceil(p * n).astype(np.int32) - all_frame_starts = torch.zeros(len(traj_idxs), self.observation_dim, device=self.device) - all_frame_ends = torch.zeros(len(traj_idxs), self.observation_dim, device=self.device) - for traj_idx in set(traj_idxs): - trajectory = self.trajectories[traj_idx] - traj_mask = traj_idxs == traj_idx - all_frame_starts[traj_mask] = trajectory[idx_low[traj_mask]] - all_frame_ends[traj_mask] = trajectory[idx_high[traj_mask]] - blend = torch.tensor(p * n - idx_low, device=self.device, dtype=torch.float32).unsqueeze(-1) - return self.slerp(all_frame_starts, all_frame_ends, blend) - - def get_full_frame_at_time(self, traj_idx, time): - """Returns full frame for the given trajectory at the specified time.""" - p = float(time) / self.trajectory_lens[traj_idx] - n = self.trajectories_full[traj_idx].shape[0] - idx_low, idx_high = int(np.floor(p * n)), int(np.ceil(p * n)) - frame_start = self.trajectories_full[traj_idx][idx_low] - frame_end = self.trajectories_full[traj_idx][idx_high] - blend = p * n - idx_low - print(idx_low, idx_high) - return self.blend_frame_pose(frame_start, frame_end, blend) - - def get_full_frame_at_time_batch(self, traj_idxs, times): - p = times / self.trajectory_lens[traj_idxs] - n = self.trajectory_num_frames[traj_idxs] - idx_low, idx_high = np.floor(p * n).astype(np.int32), np.ceil(p * n).astype(np.int32) - all_frame_pos_starts = torch.zeros(len(traj_idxs), 3, device=self.device) - all_frame_pos_ends = torch.zeros(len(traj_idxs), 3, device=self.device) - all_frame_rot_starts = torch.zeros(len(traj_idxs), 4, device=self.device) - all_frame_rot_ends = torch.zeros(len(traj_idxs), 4, device=self.device) - all_frame_amp_starts = torch.zeros(len(traj_idxs), 29, device=self.device) - all_frame_amp_ends = torch.zeros(len(traj_idxs), 29, device=self.device) - for traj_idx in set(traj_idxs): - trajectory = self.trajectories_full[traj_idx] - traj_mask = traj_idxs == traj_idx - all_frame_pos_starts[traj_mask] = G1_AMPLoader.get_root_pos_batch(trajectory[idx_low[traj_mask]]) - all_frame_pos_ends[traj_mask] = G1_AMPLoader.get_root_pos_batch(trajectory[idx_high[traj_mask]]) - all_frame_rot_starts[traj_mask] = G1_AMPLoader.get_root_rot_batch(trajectory[idx_low[traj_mask]]) - all_frame_rot_ends[traj_mask] = G1_AMPLoader.get_root_rot_batch(trajectory[idx_high[traj_mask]]) - all_frame_amp_starts[traj_mask] = trajectory[idx_low[traj_mask]][:, 7:36] # base vel3+ang3, dof vel23+ang23 - all_frame_amp_ends[traj_mask] = trajectory[idx_high[traj_mask]][:, 7:36] # base vel3+ang3, dof vel23+ang23 - blend = torch.tensor(p * n - idx_low, device=self.device, dtype=torch.float32).unsqueeze(-1) - pos_blend = self.slerp(all_frame_pos_starts, all_frame_pos_ends, blend) - rot_blend = quaternion_slerp(all_frame_rot_starts, all_frame_rot_ends, blend) - amp_blend = self.slerp(all_frame_amp_starts, all_frame_amp_ends, blend) - return torch.cat([pos_blend, rot_blend, amp_blend], dim=-1) - - def get_frame(self): - """Returns random frame.""" - traj_idx = self.weighted_traj_idx_sample() - sampled_time = self.traj_time_sample(traj_idx) - return self.get_frame_at_time(traj_idx, sampled_time) - - def get_full_frame(self): - """Returns random full frame.""" - traj_idx = self.weighted_traj_idx_sample() - sampled_time = self.traj_time_sample(traj_idx) - return self.get_full_frame_at_time(traj_idx, sampled_time) - - def get_full_frame_batch(self, num_frames): - if self.preload_transitions: - idxs = np.random.choice( - self.preloaded_s.shape[0], size=num_frames) - return self.preloaded_s[idxs] - else: - traj_idxs = self.weighted_traj_idx_sample_batch(num_frames) - times = self.traj_time_sample_batch(traj_idxs) - return self.get_full_frame_at_time_batch(traj_idxs, times) - - def blend_frame_pose(self, frame0, frame1, blend): - """Linearly interpolate between two frames, including orientation. - - Args: - frame0: First frame to be blended corresponds to (blend = 0). - frame1: Second frame to be blended corresponds to (blend = 1). - blend: Float between [0, 1], specifying the interpolation between - the two frames. - Returns: - An interpolation of the two frames. - """ - root_pos0, root_pos1 = G1_AMPLoader.get_root_pos(frame0), G1_AMPLoader.get_root_pos(frame1) - root_rot0, root_rot1 = G1_AMPLoader.get_root_rot(frame0), G1_AMPLoader.get_root_rot(frame1) - joints0, joints1 = G1_AMPLoader.get_joint_pose(frame0), G1_AMPLoader.get_joint_pose(frame1) - # tar_toe_pos_0, tar_toe_pos_1 = G1_AMPLoader.get_tar_toe_pos_local(frame0), G1_AMPLoader.get_tar_toe_pos_local(frame1) - linear_vel_0, linear_vel_1 = G1_AMPLoader.get_linear_vel(frame0), G1_AMPLoader.get_linear_vel(frame1) - angular_vel_0, angular_vel_1 = G1_AMPLoader.get_angular_vel(frame0), G1_AMPLoader.get_angular_vel(frame1) - joint_vel_0, joint_vel_1 = G1_AMPLoader.get_joint_vel(frame0), G1_AMPLoader.get_joint_vel(frame1) - - blend_root_pos = self.slerp(root_pos0, root_pos1, blend) - blend_root_rot = transformations.quaternion_slerp(root_rot0.cpu().numpy(), root_rot1.cpu().numpy(), blend) - blend_root_rot = torch.tensor(motion_util.standardize_quaternion(blend_root_rot),dtype=torch.float32, device=self.device) - blend_joints = self.slerp(joints0, joints1, blend) - # blend_tar_toe_pos = self.slerp(tar_toe_pos_0, tar_toe_pos_1, blend) - blend_linear_vel = self.slerp(linear_vel_0, linear_vel_1, blend) - blend_angular_vel = self.slerp(angular_vel_0, angular_vel_1, blend) - blend_joints_vel = self.slerp(joint_vel_0, joint_vel_1, blend) - - # return - # torch.cat([ - # blend_root_pos, blend_root_rot, blend_linear_vel, blend_angular_vel, blend_joints, blend_joints_vel]) - return torch.cat([blend_root_pos, blend_root_rot, blend_linear_vel, blend_angular_vel, blend_joints]) - - def feed_forward_generator_23dof_multi(self, num_mini_batch, mini_batch_size): - """Generates a batch of AMP transitions.""" - # import ipdb; ipdb.set_trace() - for _ in range(num_mini_batch): - if self.preload_transitions: - idxs = np.random.choice(self.preloaded_s.shape[0], size=mini_batch_size) - - frames = [] - for i in range(self.num_frames): - # 数据已在预加载时预处理,直接索引即可 - s = self.preloaded_frames[i][idxs] - frames.append(s) - else: - NotImplementedError('preload transition') - yield torch.stack(frames, dim=1) # [batch, num_frames, 16] - - - - - def quaternion_to_euler_array(self, quat): - # Ensure quaternion is in the correct format [x, y, z, w] - x, y, z, w =quat - - # Roll (x-axis rotation) - t0 = +2.0 * (w * x + y * z) - t1 = +1.0 - 2.0 * (x * x + y * y) - roll_x = np.arctan2(t0, t1) - - # Pitch (y-axis rotation) - t2 = +2.0 * (w * y - z * x) - t2 = np.clip(t2, -1.0, 1.0) - pitch_y = np.arcsin(t2) - - # Yaw (z-axis rotation) - t3 = +2.0 * (w * z + x * y) - t4 = +1.0 - 2.0 * (y * y + z * z) - yaw_z = np.arctan2(t3, t4) - - # Returns roll, pitch, yaw in a NumPy array in radians - return np.array([roll_x, pitch_y, yaw_z]) - - def euler_to_quaternion(self, root_rot): - roll, pitch, yaw = root_rot[0], root_rot[1], root_rot[2] - cy = np.cos(yaw * 0.5) - sy = np.sin(yaw * 0.5) - cp = np.cos(pitch * 0.5) - sp = np.sin(pitch * 0.5) - cr = np.cos(roll * 0.5) - sr = np.sin(roll * 0.5) - - qw = cy * cp * cr + sy * sp * sr - qx = cy * cp * sr - sy * sp * cr - qy = sy * cp * sr + cy * sp * cr - qz = sy * cp * cr - cy * sp * sr - - return np.array([qx, qy, qz, qw]) - - @property - def observation_dim(self): - """Size of AMP observations.""" - return self.trajectories[0].shape[1] + 1 - - @property - def num_motions(self): - return len(self.trajectory_names) - @staticmethod - def get_root_pos(pose): - return pose[0:3] - - @staticmethod - def get_root_pos_batch(poses): - return poses[:, 0:3] - - @staticmethod - def get_root_rot(pose): - return pose[3:7] - - @staticmethod - def get_root_rot_batch(poses): - return poses[:, 3:7] - - @staticmethod - def get_joint_pose_batch_12dof(poses): - return poses[:, 13:25] - - @staticmethod - def get_tar_toe_pos_local(pose): - return pose[G1_AMPLoader.TAR_TOE_POS_LOCAL_START_IDX:G1_AMPLoader.TAR_TOE_POS_LOCAL_END_IDX] - - @staticmethod - def get_tar_toe_pos_local_batch(poses): - return poses[:, G1_AMPLoader.TAR_TOE_POS_LOCAL_START_IDX:G1_AMPLoader.TAR_TOE_POS_LOCAL_END_IDX] diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/motion_util.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/motion_util.py deleted file mode 100644 index 3d49bfa..0000000 --- a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/motion_util.py +++ /dev/null @@ -1,97 +0,0 @@ -# coding=utf-8 -# Copyright 2020 The Google Research Authors. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Utility functions for processing motion clips.""" - -import os -import inspect -currentdir = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe()))) -parentdir = os.path.dirname(os.path.dirname(currentdir)) -os.sys.path.insert(0, parentdir) - -import numpy as np - -from rsl_rl.utils import pose3d -# from pybullet_utils import transformations - - -def standardize_quaternion(q): - """Returns a quaternion where q.w >= 0 to remove redundancy due to q = -q. - - Args: - q: A quaternion to be standardized. - - Returns: - A quaternion with q.w >= 0. - - """ - if q[-1] < 0: - q = -q - return q - - -def normalize_rotation_angle(theta): - """Returns a rotation angle normalized between [-pi, pi]. - - Args: - theta: angle of rotation (radians). - - Returns: - An angle of rotation normalized between [-pi, pi]. - - """ - norm_theta = theta - if np.abs(norm_theta) > np.pi: - norm_theta = np.fmod(norm_theta, 2 * np.pi) - if norm_theta >= 0: - norm_theta += -2 * np.pi - else: - norm_theta += 2 * np.pi - - return norm_theta - - -def calc_heading(q): - """Returns the heading of a rotation q, specified as a quaternion. - - The heading represents the rotational component of q along the vertical - axis (z axis). - - Args: - q: A quaternion that the heading is to be computed from. - - Returns: - An angle representing the rotation about the z axis. - - """ - ref_dir = np.array([1, 0, 0]) - rot_dir = pose3d.QuaternionRotatePoint(ref_dir, q) - heading = np.arctan2(rot_dir[1], rot_dir[0]) - return heading - - -# def calc_heading_rot(q): -# """Return a quaternion representing the heading rotation of q along the vertical axis (z axis). - -# Args: -# q: A quaternion that the heading is to be computed from. - -# Returns: -# A quaternion representing the rotation about the z axis. - -# """ -# heading = calc_heading(q) -# q_heading = transformations.quaternion_about_axis(heading, [0, 0, 1]) -# return q_heading diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/neptune_utils.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/neptune_utils.py deleted file mode 100644 index 3796ec8..0000000 --- a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/neptune_utils.py +++ /dev/null @@ -1,94 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -from __future__ import annotations - -import os -from dataclasses import asdict -from torch.utils.tensorboard import SummaryWriter - -try: - import neptune -except ModuleNotFoundError: - raise ModuleNotFoundError("neptune-client is required to log to Neptune.") - - -class NeptuneLogger: - def __init__(self, project, token): - self.run = neptune.init_run(project=project, api_token=token) - - def store_config(self, env_cfg, runner_cfg, alg_cfg, policy_cfg): - self.run["runner_cfg"] = runner_cfg - self.run["policy_cfg"] = policy_cfg - self.run["alg_cfg"] = alg_cfg - self.run["env_cfg"] = asdict(env_cfg) - - -class NeptuneSummaryWriter(SummaryWriter): - """Summary writer for Neptune.""" - - def __init__(self, log_dir: str, flush_secs: int, cfg): - super().__init__(log_dir, flush_secs) - - try: - project = cfg["neptune_project"] - except KeyError: - raise KeyError("Please specify neptune_project in the runner config, e.g. legged_gym.") - - try: - token = os.environ["NEPTUNE_API_TOKEN"] - except KeyError: - raise KeyError( - "Neptune api token not found. Please run or add to ~/.bashrc: export NEPTUNE_API_TOKEN=YOUR_API_TOKEN" - ) - - try: - entity = os.environ["NEPTUNE_USERNAME"] - except KeyError: - raise KeyError( - "Neptune username not found. Please run or add to ~/.bashrc: export NEPTUNE_USERNAME=YOUR_USERNAME" - ) - - neptune_project = entity + "/" + project - - self.neptune_logger = NeptuneLogger(neptune_project, token) - - self.name_map = { - "Train/mean_reward/time": "Train/mean_reward_time", - "Train/mean_episode_length/time": "Train/mean_episode_length_time", - } - - run_name = os.path.split(log_dir)[-1] - - self.neptune_logger.run["log_dir"].log(run_name) - - def _map_path(self, path): - if path in self.name_map: - return self.name_map[path] - else: - return path - - def add_scalar(self, tag, scalar_value, global_step=None, walltime=None, new_style=False): - super().add_scalar( - tag, - scalar_value, - global_step=global_step, - walltime=walltime, - new_style=new_style, - ) - self.neptune_logger.run[self._map_path(tag)].log(scalar_value, step=global_step) - - def stop(self): - self.neptune_logger.run.stop() - - def log_config(self, env_cfg, runner_cfg, alg_cfg, policy_cfg): - self.neptune_logger.store_config(env_cfg, runner_cfg, alg_cfg, policy_cfg) - - def save_model(self, model_path, iter): - self.neptune_logger.run["model/saved_model_" + str(iter)].upload(model_path) - - def save_file(self, path, iter=None): - name = path.rsplit("/", 1)[-1].split(".")[0] - self.neptune_logger.run["git_diff/" + name].upload(path) diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/pose3d.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/pose3d.py deleted file mode 100644 index a4cba1e..0000000 --- a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/pose3d.py +++ /dev/null @@ -1,283 +0,0 @@ -# coding=utf-8 -# Copyright 2020 The Google Research Authors. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Utilities for 3D pose conversion.""" -import math -import numpy as np - -# from pybullet_utils import transformations - -VECTOR3_0 = np.zeros(3, dtype=np.float64) -VECTOR3_1 = np.ones(3, dtype=np.float64) -VECTOR3_X = np.array([1, 0, 0], dtype=np.float64) -VECTOR3_Y = np.array([0, 1, 0], dtype=np.float64) -VECTOR3_Z = np.array([0, 0, 1], dtype=np.float64) - -# QUATERNION_IDENTITY is the multiplicative identity 1.0 + 0i + 0j + 0k. -# When interpreted as a rotation, it is the identity rotation. -QUATERNION_IDENTITY = np.array([0.0, 0.0, 0.0, 1.0], dtype=np.float64) - - -def Vector3RandomNormal(sigma, mu=VECTOR3_0): - """Returns a random 3D vector from a normal distribution. - - Each component is selected independently from a normal distribution. - - Args: - sigma: Scale (or stddev) of distribution for all variables. - mu: Mean of distribution for each variable. - - Returns: - A 3D vector in a numpy array. - """ - - random_v3 = np.random.normal(scale=sigma, size=3) + mu - return random_v3 - - -def Vector3RandomUniform(low=VECTOR3_0, high=VECTOR3_1): - """Returns a 3D vector selected uniformly from the input box. - - Args: - low: The min-value corner of the box. - high: The max-value corner of the box. - - Returns: - A 3D vector in a numpy array. - """ - - random_x = np.random.uniform(low=low[0], high=high[0]) - random_y = np.random.uniform(low=low[1], high=high[1]) - random_z = np.random.uniform(low=low[2], high=high[2]) - return np.array([random_x, random_y, random_z]) - - -def Vector3RandomUnit(): - """Returns a random 3D vector with unit length. - - Generates a 3D vector selected uniformly from the unit sphere. - - Returns: - A normalized 3D vector in a numpy array. - """ - longitude = np.random.uniform(low=-math.pi, high=math.pi) - sin_latitude = np.random.uniform(low=-1.0, high=1.0) - cos_latitude = math.sqrt(1.0 - sin_latitude * sin_latitude) - x = math.cos(longitude) * cos_latitude - y = math.sin(longitude) * cos_latitude - z = sin_latitude - return np.array([x, y, z], dtype=np.float64) - - -def QuaternionNormalize(q): - """Normalizes the quaternion to length 1. - - Divides the quaternion by its magnitude. If the magnitude is too - small, returns the quaternion identity value (1.0). - - Args: - q: A quaternion to be normalized. - - Raises: - ValueError: If input quaternion has length near zero. - - Returns: - A quaternion with magnitude 1 in a numpy array [x, y, z, w]. - - """ - q_norm = np.linalg.norm(q) - if np.isclose(q_norm, 0.0): - raise ValueError( - 'Quaternion may not be zero in QuaternionNormalize: |q| = %f, q = %s' % - (q_norm, q)) - return q / q_norm - - -def QuaternionFromAxisAngle(axis, angle): - """Returns a quaternion that generates the given axis-angle rotation. - - Returns the quaternion: sin(angle/2) * axis + cos(angle/2). - - Args: - axis: Axis of rotation, a 3D vector in a numpy array. - angle: The angle of rotation (radians). - - Raises: - ValueError: If input axis is not a normalizable 3D vector. - - Returns: - A unit quaternion in a numpy array. - - """ - if len(axis) != 3: - raise ValueError('Axis vector should have three components: %s' % axis) - axis_norm = np.linalg.norm(axis) - if np.isclose(axis_norm, 0.0): - raise ValueError('Axis vector may not have zero length: |v| = %f, v = %s' % - (axis_norm, axis)) - half_angle = angle * 0.5 - q = np.zeros(4, dtype=np.float64) - q[0:3] = axis - q[0:3] *= math.sin(half_angle) / axis_norm - q[3] = math.cos(half_angle) - return q - - -def QuaternionToAxisAngle(quat, default_axis=VECTOR3_Z, direction_axis=None): - """Calculates axis and angle of rotation performed by a quaternion. - - Calculates the axis and angle of the rotation performed by the quaternion. - The quaternion should have four values and be normalized. - - Args: - quat: Unit quaternion in a numpy array. - default_axis: 3D vector axis used if the rotation is near to zero. Without - this default, small rotations would result in an exception. It is - reasonable to use a default axis for tiny rotations, because zero angle - rotations about any axis are equivalent. - direction_axis: Used to disambiguate rotation directions. If the - direction_axis is specified, the axis of the rotation will be chosen such - that its inner product with the direction_axis is non-negative. - - Raises: - ValueError: If quat is not a normalized quaternion. - - Returns: - axis: Axis of rotation. - angle: Angle in radians. - """ - if len(quat) != 4: - raise ValueError( - 'Quaternion should have four components [x, y, z, w]: %s' % quat) - if not np.isclose(1.0, np.linalg.norm(quat)): - raise ValueError('Quaternion should have unit length: |q| = %f, q = %s' % - (np.linalg.norm(quat), quat)) - axis = quat[:3].copy() - axis_norm = np.linalg.norm(axis) - min_axis_norm = 1e-8 - if axis_norm < min_axis_norm: - axis = default_axis - if len(default_axis) != 3: - raise ValueError('Axis vector should have three components: %s' % axis) - if not np.isclose(np.linalg.norm(axis), 1.0): - raise ValueError('Axis vector should have unit length: |v| = %f, v = %s' % - (np.linalg.norm(axis), axis)) - else: - axis /= axis_norm - sin_half_angle = axis_norm - if direction_axis is not None and np.inner(axis, direction_axis) < 0: - sin_half_angle = -sin_half_angle - axis = -axis - cos_half_angle = quat[3] - half_angle = math.atan2(sin_half_angle, cos_half_angle) - angle = half_angle * 2 - return axis, angle - - -def QuaternionRandomRotation(max_angle=math.pi): - """Creates a random small rotation around a random axis. - - Generates a small rotation with the axis vector selected uniformly - from the unit sphere and an angle selected from a uniform - distribution over [0, max_angle]. - - If the max_angle is not specified, the rotation should be selected - uniformly over all possible rotation angles. - - Args: - max_angle: The maximum angle of rotation (radians). - - Returns: - A unit quaternion in a numpy array. - - """ - - angle = np.random.uniform(low=0, high=max_angle) - axis = Vector3RandomUnit() - return QuaternionFromAxisAngle(axis, angle) - - -# def QuaternionRotatePoint(point, quat): -# """Performs a rotation by quaternion. - -# Rotate the point by the quaternion using quaternion multiplication, -# (q * p * q^-1), without constructing the rotation matrix. - -# Args: -# point: The point to be rotated. -# quat: The rotation represented as a quaternion [x, y, z, w]. - -# Returns: -# A 3D vector in a numpy array. -# """ - -# q_point = np.array([point[0], point[1], point[2], 0.0]) -# quat_inverse = transformations.quaternion_inverse(quat) -# q_point_rotated = transformations.quaternion_multiply( -# transformations.quaternion_multiply(quat, q_point), quat_inverse) -# return q_point_rotated[:3] - - -def IsRotationMatrix(m): - """Returns true if the 3x3 submatrix represents a rotation. - - Args: - m: A transformation matrix. - - Raises: - ValueError: If input is not a matrix of size at least 3x3. - - Returns: - True if the 3x3 submatrix is a rotation (orthogonal). - """ - if len(m.shape) != 2 or m.shape[0] < 3 or m.shape[1] < 3: - raise ValueError('Matrix should be 3x3 or 4x4: %s\n %s' % (m.shape, m)) - rot = m[:3, :3] - eye = np.matmul(rot, np.transpose(rot)) - return np.isclose(eye, np.identity(3), atol=1e-4).all() - -# def ZAxisAlignedRobotPoseTool(robot_pose_tool): -# """Returns the current gripper pose rotated for alignment with the z-axis. - -# Args: -# robot_pose_tool: a pose3d.Pose3d() instance. - -# Returns: -# An instance of pose.Transform representing the current gripper pose -# rotated for alignment with the z-axis. -# """ -# # Align the current pose to the z-axis. -# robot_pose_tool.quaternion = transformations.quaternion_multiply( -# RotationBetween( -# robot_pose_tool.matrix4x4[0:3, 0:3].dot(np.array([0, 0, 1])), -# np.array([0.0, 0.0, -1.0])), robot_pose_tool.quaternion) -# return robot_pose_tool - -# def RotationBetween(a_translation_b, a_translation_c): -# """Computes the rotation from one vector to another. - -# The computed rotation has the property that: - -# a_translation_c = a_rotation_b_to_c * a_translation_b - -# Args: -# a_translation_b: vec3, vector to rotate from -# a_translation_c: vec3, vector to rotate to - -# Returns: -# a_rotation_b_to_c: new Orientation -# """ -# rotation = rotation3.Rotation3.rotation_between( -# a_translation_b, a_translation_c, err_msg='RotationBetween') -# return rotation.quaternion.xyzw diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/utils.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/utils.py deleted file mode 100644 index 3605622..0000000 --- a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/utils.py +++ /dev/null @@ -1,360 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -from __future__ import annotations - -import git -import importlib -import os -import pathlib -import torch -import warnings -from tensordict import TensorDict -from typing import Callable -import numpy as np -class RunningMeanStd: - def __init__(self, epsilon: float = 1e-4, shape: Tuple[int, ...] = ()): - """ - Calculates the running mean and std of a data stream - https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Parallel_algorithm - :param epsilon: helps with arithmetic issues - :param shape: the shape of the data stream's output - """ - self.mean = np.zeros(shape, np.float64) - self.var = np.ones(shape, np.float64) - self.count = epsilon - - def update(self, arr: np.ndarray) -> None: - batch_mean = np.mean(arr, axis=0) - batch_var = np.var(arr, axis=0) - batch_count = arr.shape[0] - self.update_from_moments(batch_mean, batch_var, batch_count) - - def update_from_moments(self, batch_mean: np.ndarray, batch_var: np.ndarray, batch_count: int) -> None: - delta = batch_mean - self.mean - tot_count = self.count + batch_count - - new_mean = self.mean + delta * batch_count / tot_count - m_a = self.var * self.count - m_b = batch_var * batch_count - m_2 = m_a + m_b + np.square(delta) * self.count * batch_count / (self.count + batch_count) - new_var = m_2 / (self.count + batch_count) - - new_count = batch_count + self.count - - self.mean = new_mean - self.var = new_var - self.count = new_count - - -class Normalizer(RunningMeanStd): - def __init__(self, input_dim, epsilon=1e-4, clip_obs=10.0): - super().__init__(shape=input_dim) - self.epsilon = epsilon - self.clip_obs = clip_obs - - def normalize(self, input): - return np.clip((input - self.mean) / np.sqrt(self.var + self.epsilon), -self.clip_obs, self.clip_obs) - - def normalize_torch(self, input, device): - mean_torch = torch.tensor(self.mean, device=device, dtype=torch.float32) - std_torch = torch.sqrt(torch.tensor(self.var + self.epsilon, device=device, dtype=torch.float32)) - return torch.clamp((input - mean_torch) / std_torch, -self.clip_obs, self.clip_obs) - - def update_normalizer(self, rollouts, expert_loader): - policy_data_generator = rollouts.feed_forward_generator_amp(None, mini_batch_size=expert_loader.batch_size) - expert_data_generator = expert_loader.dataset.feed_forward_generator_amp(expert_loader.batch_size) - - for expert_batch, policy_batch in zip(expert_data_generator, policy_data_generator): - self.update(torch.vstack(tuple(policy_batch) + tuple(expert_batch)).cpu().numpy()) - - -def resolve_nn_activation(act_name: str) -> torch.nn.Module: - """Resolves the activation function from the name. - - Args: - act_name: The name of the activation function. - - Returns: - The activation function. - - Raises: - ValueError: If the activation function is not found. - """ - act_dict = { - "elu": torch.nn.ELU(), - "selu": torch.nn.SELU(), - "relu": torch.nn.ReLU(), - "crelu": torch.nn.CELU(), - "lrelu": torch.nn.LeakyReLU(), - "tanh": torch.nn.Tanh(), - "sigmoid": torch.nn.Sigmoid(), - "softplus": torch.nn.Softplus(), - "gelu": torch.nn.GELU(), - "swish": torch.nn.SiLU(), - "mish": torch.nn.Mish(), - "identity": torch.nn.Identity(), - } - - act_name = act_name.lower() - if act_name in act_dict: - return act_dict[act_name] - else: - raise ValueError(f"Invalid activation function '{act_name}'. Valid activations are: {list(act_dict.keys())}") - - -def resolve_optimizer(optimizer_name: str) -> torch.optim.Optimizer: - """Resolves the optimizer from the name. - - Args: - optimizer_name: The name of the optimizer. - - Returns: - The optimizer. - - Raises: - ValueError: If the optimizer is not found. - """ - optimizer_dict = { - "adam": torch.optim.Adam, - "adamw": torch.optim.AdamW, - "sgd": torch.optim.SGD, - "rmsprop": torch.optim.RMSprop, - } - - optimizer_name = optimizer_name.lower() - if optimizer_name in optimizer_dict: - return optimizer_dict[optimizer_name] - else: - raise ValueError(f"Invalid optimizer '{optimizer_name}'. Valid optimizers are: {list(optimizer_dict.keys())}") - - -def split_and_pad_trajectories( - tensor: torch.Tensor | TensorDict, dones: torch.Tensor -) -> tuple[torch.Tensor | TensorDict, torch.Tensor]: - """Splits trajectories at done indices. Then concatenates them and pads with zeros up to the length of the longest - trajectory. Returns masks corresponding to valid parts of the trajectories. - - Example: - Input: [[a1, a2, a3, a4 | a5, a6], - [b1, b2 | b3, b4, b5 | b6]] - - Output:[[a1, a2, a3, a4], | [[True, True, True, True], - [a5, a6, 0, 0], | [True, True, False, False], - [b1, b2, 0, 0], | [True, True, False, False], - [b3, b4, b5, 0], | [True, True, True, False], - [b6, 0, 0, 0]] | [True, False, False, False]] - - Assumes that the input has the following order of dimensions: [time, number of envs, additional dimensions] - """ - - dones = dones.clone() - dones[-1] = 1 - # Permute the buffers to have order (num_envs, num_transitions_per_env, ...), for correct reshaping - flat_dones = dones.transpose(1, 0).reshape(-1, 1) - # Get length of trajectory by counting the number of successive not done elements - done_indices = torch.cat((flat_dones.new_tensor([-1], dtype=torch.int64), flat_dones.nonzero()[:, 0])) - trajectory_lengths = done_indices[1:] - done_indices[:-1] - trajectory_lengths_list = trajectory_lengths.tolist() - # Extract the individual trajectories - if isinstance(tensor, TensorDict): - padded_trajectories = {} - for k, v in tensor.items(): - # split the tensor into trajectories - trajectories = torch.split(v.transpose(1, 0).flatten(0, 1), trajectory_lengths_list) - # add at least one full length trajectory - trajectories = trajectories + (torch.zeros(v.shape[0], *v.shape[2:], device=v.device),) - # pad the trajectories to the length of the longest trajectory - padded_trajectories[k] = torch.nn.utils.rnn.pad_sequence(trajectories) - # remove the added tensor - padded_trajectories[k] = padded_trajectories[k][:, :-1] - padded_trajectories = TensorDict( - padded_trajectories, batch_size=[tensor.batch_size[0], len(trajectory_lengths_list)] - ) - else: - # split the tensor into trajectories - trajectories = torch.split(tensor.transpose(1, 0).flatten(0, 1), trajectory_lengths_list) - # add at least one full length trajectory - trajectories = trajectories + (torch.zeros(tensor.shape[0], *tensor.shape[2:], device=tensor.device),) - # pad the trajectories to the length of the longest trajectory - padded_trajectories = torch.nn.utils.rnn.pad_sequence(trajectories) - # remove the added tensor - padded_trajectories = padded_trajectories[:, :-1] - # create masks for the valid parts of the trajectories - trajectory_masks = trajectory_lengths > torch.arange(0, tensor.shape[0], device=tensor.device).unsqueeze(1) - return padded_trajectories, trajectory_masks - - -def unpad_trajectories(trajectories, masks): - """Does the inverse operation of split_and_pad_trajectories()""" - # Need to transpose before and after the masking to have proper reshaping - return ( - trajectories.transpose(1, 0)[masks.transpose(1, 0)] - .view(-1, trajectories.shape[0], trajectories.shape[-1]) - .transpose(1, 0) - ) - - -def store_code_state(logdir, repositories) -> list: - git_log_dir = os.path.join(logdir, "git") - os.makedirs(git_log_dir, exist_ok=True) - file_paths = [] - for repository_file_path in repositories: - try: - repo = git.Repo(repository_file_path, search_parent_directories=True) - t = repo.head.commit.tree - except Exception: - print(f"Could not find git repository in {repository_file_path}. Skipping.") - # skip if not a git repository - continue - # get the name of the repository - repo_name = pathlib.Path(repo.working_dir).name - diff_file_name = os.path.join(git_log_dir, f"{repo_name}.diff") - # check if the diff file already exists - if os.path.isfile(diff_file_name): - continue - # write the diff file - print(f"Storing git diff for '{repo_name}' in: {diff_file_name}") - with open(diff_file_name, "x", encoding="utf-8") as f: - content = f"--- git status ---\n{repo.git.status()} \n\n\n--- git diff ---\n{repo.git.diff(t)}" - f.write(content) - # add the file path to the list of files to be uploaded - file_paths.append(diff_file_name) - return file_paths - - -def string_to_callable(name: str) -> Callable: - """Resolves the module and function names to return the function. - - Args: - name: The function name. The format should be 'module:attribute_name'. - - Raises: - ValueError: When the resolved attribute is not a function. - ValueError: When unable to resolve the attribute. - - Returns: - The function loaded from the module. - """ - try: - mod_name, attr_name = name.split(":") - mod = importlib.import_module(mod_name) - callable_object = getattr(mod, attr_name) - # check if attribute is callable - if callable(callable_object): - return callable_object - else: - raise ValueError(f"The imported object is not callable: '{name}'") - except AttributeError as e: - msg = ( - "We could not interpret the entry as a callable object. The format of input should be" - f" 'module:attribute_name'\nWhile processing input '{name}', received the error:\n {e}." - ) - raise ValueError(msg) - - -def resolve_obs_groups( - obs: TensorDict, obs_groups: dict[str, list[str]], default_sets: list[str] -) -> dict[str, list[str]]: - """Validates the observation configuration and defaults missing observation sets. - - The input is an observation dictionary `obs` containing observation groups and a configuration dictionary - `obs_groups` where the keys are the observation sets and the values are lists of observation groups. - - The configuration dictionary could for example look like: - { - "policy": ["group_1", "group_2"], - "critic": ["group_1", "group_3"] - } - - This means that the 'policy' observation set will contain the observations "group_1" and "group_2" and the - 'critic' observation set will contain the observations "group_1" and "group_3". This function will check that all - the observations in the 'policy' and 'critic' observation sets are present in the observation dictionary from the - environment. - - Additionally, if one of the `default_sets`, e.g. "critic", is not present in the configuration dictionary, - this function will: - - 1. Check if a group with the same name exists in the observations and assign this group to the observation set. - 2. If 1. fails, it will assign the observations from the 'policy' observation set to the default observation set. - - Args: - obs: Observations from the environment in the form of a dictionary. - obs_groups: Observation sets configuration. - default_sets: Reserved observation set names used by the algorithm (besides 'policy'). - If not provided in 'obs_groups', a default behavior gets triggered. - - Returns: - The resolved observation groups. - - Raises: - ValueError: If any observation set is an empty list. - ValueError: If any observation set contains an observation term that is not present in the observations. - """ - # check if policy observation set exists - if "policy" not in obs_groups.keys(): - if "policy" in obs: - obs_groups["policy"] = ["policy"] - warnings.warn( - "The observation configuration dictionary 'obs_groups' must contain the 'policy' key." - " As an observation group with the name 'policy' was found, this is assumed to be the observation set." - " Consider adding the 'policy' key to the 'obs_groups' dictionary for clarity." - " This behavior will be removed in a future version." - ) - else: - raise ValueError( - "The observation configuration dictionary 'obs_groups' must contain the 'policy' key." - f" Found keys: {list(obs_groups.keys())}" - ) - - # check all observation sets for valid observation groups - for set_name, groups in obs_groups.items(): - # check if the list is empty - if len(groups) == 0: - msg = f"The '{set_name}' key in the 'obs_groups' dictionary can not be an empty list." - if set_name in default_sets: - if set_name not in obs: - msg += " Consider removing the key to default to the observations used for the 'policy' set." - else: - msg += ( - f" Consider removing the key to default to the observation '{set_name}' from the environment." - ) - raise ValueError(msg) - # check groups exist inside the observations from the environment - for group in groups: - if group not in obs: - raise ValueError( - f"Observation '{group}' in observation set '{set_name}' not found in the observations from the" - f" environment. Available observations from the environment: {list(obs.keys())}" - ) - - # fill missing observation sets - for default_set_name in default_sets: - if default_set_name not in obs_groups.keys(): - if default_set_name in obs: - obs_groups[default_set_name] = [default_set_name] - warnings.warn( - f"The observation configuration dictionary 'obs_groups' must contain the '{default_set_name}' key." - f" As an observation group with the name '{default_set_name}' was found, this is assumed to be the" - f" observation set. Consider adding the '{default_set_name}' key to the 'obs_groups' dictionary for" - " clarity. This behavior will be removed in a future version." - ) - else: - obs_groups[default_set_name] = obs_groups["policy"].copy() - warnings.warn( - f"The observation configuration dictionary 'obs_groups' must contain the '{default_set_name}' key." - f" As the configuration for '{default_set_name}' is missing, the observations from the 'policy' set" - f" are used. Consider adding the '{default_set_name}' key to the 'obs_groups' dictionary for" - " clarity. This behavior will be removed in a future version." - ) - - # print the final parsed observation sets - print("-" * 80) - print("Resolved observation sets: ") - for set_name, groups in obs_groups.items(): - print("\t", set_name, ": ", groups) - print("-" * 80) - - return obs_groups diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/wandb_utils.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/wandb_utils.py deleted file mode 100644 index 243e82d..0000000 --- a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/wandb_utils.py +++ /dev/null @@ -1,87 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -from __future__ import annotations - -import os -from dataclasses import asdict -from torch.utils.tensorboard import SummaryWriter - -try: - import wandb -except ModuleNotFoundError: - raise ModuleNotFoundError("Wandb is required to log to Weights and Biases.") - - -class WandbSummaryWriter(SummaryWriter): - """Summary writer for Weights and Biases.""" - - def __init__(self, log_dir: str, flush_secs: int, cfg): - super().__init__(log_dir, flush_secs) - - # Get the run name - run_name = os.path.split(log_dir)[-1] - - try: - project = cfg["wandb_project"] - except KeyError: - raise KeyError("Please specify wandb_project in the runner config, e.g. legged_gym.") - - try: - entity = os.environ["WANDB_USERNAME"] - except KeyError: - entity = None - - # Initialize wandb - wandb.init(project=project, entity=entity, name=run_name) - - # Add log directory to wandb - wandb.config.update({"log_dir": log_dir}) - - self.name_map = { - "Train/mean_reward/time": "Train/mean_reward_time", - "Train/mean_episode_length/time": "Train/mean_episode_length_time", - } - - def store_config(self, env_cfg, runner_cfg, alg_cfg, policy_cfg): - wandb.config.update({"runner_cfg": runner_cfg}) - wandb.config.update({"policy_cfg": policy_cfg}) - wandb.config.update({"alg_cfg": alg_cfg}) - try: - wandb.config.update({"env_cfg": env_cfg.to_dict()}) - except Exception: - wandb.config.update({"env_cfg": asdict(env_cfg)}) - - def add_scalar(self, tag, scalar_value, global_step=None, walltime=None, new_style=False): - super().add_scalar( - tag, - scalar_value, - global_step=global_step, - walltime=walltime, - new_style=new_style, - ) - wandb.log({self._map_path(tag): scalar_value}, step=global_step) - - def stop(self): - wandb.finish() - - def log_config(self, env_cfg, runner_cfg, alg_cfg, policy_cfg): - self.store_config(env_cfg, runner_cfg, alg_cfg, policy_cfg) - - def save_model(self, model_path, iter): - wandb.save(model_path, base_path=os.path.dirname(model_path)) - - def save_file(self, path, iter=None): - wandb.save(path, base_path=os.path.dirname(path)) - - """ - Private methods. - """ - - def _map_path(self, path): - if path in self.name_map: - return self.name_map[path] - else: - return path diff --git a/rsl_rl/build/lib/rsl_rl/env/__init__.py b/rsl_rl/build/lib/rsl_rl/env/__init__.py deleted file mode 100644 index ab7c056..0000000 --- a/rsl_rl/build/lib/rsl_rl/env/__init__.py +++ /dev/null @@ -1,10 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -"""Submodule defining the environment definitions.""" - -from .vec_env import VecEnv - -__all__ = ["VecEnv"] diff --git a/rsl_rl/build/lib/rsl_rl/env/vec_env.py b/rsl_rl/build/lib/rsl_rl/env/vec_env.py deleted file mode 100644 index ae471cf..0000000 --- a/rsl_rl/build/lib/rsl_rl/env/vec_env.py +++ /dev/null @@ -1,113 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -from __future__ import annotations - -import torch -from abc import ABC, abstractmethod -from tensordict import TensorDict - - -class VecEnv(ABC): - """Abstract class for a vectorized environment. - - The vectorized environment is a collection of environments that are synchronized. This means that - the same type of action is applied to all environments and the same type of observation is returned from all - environments. - """ - - num_envs: int - """Number of environments.""" - - num_actions: int - """Number of actions.""" - - max_episode_length: int | torch.Tensor - - max_episode_length_s: float - """Maximum episode length. - - The maximum episode length can be a scalar or a tensor. If it is a scalar, it is the same for all environments. - If it is a tensor, it is the maximum episode length for each environment. This is useful for dynamic episode - lengths. - """ - - episode_length_buf: torch.Tensor - """Buffer for current episode lengths.""" - - device: torch.device | str - """Device to use.""" - - cfg: dict | object - """Configuration object.""" - - reset_env_ids: torch.Tensor | None = None - - contact_phase: torch.Tensor | None = None - """ - Operations. - """ - - @abstractmethod - def get_observations(self) -> TensorDict: - """Return the current observations. - - Returns: - observations (TensorDict): Observations from the environment. - """ - raise NotImplementedError - - @abstractmethod - def get_amp_observations(self) -> TensorDict: - """Return the current AMP observations. - - Returns: - observations (TensorDict): Observations from the environment. - """ - raise NotImplementedError - - @abstractmethod - def step(self, actions: torch.Tensor) -> tuple[TensorDict, torch.Tensor, torch.Tensor, dict]: - """Apply input action to the environment. - - Args: - actions (torch.Tensor): Input actions to apply. Shape: (num_envs, num_actions) - - Returns: - observations (TensorDict): Observations from the environment. - rewards (torch.Tensor): Rewards from the environment. Shape: (num_envs,) - dones (torch.Tensor): Done flags from the environment. Shape: (num_envs,) - extras (dict): Extra information from the environment. - - Observations: - - The observations TensorDict usually contains multiple observation groups. The `obs_groups` - dictionary of the runner configuration specifies which observation groups are used for which - purpose, i.e., it maps the available observation groups to observation sets. The observation sets - (keys of the `obs_groups` dictionary) currently used by rsl_rl are: - - - "policy": Specified observation groups are used as input to the actor/student network. - - "critic": Specified observation groups are used as input to the critic network. - - "teacher": Specified observation groups are used as input to the teacher network. - - "rnd_state": Specified observation groups are used as input to the RND network. - - Incomplete or incorrect configurations are handled in the `resolve_obs_groups()` function in - `rsl_rl/utils/utils.py`. - - Extras: - - The extras dictionary includes metrics such as the episode reward, episode length, etc. The following - dictionary keys are used by rsl_rl: - - - "time_outs" (torch.Tensor): Timeouts for the environments. These correspond to terminations that - happen due to time limits and not due to the environment reaching a terminal state. This is useful - for environments that have a fixed episode length. - - - "log" (dict[str, float | torch.Tensor]): Additional information for logging and debugging purposes. - The key should be a string and start with "/" for namespacing. The value can be a scalar or a - tensor. If it is a tensor, the mean of the tensor is used for logging. - """ - raise NotImplementedError - \ No newline at end of file diff --git a/rsl_rl/build/lib/rsl_rl/modules/__init__.py b/rsl_rl/build/lib/rsl_rl/modules/__init__.py deleted file mode 100644 index 9afe1a2..0000000 --- a/rsl_rl/build/lib/rsl_rl/modules/__init__.py +++ /dev/null @@ -1,21 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -"""Definitions for neural-network components for RL-agents.""" - -from .actor_critic import ActorCritic -from .actor_critic_recurrent import ActorCriticRecurrent -from .rnd import * -from .student_teacher import StudentTeacher -from .student_teacher_recurrent import StudentTeacherRecurrent -from .symmetry import * -from .discriminator_multi import DiscriminatorMulti -__all__ = [ - "ActorCritic", - "ActorCriticRecurrent", - "StudentTeacher", - "StudentTeacherRecurrent", - "DiscriminatorMulti", -] diff --git a/rsl_rl/build/lib/rsl_rl/modules/actor_critic.py b/rsl_rl/build/lib/rsl_rl/modules/actor_critic.py deleted file mode 100644 index 0efc36b..0000000 --- a/rsl_rl/build/lib/rsl_rl/modules/actor_critic.py +++ /dev/null @@ -1,195 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -from __future__ import annotations - -import torch -import torch.nn as nn -from torch.distributions import Normal - -from rsl_rl.networks import MLP, EmpiricalNormalization - - -class ActorCritic(nn.Module): - is_recurrent = False - - def __init__( - self, - obs, - obs_groups, - num_actions, - actor_obs_normalization=False, - critic_obs_normalization=False, - actor_hidden_dims=[256, 256, 256], - critic_hidden_dims=[256, 256, 256], - activation="elu", - init_noise_std=1.0, - noise_std_type: str = "scalar", - state_dependent_std=False, - **kwargs, - ): - if kwargs: - print( - "ActorCritic.__init__ got unexpected arguments, which will be ignored: " - + str([key for key in kwargs.keys()]) - ) - super().__init__() - - # get the observation dimensions - self.obs_groups = obs_groups - num_actor_obs = 0 - for obs_group in obs_groups["policy"]: - assert len(obs[obs_group].shape) == 2, "The ActorCritic module only supports 1D observations." - num_actor_obs += obs[obs_group].shape[-1] - num_critic_obs = 0 - for obs_group in obs_groups["critic"]: - assert len(obs[obs_group].shape) == 2, "The ActorCritic module only supports 1D observations." - num_critic_obs += obs[obs_group].shape[-1] - - self.state_dependent_std = state_dependent_std - # actor - if self.state_dependent_std: - self.actor = MLP(num_actor_obs, [2, num_actions], actor_hidden_dims, activation) - else: - self.actor = MLP(num_actor_obs, num_actions, actor_hidden_dims, activation) - # actor observation normalization - self.actor_obs_normalization = actor_obs_normalization - if actor_obs_normalization: - self.actor_obs_normalizer = EmpiricalNormalization(num_actor_obs) - else: - self.actor_obs_normalizer = torch.nn.Identity() - print(f"Actor MLP: {self.actor}") - - # critic - self.critic = MLP(num_critic_obs, 1, critic_hidden_dims, activation) - # critic observation normalization - self.critic_obs_normalization = critic_obs_normalization - if critic_obs_normalization: - self.critic_obs_normalizer = EmpiricalNormalization(num_critic_obs) - else: - self.critic_obs_normalizer = torch.nn.Identity() - print(f"Critic MLP: {self.critic}") - - # Action noise - self.noise_std_type = noise_std_type - if self.state_dependent_std: - torch.nn.init.zeros_(self.actor[-2].weight[num_actions:]) - if self.noise_std_type == "scalar": - torch.nn.init.constant_(self.actor[-2].bias[num_actions:], init_noise_std) - elif self.noise_std_type == "log": - torch.nn.init.constant_( - self.actor[-2].bias[num_actions:], torch.log(torch.tensor(init_noise_std + 1e-7)) - ) - else: - raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") - else: - if self.noise_std_type == "scalar": - self.std = nn.Parameter(init_noise_std * torch.ones(num_actions)) - elif self.noise_std_type == "log": - self.log_std = nn.Parameter(torch.log(init_noise_std * torch.ones(num_actions))) - else: - raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") - - # Action distribution (populated in update_distribution) - self.distribution = None - # disable args validation for speedup - Normal.set_default_validate_args(False) - - def reset(self, dones=None): - pass - - def forward(self): - raise NotImplementedError - - @property - def action_mean(self): - return self.distribution.mean - - @property - def action_std(self): - return self.distribution.stddev - - @property - def entropy(self): - return self.distribution.entropy().sum(dim=-1) - - def update_distribution(self, obs): - if self.state_dependent_std: - # compute mean and standard deviation - mean_and_std = self.actor(obs) - if self.noise_std_type == "scalar": - mean, std = torch.unbind(mean_and_std, dim=-2) - elif self.noise_std_type == "log": - mean, log_std = torch.unbind(mean_and_std, dim=-2) - std = torch.exp(log_std) - else: - raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") - else: - # compute mean - mean = self.actor(obs) - # compute standard deviation - if self.noise_std_type == "scalar": - std = self.std.expand_as(mean) - elif self.noise_std_type == "log": - std = torch.exp(self.log_std).expand_as(mean) - else: - raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") - # create distribution - self.distribution = Normal(mean, std) - - def act(self, obs, **kwargs): - obs = self.get_actor_obs(obs) - obs = self.actor_obs_normalizer(obs) - self.update_distribution(obs) - return self.distribution.sample() - - def act_inference(self, obs): - obs = self.get_actor_obs(obs) - obs = self.actor_obs_normalizer(obs) - return self.actor(obs) - - def evaluate(self, obs, **kwargs): - obs = self.get_critic_obs(obs) - obs = self.critic_obs_normalizer(obs) - return self.critic(obs) - - def get_actor_obs(self, obs): - obs_list = [] - for obs_group in self.obs_groups["policy"]: - obs_list.append(obs[obs_group]) - return torch.cat(obs_list, dim=-1) - - def get_critic_obs(self, obs): - obs_list = [] - for obs_group in self.obs_groups["critic"]: - obs_list.append(obs[obs_group]) - return torch.cat(obs_list, dim=-1) - - def get_actions_log_prob(self, actions): - return self.distribution.log_prob(actions).sum(dim=-1) - - def update_normalization(self, obs): - if self.actor_obs_normalization: - actor_obs = self.get_actor_obs(obs) - self.actor_obs_normalizer.update(actor_obs) - if self.critic_obs_normalization: - critic_obs = self.get_critic_obs(obs) - self.critic_obs_normalizer.update(critic_obs) - - def load_state_dict(self, state_dict, strict=True): - """Load the parameters of the actor-critic model. - - Args: - state_dict (dict): State dictionary of the model. - strict (bool): Whether to strictly enforce that the keys in state_dict match the keys returned by this - module's state_dict() function. - - Returns: - bool: Whether this training resumes a previous training. This flag is used by the `load()` function of - `OnPolicyRunner` to determine how to load further parameters (relevant for, e.g., distillation). - """ - - super().load_state_dict(state_dict, strict=strict) - return True # training resumes diff --git a/rsl_rl/build/lib/rsl_rl/modules/actor_critic_recurrent.py b/rsl_rl/build/lib/rsl_rl/modules/actor_critic_recurrent.py deleted file mode 100644 index bba46ca..0000000 --- a/rsl_rl/build/lib/rsl_rl/modules/actor_critic_recurrent.py +++ /dev/null @@ -1,218 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -from __future__ import annotations - -import torch -import torch.nn as nn -import warnings -from torch.distributions import Normal - -from rsl_rl.networks import MLP, EmpiricalNormalization, Memory - - -class ActorCriticRecurrent(nn.Module): - is_recurrent = True - - def __init__( - self, - obs, - obs_groups, - num_actions, - actor_obs_normalization=False, - critic_obs_normalization=False, - actor_hidden_dims=[256, 256, 256], - critic_hidden_dims=[256, 256, 256], - activation="elu", - init_noise_std=1.0, - noise_std_type: str = "scalar", - state_dependent_std=False, - rnn_type="lstm", - rnn_hidden_dim=256, - rnn_num_layers=1, - **kwargs, - ): - if "rnn_hidden_size" in kwargs: - warnings.warn( - "The argument `rnn_hidden_size` is deprecated and will be removed in a future version. " - "Please use `rnn_hidden_dim` instead.", - DeprecationWarning, - ) - if rnn_hidden_dim == 256: # Only override if the new argument is at its default - rnn_hidden_dim = kwargs.pop("rnn_hidden_size") - if kwargs: - print( - "ActorCriticRecurrent.__init__ got unexpected arguments, which will be ignored: " + str(kwargs.keys()), - ) - super().__init__() - - # get the observation dimensions - self.obs_groups = obs_groups - num_actor_obs = 0 - for obs_group in obs_groups["policy"]: - assert len(obs[obs_group].shape) == 2, "The ActorCriticRecurrent module only supports 1D observations." - num_actor_obs += obs[obs_group].shape[-1] - num_critic_obs = 0 - for obs_group in obs_groups["critic"]: - assert len(obs[obs_group].shape) == 2, "The ActorCriticRecurrent module only supports 1D observations." - num_critic_obs += obs[obs_group].shape[-1] - - self.state_dependent_std = state_dependent_std - # actor - self.memory_a = Memory(num_actor_obs, type=rnn_type, num_layers=rnn_num_layers, hidden_size=rnn_hidden_dim) - if self.state_dependent_std: - self.actor = MLP(rnn_hidden_dim, [2, num_actions], actor_hidden_dims, activation) - else: - self.actor = MLP(rnn_hidden_dim, num_actions, actor_hidden_dims, activation) - - # actor observation normalization - self.actor_obs_normalization = actor_obs_normalization - if actor_obs_normalization: - self.actor_obs_normalizer = EmpiricalNormalization(num_actor_obs) - else: - self.actor_obs_normalizer = torch.nn.Identity() - print(f"Actor RNN: {self.memory_a}") - print(f"Actor MLP: {self.actor}") - - # critic - self.memory_c = Memory(num_critic_obs, type=rnn_type, num_layers=rnn_num_layers, hidden_size=rnn_hidden_dim) - self.critic = MLP(rnn_hidden_dim, 1, critic_hidden_dims, activation) - # critic observation normalization - self.critic_obs_normalization = critic_obs_normalization - if critic_obs_normalization: - self.critic_obs_normalizer = EmpiricalNormalization(num_critic_obs) - else: - self.critic_obs_normalizer = torch.nn.Identity() - print(f"Critic RNN: {self.memory_c}") - print(f"Critic MLP: {self.critic}") - - # Action noise - self.noise_std_type = noise_std_type - if self.state_dependent_std: - torch.nn.init.zeros_(self.actor[-2].weight[num_actions:]) - if self.noise_std_type == "scalar": - torch.nn.init.constant_(self.actor[-2].bias[num_actions:], init_noise_std) - elif self.noise_std_type == "log": - torch.nn.init.constant_( - self.actor[-2].bias[num_actions:], torch.log(torch.tensor(init_noise_std + 1e-7)) - ) - else: - raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") - else: - if self.noise_std_type == "scalar": - self.std = nn.Parameter(init_noise_std * torch.ones(num_actions)) - elif self.noise_std_type == "log": - self.log_std = nn.Parameter(torch.log(init_noise_std * torch.ones(num_actions))) - else: - raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") - - # Action distribution (populated in update_distribution) - self.distribution = None - # disable args validation for speedup - Normal.set_default_validate_args(False) - - @property - def action_mean(self): - return self.distribution.mean - - @property - def action_std(self): - return self.distribution.stddev - - @property - def entropy(self): - return self.distribution.entropy().sum(dim=-1) - - def reset(self, dones=None): - self.memory_a.reset(dones) - self.memory_c.reset(dones) - - def forward(self): - raise NotImplementedError - - def update_distribution(self, obs): - if self.state_dependent_std: - # compute mean and standard deviation - mean_and_std = self.actor(obs) - if self.noise_std_type == "scalar": - mean, std = torch.unbind(mean_and_std, dim=-2) - elif self.noise_std_type == "log": - mean, log_std = torch.unbind(mean_and_std, dim=-2) - std = torch.exp(log_std) - else: - raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") - else: - # compute mean - mean = self.actor(obs) - # compute standard deviation - if self.noise_std_type == "scalar": - std = self.std.expand_as(mean) - elif self.noise_std_type == "log": - std = torch.exp(self.log_std).expand_as(mean) - else: - raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") - # create distribution - self.distribution = Normal(mean, std) - - def act(self, obs, masks=None, hidden_states=None): - obs = self.get_actor_obs(obs) - obs = self.actor_obs_normalizer(obs) - out_mem = self.memory_a(obs, masks, hidden_states).squeeze(0) - self.update_distribution(out_mem) - return self.distribution.sample() - - def act_inference(self, obs): - obs = self.get_actor_obs(obs) - obs = self.actor_obs_normalizer(obs) - out_mem = self.memory_a(obs).squeeze(0) - return self.actor(out_mem) - - def evaluate(self, obs, masks=None, hidden_states=None): - obs = self.get_critic_obs(obs) - obs = self.critic_obs_normalizer(obs) - out_mem = self.memory_c(obs, masks, hidden_states).squeeze(0) - return self.critic(out_mem) - - def get_actor_obs(self, obs): - obs_list = [] - for obs_group in self.obs_groups["policy"]: - obs_list.append(obs[obs_group]) - return torch.cat(obs_list, dim=-1) - - def get_critic_obs(self, obs): - obs_list = [] - for obs_group in self.obs_groups["critic"]: - obs_list.append(obs[obs_group]) - return torch.cat(obs_list, dim=-1) - - def get_actions_log_prob(self, actions): - return self.distribution.log_prob(actions).sum(dim=-1) - - def get_hidden_states(self): - return self.memory_a.hidden_states, self.memory_c.hidden_states - - def update_normalization(self, obs): - if self.actor_obs_normalization: - actor_obs = self.get_actor_obs(obs) - self.actor_obs_normalizer.update(actor_obs) - if self.critic_obs_normalization: - critic_obs = self.get_critic_obs(obs) - self.critic_obs_normalizer.update(critic_obs) - - def load_state_dict(self, state_dict, strict=True): - """Load the parameters of the actor-critic model. - - Args: - state_dict (dict): State dictionary of the model. - strict (bool): Whether to strictly enforce that the keys in state_dict match the keys returned by this - module's state_dict() function. - - Returns: - bool: Whether this training resumes a previous training. This flag is used by the `load()` function of - `OnPolicyRunner` to determine how to load further parameters (relevant for, e.g., distillation). - """ - - super().load_state_dict(state_dict, strict=strict) - return True diff --git a/rsl_rl/build/lib/rsl_rl/modules/discriminator_multi.py b/rsl_rl/build/lib/rsl_rl/modules/discriminator_multi.py deleted file mode 100644 index 7899d23..0000000 --- a/rsl_rl/build/lib/rsl_rl/modules/discriminator_multi.py +++ /dev/null @@ -1,102 +0,0 @@ -import torch -import torch.nn as nn -from torch import autograd -import torch.nn.utils.spectral_norm as spectral_norm - - -class DiscriminatorMulti(nn.Module): - def __init__( - self, state_dim, amp_reward_coef, hidden_layer_sizes, device, - num_frames=2, task_reward_lerp=0.0, use_lerp=True): - super(DiscriminatorMulti, self).__init__() - - self.device = device - self.state_dim = state_dim - self.use_lerp = use_lerp - self.num_frames = num_frames # 存储帧数参数 - - self.amp_reward_coef = amp_reward_coef - amp_layers = [] - - curr_in_dim = state_dim * num_frames - for hidden_dim in hidden_layer_sizes: - amp_layers.append(spectral_norm(nn.Linear(curr_in_dim, hidden_dim))) - amp_layers.append(nn.ReLU()) - curr_in_dim = hidden_dim - self.trunk = nn.Sequential(*amp_layers).to(device) - self.amp_linear = spectral_norm(nn.Linear(hidden_layer_sizes[-1], 1)).to(device) - - self.trunk.train() - self.amp_linear.train() - - self.task_reward_lerp = task_reward_lerp - - def forward(self, x): - h = self.trunk(x) - d = self.amp_linear(h) - return d - - def compute_grad_pen(self, - expert_states, # 改为接收多帧状态列表 - lambda_=10): - # 将多帧状态沿最后一个维度拼接 - expert_data = expert_states.flatten(1) - expert_data.requires_grad = True - - disc = self.amp_linear(self.trunk(expert_data)) - ones = torch.ones(disc.size(), device=disc.device) - grad = autograd.grad( - outputs=disc, inputs=expert_data, - grad_outputs=ones, create_graph=True, - retain_graph=True, only_inputs=True)[0] - - # Enforce that the grad norm approaches 0. - grad_pen = lambda_ * (grad.norm(2, dim=1) - 0).pow(2).mean() - return grad_pen - - - def get_disc_weights(self): - weights = [] - for m in self.trunk.modules(): - if isinstance(m, nn.Linear): - weights.append(torch.flatten(m.weight)) - - weights.append(torch.flatten(self.amp_linear.weight)) - return weights - - def get_disc_logit_weights(self): - return torch.flatten(self.amp_linear.weight) - - def predict_amp_reward( - self, states, # 改为接收多帧状态列表 - task_reward, normalizer=None): - """ - states: torch.Tensor, shape=(num_envs, num_frames, state_dim) - task_reward: torch.Tensor, shape=(num_envs, 1) - """ - # import ipdb; ipdb.set_trace() - with torch.no_grad(): - self.eval() - if normalizer is not None: - # 对每一帧状态进行归一化 - states = normalizer.normalize_torch(states, self.device) - - # 拼接多帧状态 - state_cat = states.flatten(1) - d = self.amp_linear(self.trunk(state_cat)) - disc_reward = self.amp_reward_coef * torch.clamp(1 - (1/4) * torch.square(d - 1), min=0) - - if self.use_lerp: - if self.task_reward_lerp > 0: - reward = self._lerp_reward(disc_reward, task_reward.unsqueeze(-1)) - self.train() - return reward.squeeze(), d, disc_reward.squeeze() * (1.0 - self.task_reward_lerp) - else: - disc_reward *= 0.02 - reward = task_reward.unsqueeze(-1) + disc_reward - self.train() - return reward.squeeze(), d, disc_reward.squeeze() - - def _lerp_reward(self, disc_r, task_r): - r = (1.0 - self.task_reward_lerp) * disc_r + self.task_reward_lerp * task_r - return r \ No newline at end of file diff --git a/rsl_rl/build/lib/rsl_rl/modules/rnd.py b/rsl_rl/build/lib/rsl_rl/modules/rnd.py deleted file mode 100644 index 8e65c43..0000000 --- a/rsl_rl/build/lib/rsl_rl/modules/rnd.py +++ /dev/null @@ -1,209 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -from __future__ import annotations - -import torch -import torch.nn as nn - -from rsl_rl.networks import MLP, EmpiricalDiscountedVariationNormalization, EmpiricalNormalization - - -class RandomNetworkDistillation(nn.Module): - """Implementation of Random Network Distillation (RND) [1] - - References: - .. [1] Burda, Yuri, et al. "Exploration by random network distillation." arXiv preprint arXiv:1810.12894 (2018). - """ - - def __init__( - self, - num_states: int, - obs_groups: dict, - num_outputs: int, - predictor_hidden_dims: list[int], - target_hidden_dims: list[int], - activation: str = "elu", - weight: float = 0.0, - state_normalization: bool = False, - reward_normalization: bool = False, - device: str = "cpu", - weight_schedule: dict | None = None, - ): - """Initialize the RND module. - - - If :attr:`state_normalization` is True, then the input state is normalized using an Empirical Normalization layer. - - If :attr:`reward_normalization` is True, then the intrinsic reward is normalized using an Empirical Discounted - Variation Normalization layer. - - .. note:: - If the hidden dimensions are -1 in the predictor and target networks configuration, then the number of states - is used as the hidden dimension. - - Args: - num_states: Number of states/inputs to the predictor and target networks. - num_outputs: Number of outputs (embedding size) of the predictor and target networks. - predictor_hidden_dims: List of hidden dimensions of the predictor network. - target_hidden_dims: List of hidden dimensions of the target network. - activation: Activation function. Defaults to "elu". - weight: Scaling factor of the intrinsic reward. Defaults to 0.0. - state_normalization: Whether to normalize the input state. Defaults to False. - reward_normalization: Whether to normalize the intrinsic reward. Defaults to False. - device: Device to use. Defaults to "cpu". - weight_schedule: The type of schedule to use for the RND weight parameter. - Defaults to None, in which case the weight parameter is constant. - It is a dictionary with the following keys: - - - "mode": The type of schedule to use for the RND weight parameter. - - "constant": Constant weight schedule. - - "step": Step weight schedule. - - "linear": Linear weight schedule. - - For the "step" weight schedule, the following parameters are required: - - - "final_step": The step at which the weight parameter is set to the final value. - - "final_value": The final value of the weight parameter. - - For the "linear" weight schedule, the following parameters are required: - - "initial_step": The step at which the weight parameter is set to the initial value. - - "final_step": The step at which the weight parameter is set to the final value. - - "final_value": The final value of the weight parameter. - """ - # initialize parent class - super().__init__() - - # Store parameters - self.num_states = num_states - self.obs_groups = obs_groups - self.num_outputs = num_outputs - self.initial_weight = weight - self.device = device - self.state_normalization = state_normalization - self.reward_normalization = reward_normalization - - # Normalization of input gates - if state_normalization: - self.state_normalizer = EmpiricalNormalization(shape=[self.num_states], until=1.0e8).to(self.device) - else: - self.state_normalizer = torch.nn.Identity() - # Normalization of intrinsic reward - if reward_normalization: - self.reward_normalizer = EmpiricalDiscountedVariationNormalization(shape=[], until=1.0e8).to(self.device) - else: - self.reward_normalizer = torch.nn.Identity() - - # counter for the number of updates - self.update_counter = 0 - - # resolve weight schedule - if weight_schedule is not None: - self.weight_scheduler_params = weight_schedule - self.weight_scheduler = getattr(self, f"_{weight_schedule['mode']}_weight_schedule") - else: - self.weight_scheduler = None - # Create network architecture - self.predictor = MLP(num_states, num_outputs, predictor_hidden_dims, activation).to(self.device) - self.target = MLP(num_states, num_outputs, target_hidden_dims, activation).to(self.device) - - # make target network not trainable - self.target.eval() - - def get_intrinsic_reward(self, obs) -> torch.Tensor: - # Note: the counter is updated number of env steps per learning iteration - self.update_counter += 1 - # Extract the rnd state from the observation - rnd_state = self.get_rnd_state(obs) - rnd_state = self.state_normalizer(rnd_state) - # Obtain the embedding of the rnd state from the target and predictor networks - target_embedding = self.target(rnd_state).detach() - predictor_embedding = self.predictor(rnd_state).detach() - # Compute the intrinsic reward as the distance between the embeddings - intrinsic_reward = torch.linalg.norm(target_embedding - predictor_embedding, dim=1) - # Normalize intrinsic reward - intrinsic_reward = self.reward_normalizer(intrinsic_reward) - - # Check the weight schedule - if self.weight_scheduler is not None: - self.weight = self.weight_scheduler(step=self.update_counter, **self.weight_scheduler_params) - else: - self.weight = self.initial_weight - # Scale intrinsic reward - intrinsic_reward *= self.weight - - return intrinsic_reward - - def forward(self, *args, **kwargs): - raise RuntimeError("Forward method is not implemented. Use get_intrinsic_reward instead.") - - def train(self, mode: bool = True): - # sets module into training mode - self.predictor.train(mode) - if self.state_normalization: - self.state_normalizer.train(mode) - if self.reward_normalization: - self.reward_normalizer.train(mode) - return self - - def eval(self): - return self.train(False) - - def get_rnd_state(self, obs): - obs_list = [] - for obs_group in self.obs_groups["rnd_state"]: - obs_list.append(obs[obs_group]) - return torch.cat(obs_list, dim=-1) - - def update_normalization(self, obs): - # Normalize the state - if self.state_normalization: - rnd_state = self.get_rnd_state(obs) - self.state_normalizer.update(rnd_state) - - """ - Different weight schedules. - """ - - def _constant_weight_schedule(self, step: int, **kwargs): - return self.initial_weight - - def _step_weight_schedule(self, step: int, final_step: int, final_value: float, **kwargs): - return self.initial_weight if step < final_step else final_value - - def _linear_weight_schedule(self, step: int, initial_step: int, final_step: int, final_value: float, **kwargs): - if step < initial_step: - return self.initial_weight - elif step > final_step: - return final_value - else: - return self.initial_weight + (final_value - self.initial_weight) * (step - initial_step) / ( - final_step - initial_step - ) - - -def resolve_rnd_config(alg_cfg, obs, obs_groups, env): - """Resolve the RND configuration. - - Args: - alg_cfg: The algorithm configuration dictionary. - obs: The observation dictionary. - obs_groups: The observation groups dictionary. - env: The environment. - - Returns: - The resolved algorithm configuration dictionary. - """ - # resolve dimension of rnd gated state - if "rnd_cfg" in alg_cfg and alg_cfg["rnd_cfg"] is not None: - # get dimension of rnd gated state - num_rnd_state = 0 - for obs_group in obs_groups["rnd_state"]: - assert len(obs[obs_group].shape) == 2, "The RND module only supports 1D observations." - num_rnd_state += obs[obs_group].shape[-1] - # add rnd gated state to config - alg_cfg["rnd_cfg"]["num_states"] = num_rnd_state - alg_cfg["rnd_cfg"]["obs_groups"] = obs_groups - # scale down the rnd weight with timestep - alg_cfg["rnd_cfg"]["weight"] *= env.unwrapped.step_dt - return alg_cfg diff --git a/rsl_rl/build/lib/rsl_rl/modules/student_teacher.py b/rsl_rl/build/lib/rsl_rl/modules/student_teacher.py deleted file mode 100644 index 6bf1380..0000000 --- a/rsl_rl/build/lib/rsl_rl/modules/student_teacher.py +++ /dev/null @@ -1,206 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -from __future__ import annotations - -import torch -import torch.nn as nn -from torch.distributions import Normal - -from rsl_rl.networks import MLP, EmpiricalNormalization - - -class StudentTeacher(nn.Module): - is_recurrent = False - - def __init__( - self, - obs, - obs_groups, - num_actions, - student_obs_normalization=False, - teacher_obs_normalization=False, - student_hidden_dims=[256, 256, 256], - teacher_hidden_dims=[256, 256, 256], - activation="elu", - init_noise_std=0.1, - noise_std_type: str = "scalar", - **kwargs, - ): - if kwargs: - print( - "StudentTeacher.__init__ got unexpected arguments, which will be ignored: " - + str([key for key in kwargs.keys()]) - ) - super().__init__() - - self.loaded_teacher = False # indicates if teacher has been loaded - - # get the observation dimensions - self.obs_groups = obs_groups - num_student_obs = 0 - for obs_group in obs_groups["policy"]: - assert len(obs[obs_group].shape) == 2, "The StudentTeacher module only supports 1D observations." - num_student_obs += obs[obs_group].shape[-1] - num_teacher_obs = 0 - for obs_group in obs_groups["teacher"]: - assert len(obs[obs_group].shape) == 2, "The StudentTeacher module only supports 1D observations." - num_teacher_obs += obs[obs_group].shape[-1] - - # student - self.student = MLP(num_student_obs, num_actions, student_hidden_dims, activation) - - # student observation normalization - self.student_obs_normalization = student_obs_normalization - if student_obs_normalization: - self.student_obs_normalizer = EmpiricalNormalization(num_student_obs) - else: - self.student_obs_normalizer = torch.nn.Identity() - - print(f"Student MLP: {self.student}") - - # teacher - self.teacher = MLP(num_teacher_obs, num_actions, teacher_hidden_dims, activation) - self.teacher.eval() - - # teacher observation normalization - self.teacher_obs_normalization = teacher_obs_normalization - if teacher_obs_normalization: - self.teacher_obs_normalizer = EmpiricalNormalization(num_teacher_obs) - else: - self.teacher_obs_normalizer = torch.nn.Identity() - - print(f"Teacher MLP: {self.teacher}") - - # action noise - self.noise_std_type = noise_std_type - if self.noise_std_type == "scalar": - self.std = nn.Parameter(init_noise_std * torch.ones(num_actions)) - elif self.noise_std_type == "log": - self.log_std = nn.Parameter(torch.log(init_noise_std * torch.ones(num_actions))) - else: - raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") - - # action distribution (populated in update_distribution) - self.distribution = None - # disable args validation for speedup - Normal.set_default_validate_args(False) - - def reset(self, dones=None, hidden_states=None): - pass - - def forward(self): - raise NotImplementedError - - @property - def action_mean(self): - return self.distribution.mean - - @property - def action_std(self): - return self.distribution.stddev - - @property - def entropy(self): - return self.distribution.entropy().sum(dim=-1) - - def update_distribution(self, obs): - # compute mean - mean = self.student(obs) - # compute standard deviation - if self.noise_std_type == "scalar": - std = self.std.expand_as(mean) - elif self.noise_std_type == "log": - std = torch.exp(self.log_std).expand_as(mean) - else: - raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") - # create distribution - self.distribution = Normal(mean, std) - - def act(self, obs): - obs = self.get_student_obs(obs) - obs = self.student_obs_normalizer(obs) - self.update_distribution(obs) - return self.distribution.sample() - - def act_inference(self, obs): - obs = self.get_student_obs(obs) - obs = self.student_obs_normalizer(obs) - return self.student(obs) - - def evaluate(self, obs): - obs = self.get_teacher_obs(obs) - obs = self.teacher_obs_normalizer(obs) - with torch.no_grad(): - return self.teacher(obs) - - def get_student_obs(self, obs): - obs_list = [] - for obs_group in self.obs_groups["policy"]: - obs_list.append(obs[obs_group]) - return torch.cat(obs_list, dim=-1) - - def get_teacher_obs(self, obs): - obs_list = [] - for obs_group in self.obs_groups["teacher"]: - obs_list.append(obs[obs_group]) - return torch.cat(obs_list, dim=-1) - - def get_hidden_states(self): - return None - - def detach_hidden_states(self, dones=None): - pass - - def train(self, mode=True): - super().train(mode) - # make sure teacher is in eval mode - self.teacher.eval() - self.teacher_obs_normalizer.eval() - - def update_normalization(self, obs): - if self.student_obs_normalization: - student_obs = self.get_student_obs(obs) - self.student_obs_normalizer.update(student_obs) - - def load_state_dict(self, state_dict, strict=True): - """Load the parameters of the student and teacher networks. - - Args: - state_dict (dict): State dictionary of the model. - strict (bool): Whether to strictly enforce that the keys in state_dict match the keys returned by this - module's state_dict() function. - - Returns: - bool: Whether this training resumes a previous training. This flag is used by the `load()` function of - `OnPolicyRunner` to determine how to load further parameters. - """ - - # check if state_dict contains teacher and student or just teacher parameters - if any("actor" in key for key in state_dict.keys()): # loading parameters from rl training - # rename keys to match teacher and remove critic parameters - teacher_state_dict = {} - teacher_obs_normalizer_state_dict = {} - for key, value in state_dict.items(): - if "actor." in key: - teacher_state_dict[key.replace("actor.", "")] = value - if "actor_obs_normalizer." in key: - teacher_obs_normalizer_state_dict[key.replace("actor_obs_normalizer.", "")] = value - self.teacher.load_state_dict(teacher_state_dict, strict=strict) - self.teacher_obs_normalizer.load_state_dict(teacher_obs_normalizer_state_dict, strict=strict) - # set flag for successfully loading the parameters - self.loaded_teacher = True - self.teacher.eval() - self.teacher_obs_normalizer.eval() - return False # training does not resume - elif any("student" in key for key in state_dict.keys()): # loading parameters from distillation training - super().load_state_dict(state_dict, strict=strict) - # set flag for successfully loading the parameters - self.loaded_teacher = True - self.teacher.eval() - self.teacher_obs_normalizer.eval() - return True # training resumes - else: - raise ValueError("state_dict does not contain student or teacher parameters") diff --git a/rsl_rl/build/lib/rsl_rl/modules/student_teacher_recurrent.py b/rsl_rl/build/lib/rsl_rl/modules/student_teacher_recurrent.py deleted file mode 100644 index 964a2dc..0000000 --- a/rsl_rl/build/lib/rsl_rl/modules/student_teacher_recurrent.py +++ /dev/null @@ -1,249 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -from __future__ import annotations - -import torch -import torch.nn as nn -import warnings -from torch.distributions import Normal - -from rsl_rl.networks import MLP, EmpiricalNormalization, Memory - - -class StudentTeacherRecurrent(nn.Module): - is_recurrent = True - - def __init__( - self, - obs, - obs_groups, - num_actions, - student_obs_normalization=False, - teacher_obs_normalization=False, - student_hidden_dims=[256, 256, 256], - teacher_hidden_dims=[256, 256, 256], - activation="elu", - init_noise_std=0.1, - noise_std_type: str = "scalar", - rnn_type="lstm", - rnn_hidden_dim=256, - rnn_num_layers=1, - teacher_recurrent=False, - **kwargs, - ): - if "rnn_hidden_size" in kwargs: - warnings.warn( - "The argument `rnn_hidden_size` is deprecated and will be removed in a future version. " - "Please use `rnn_hidden_dim` instead.", - DeprecationWarning, - ) - if rnn_hidden_dim == 256: # Only override if the new argument is at its default - rnn_hidden_dim = kwargs.pop("rnn_hidden_size") - if kwargs: - print( - "StudentTeacherRecurrent.__init__ got unexpected arguments, which will be ignored: " - + str(kwargs.keys()), - ) - super().__init__() - - self.loaded_teacher = False # indicates if teacher has been loaded - self.teacher_recurrent = teacher_recurrent # indicates if teacher is recurrent too - - # get the observation dimensions - self.obs_groups = obs_groups - num_student_obs = 0 - for obs_group in obs_groups["policy"]: - assert len(obs[obs_group].shape) == 2, "The StudentTeacher module only supports 1D observations." - num_student_obs += obs[obs_group].shape[-1] - num_teacher_obs = 0 - for obs_group in obs_groups["teacher"]: - assert len(obs[obs_group].shape) == 2, "The StudentTeacher module only supports 1D observations." - num_teacher_obs += obs[obs_group].shape[-1] - - # student - self.memory_s = Memory(num_student_obs, type=rnn_type, num_layers=rnn_num_layers, hidden_size=rnn_hidden_dim) - self.student = MLP(rnn_hidden_dim, num_actions, student_hidden_dims, activation) - - # student observation normalization - self.student_obs_normalization = student_obs_normalization - if student_obs_normalization: - self.student_obs_normalizer = EmpiricalNormalization(num_student_obs) - else: - self.student_obs_normalizer = torch.nn.Identity() - - print(f"Student RNN: {self.memory_s}") - print(f"Student MLP: {self.student}") - - # teacher - if self.teacher_recurrent: - self.memory_t = Memory( - num_teacher_obs, type=rnn_type, num_layers=rnn_num_layers, hidden_size=rnn_hidden_dim - ) - num_teacher_obs = rnn_hidden_dim - self.teacher = MLP(num_teacher_obs, num_actions, teacher_hidden_dims, activation) - - # teacher observation normalization - self.teacher_obs_normalization = teacher_obs_normalization - if teacher_obs_normalization: - self.teacher_obs_normalizer = EmpiricalNormalization(num_teacher_obs) - else: - self.teacher_obs_normalizer = torch.nn.Identity() - - if self.teacher_recurrent: - print(f"Teacher RNN: {self.memory_t}") - print(f"Teacher MLP: {self.teacher}") - - # action noise - self.noise_std_type = noise_std_type - if self.noise_std_type == "scalar": - self.std = nn.Parameter(init_noise_std * torch.ones(num_actions)) - elif self.noise_std_type == "log": - self.log_std = nn.Parameter(torch.log(init_noise_std * torch.ones(num_actions))) - else: - raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") - - # action distribution (populated in update_distribution) - self.distribution = None - # disable args validation for speedup - Normal.set_default_validate_args(False) - - def reset(self, dones=None, hidden_states=None): - if hidden_states is None: - hidden_states = (None, None) - self.memory_s.reset(dones, hidden_states[0]) - if self.teacher_recurrent: - self.memory_t.reset(dones, hidden_states[1]) - - def forward(self): - raise NotImplementedError - - @property - def action_mean(self): - return self.distribution.mean - - @property - def action_std(self): - return self.distribution.stddev - - @property - def entropy(self): - return self.distribution.entropy().sum(dim=-1) - - def update_distribution(self, obs): - # compute mean - mean = self.student(obs) - # compute standard deviation - if self.noise_std_type == "scalar": - std = self.std.expand_as(mean) - elif self.noise_std_type == "log": - std = torch.exp(self.log_std).expand_as(mean) - else: - raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") - # create distribution - self.distribution = Normal(mean, std) - - def act(self, obs): - obs = self.get_student_obs(obs) - obs = self.student_obs_normalizer(obs) - out_mem = self.memory_s(obs).squeeze(0) - self.update_distribution(out_mem) - return self.distribution.sample() - - def act_inference(self, obs): - obs = self.get_student_obs(obs) - obs = self.student_obs_normalizer(obs) - out_mem = self.memory_s(obs).squeeze(0) - return self.student(out_mem) - - def evaluate(self, obs): - obs = self.get_teacher_obs(obs) - obs = self.teacher_obs_normalizer(obs) - with torch.no_grad(): - if self.teacher_recurrent: - self.memory_t.eval() - obs = self.memory_t(obs).squeeze(0) - return self.teacher(obs) - - def get_student_obs(self, obs): - obs_list = [] - for obs_group in self.obs_groups["policy"]: - obs_list.append(obs[obs_group]) - return torch.cat(obs_list, dim=-1) - - def get_teacher_obs(self, obs): - obs_list = [] - for obs_group in self.obs_groups["teacher"]: - obs_list.append(obs[obs_group]) - return torch.cat(obs_list, dim=-1) - - def get_hidden_states(self): - if self.teacher_recurrent: - return self.memory_s.hidden_states, self.memory_t.hidden_states - else: - return self.memory_s.hidden_states, None - - def detach_hidden_states(self, dones=None): - self.memory_s.detach_hidden_states(dones) - if self.teacher_recurrent: - self.memory_t.detach_hidden_states(dones) - - def train(self, mode=True): - super().train(mode) - # make sure teacher is in eval mode - self.teacher.eval() - self.teacher_obs_normalizer.eval() - - def update_normalization(self, obs): - if self.student_obs_normalization: - student_obs = self.get_student_obs(obs) - self.student_obs_normalizer.update(student_obs) - - def load_state_dict(self, state_dict, strict=True): - """Load the parameters of the student and teacher networks. - - Args: - state_dict (dict): State dictionary of the model. - strict (bool): Whether to strictly enforce that the keys in state_dict match the keys returned by this - module's state_dict() function. - - Returns: - bool: Whether this training resumes a previous training. This flag is used by the `load()` function of - `OnPolicyRunner` to determine how to load further parameters. - """ - - # check if state_dict contains teacher and student or just teacher parameters - if any("actor" in key for key in state_dict.keys()): # loading parameters from rl training - # rename keys to match teacher and remove critic parameters - teacher_state_dict = {} - teacher_obs_normalizer_state_dict = {} - for key, value in state_dict.items(): - if "actor." in key: - teacher_state_dict[key.replace("actor.", "")] = value - if "actor_obs_normalizer." in key: - teacher_obs_normalizer_state_dict[key.replace("actor_obs_normalizer.", "")] = value - self.teacher.load_state_dict(teacher_state_dict, strict=strict) - self.teacher_obs_normalizer.load_state_dict(teacher_obs_normalizer_state_dict, strict=strict) - # also load recurrent memory if teacher is recurrent - if self.teacher_recurrent: - memory_t_state_dict = {} - for key, value in state_dict.items(): - if "memory_a." in key: - memory_t_state_dict[key.replace("memory_a.", "")] = value - self.memory_t.load_state_dict(memory_t_state_dict, strict=strict) - # set flag for successfully loading the parameters - self.loaded_teacher = True - self.teacher.eval() - self.teacher_obs_normalizer.eval() - return False # training does not resume - elif any("student" in key for key in state_dict.keys()): # loading parameters from distillation training - super().load_state_dict(state_dict, strict=strict) - # set flag for successfully loading the parameters - self.loaded_teacher = True - self.teacher.eval() - self.teacher_obs_normalizer.eval() - return True # training resumes - else: - raise ValueError("state_dict does not contain student or teacher parameters") diff --git a/rsl_rl/build/lib/rsl_rl/modules/symmetry.py b/rsl_rl/build/lib/rsl_rl/modules/symmetry.py deleted file mode 100644 index b017515..0000000 --- a/rsl_rl/build/lib/rsl_rl/modules/symmetry.py +++ /dev/null @@ -1,24 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -from __future__ import annotations - - -def resolve_symmetry_config(alg_cfg, env): - """Resolve the symmetry configuration. - - Args: - alg_cfg: The algorithm configuration dictionary. - env: The environment. - - Returns: - The resolved algorithm configuration dictionary. - """ - - # if using symmetry then pass the environment config object - if "symmetry_cfg" in alg_cfg and alg_cfg["symmetry_cfg"] is not None: - # this is used by the symmetry function for handling different observation terms - alg_cfg["symmetry_cfg"]["_env"] = env - return alg_cfg diff --git a/rsl_rl/build/lib/rsl_rl/networks/__init__.py b/rsl_rl/build/lib/rsl_rl/networks/__init__.py deleted file mode 100644 index c18f487..0000000 --- a/rsl_rl/build/lib/rsl_rl/networks/__init__.py +++ /dev/null @@ -1,10 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -"""Definitions for components of modules.""" - -from .memory import Memory -from .mlp import MLP -from .normalization import EmpiricalDiscountedVariationNormalization, EmpiricalNormalization diff --git a/rsl_rl/build/lib/rsl_rl/networks/memory.py b/rsl_rl/build/lib/rsl_rl/networks/memory.py deleted file mode 100644 index 7577357..0000000 --- a/rsl_rl/build/lib/rsl_rl/networks/memory.py +++ /dev/null @@ -1,70 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -from __future__ import annotations - -import torch.nn as nn - -from rsl_rl.utils import unpad_trajectories - - -class Memory(nn.Module): - """Memory module for recurrent networks. - - This module is used to store the hidden states of the policy. - Currently only supports GRU and LSTM. - """ - - def __init__(self, input_size, type="lstm", num_layers=1, hidden_size=256): - super().__init__() - # RNN - rnn_cls = nn.GRU if type.lower() == "gru" else nn.LSTM - self.rnn = rnn_cls(input_size=input_size, hidden_size=hidden_size, num_layers=num_layers) - self.hidden_states = None - - def forward(self, input, masks=None, hidden_states=None): - batch_mode = masks is not None - if batch_mode: - # batch mode: needs saved hidden states - if hidden_states is None: - raise ValueError("Hidden states not passed to memory module during policy update") - out, _ = self.rnn(input, hidden_states) - out = unpad_trajectories(out, masks) - else: - # inference/distillation mode: uses hidden states of last step - out, self.hidden_states = self.rnn(input.unsqueeze(0), self.hidden_states) - return out - - def reset(self, dones=None, hidden_states=None): - if dones is None: # reset all hidden states - if hidden_states is None: - self.hidden_states = None - else: - self.hidden_states = hidden_states - elif self.hidden_states is not None: # reset hidden states of done environments - if hidden_states is None: - if isinstance(self.hidden_states, tuple): # tuple in case of LSTM - for hidden_state in self.hidden_states: - hidden_state[..., dones == 1, :] = 0.0 - else: - self.hidden_states[..., dones == 1, :] = 0.0 - else: - NotImplementedError( - "Resetting hidden states of done environments with custom hidden states is not implemented" - ) - - def detach_hidden_states(self, dones=None): - if self.hidden_states is not None: - if dones is None: # detach all hidden states - if isinstance(self.hidden_states, tuple): # tuple in case of LSTM - self.hidden_states = tuple(hidden_state.detach() for hidden_state in self.hidden_states) - else: - self.hidden_states = self.hidden_states.detach() - else: # detach hidden states of done environments - if isinstance(self.hidden_states, tuple): # tuple in case of LSTM - for hidden_state in self.hidden_states: - hidden_state[..., dones == 1, :] = hidden_state[..., dones == 1, :].detach() - else: - self.hidden_states[..., dones == 1, :] = self.hidden_states[..., dones == 1, :].detach() diff --git a/rsl_rl/build/lib/rsl_rl/networks/mlp.py b/rsl_rl/build/lib/rsl_rl/networks/mlp.py deleted file mode 100644 index e91574e..0000000 --- a/rsl_rl/build/lib/rsl_rl/networks/mlp.py +++ /dev/null @@ -1,120 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -from __future__ import annotations - -import torch -import torch.nn as nn -from functools import reduce - -from rsl_rl.utils import resolve_nn_activation - - -class MLP(nn.Sequential): - """Multi-layer perceptron. - - The MLP network is a sequence of linear layers and activation functions. The - last layer is a linear layer that outputs the desired dimension unless the - last activation function is specified. - - It provides additional conveniences: - - - If the hidden dimensions have a value of ``-1``, the dimension is inferred - from the input dimension. - - If the output dimension is a tuple, the output is reshaped to the desired - shape. - - """ - - def __init__( - self, - input_dim: int, - output_dim: int | tuple[int] | list[int], - hidden_dims: tuple[int] | list[int], - activation: str = "elu", - last_activation: str | None = None, - ): - """Initialize the MLP. - - Args: - input_dim: Dimension of the input. - output_dim: Dimension of the output. - hidden_dims: Dimensions of the hidden layers. A value of ``-1`` indicates - that the dimension should be inferred from the input dimension. - activation: Activation function. Defaults to "elu". - last_activation: Activation function of the last layer. Defaults to None, - in which case the last layer is linear. - """ - super().__init__() - - # resolve activation functions - activation_mod = resolve_nn_activation(activation) - last_activation_mod = resolve_nn_activation(last_activation) if last_activation is not None else None - # resolve number of hidden dims if they are -1 - hidden_dims_processed = [input_dim if dim == -1 else dim for dim in hidden_dims] - - # create layers sequentially - layers = [] - layers.append(nn.Linear(input_dim, hidden_dims_processed[0])) - layers.append(activation_mod) - - for layer_index in range(len(hidden_dims_processed) - 1): - layers.append(nn.Linear(hidden_dims_processed[layer_index], hidden_dims_processed[layer_index + 1])) - layers.append(activation_mod) - - # add last layer - if isinstance(output_dim, int): - layers.append(nn.Linear(hidden_dims_processed[-1], output_dim)) - else: - # compute the total output dimension - total_out_dim = reduce(lambda x, y: x * y, output_dim) - # add a layer to reshape the output to the desired shape - layers.append(nn.Linear(hidden_dims_processed[-1], total_out_dim)) - layers.append(nn.Unflatten(dim=-1, unflattened_size=output_dim)) - - # add last activation function if specified - if last_activation_mod is not None: - layers.append(last_activation_mod) - - # register the layers - for idx, layer in enumerate(layers): - self.add_module(f"{idx}", layer) - - def init_weights(self, scales: float | tuple[float]): - """Initialize the weights of the MLP. - - Args: - scales: Scale factor for the weights. - """ - - def get_scale(idx) -> float: - """Get the scale factor for the weights of the MLP. - - Args: - idx: Index of the layer. - """ - return scales[idx] if isinstance(scales, (list, tuple)) else scales - - # initialize the weights - for idx, module in enumerate(self): - if isinstance(module, nn.Linear): - nn.init.orthogonal_(module.weight, gain=get_scale(idx)) - nn.init.zeros_(module.bias) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - """Forward pass of the MLP. - - Args: - x: Input tensor. - """ - for layer in self: - x = layer(x) - return x - - def reset(self, dones=None, hidden_states=None): - pass - - def detach_hidden_states(self, dones=None): - pass diff --git a/rsl_rl/build/lib/rsl_rl/networks/normalization.py b/rsl_rl/build/lib/rsl_rl/networks/normalization.py deleted file mode 100644 index 5fd9692..0000000 --- a/rsl_rl/build/lib/rsl_rl/networks/normalization.py +++ /dev/null @@ -1,130 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -# Copyright (c) 2020 Preferred Networks, Inc. - -from __future__ import annotations - -import torch -from torch import nn - - -class EmpiricalNormalization(nn.Module): - """Normalize mean and variance of values based on empirical values.""" - - def __init__(self, shape, eps=1e-2, until=None): - """Initialize EmpiricalNormalization module. - - Args: - shape (int or tuple of int): Shape of input values except batch axis. - eps (float): Small value for stability. - until (int or None): If this arg is specified, the module learns input values until the sum of batch sizes - exceeds it. - - Note: The normalization parameters are computed over the whole batch, not for each environment separately. - """ - super().__init__() - self.eps = eps - self.until = until - self.register_buffer("_mean", torch.zeros(shape).unsqueeze(0)) - self.register_buffer("_var", torch.ones(shape).unsqueeze(0)) - self.register_buffer("_std", torch.ones(shape).unsqueeze(0)) - self.register_buffer("count", torch.tensor(0, dtype=torch.long)) - - @property - def mean(self): - return self._mean.squeeze(0).clone() - - @property - def std(self): - return self._std.squeeze(0).clone() - - def forward(self, x): - """Normalize mean and variance of values based on empirical values.""" - - return (x - self._mean) / (self._std + self.eps) - - @torch.jit.unused - def update(self, x): - """Learn input values without computing the output values of them""" - - if not self.training: - return - if self.until is not None and self.count >= self.until: - return - - count_x = x.shape[0] - self.count += count_x - rate = count_x / self.count - var_x = torch.var(x, dim=0, unbiased=False, keepdim=True) - mean_x = torch.mean(x, dim=0, keepdim=True) - delta_mean = mean_x - self._mean - self._mean += rate * delta_mean - self._var += rate * (var_x - self._var + delta_mean * (mean_x - self._mean)) - self._std = torch.sqrt(self._var) - - @torch.jit.unused - def inverse(self, y): - """De-normalize values based on empirical values.""" - - return y * (self._std + self.eps) + self._mean - - -class EmpiricalDiscountedVariationNormalization(nn.Module): - """Reward normalization from Pathak's large scale study on PPO. - - Reward normalization. Since the reward function is non-stationary, it is useful to normalize - the scale of the rewards so that the value function can learn quickly. We did this by dividing - the rewards by a running estimate of the standard deviation of the sum of discounted rewards. - """ - - def __init__(self, shape, eps=1e-2, gamma=0.99, until=None): - super().__init__() - - self.emp_norm = EmpiricalNormalization(shape, eps, until) - self.disc_avg = _DiscountedAverage(gamma) - - def forward(self, rew): - if self.training: - # update discounted rewards - avg = self.disc_avg.update(rew) - # update moments from discounted rewards - self.emp_norm.update(avg) - - # normalize rewards with the empirical std - if self.emp_norm._std > 0: - return rew / self.emp_norm._std - else: - return rew - - -""" -Helper class. -""" - - -class _DiscountedAverage: - r"""Discounted average of rewards. - - The discounted average is defined as: - - .. math:: - - \bar{R}_t = \gamma \bar{R}_{t-1} + r_t - - Args: - gamma (float): Discount factor. - """ - - def __init__(self, gamma): - self.avg = None - self.gamma = gamma - - def update(self, rew: torch.Tensor) -> torch.Tensor: - if self.avg is None: - self.avg = rew - else: - self.avg = self.avg * self.gamma + rew - return self.avg diff --git a/rsl_rl/build/lib/rsl_rl/runners/__init__.py b/rsl_rl/build/lib/rsl_rl/runners/__init__.py deleted file mode 100644 index 61f1682..0000000 --- a/rsl_rl/build/lib/rsl_rl/runners/__init__.py +++ /dev/null @@ -1,12 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -"""Implementation of runners for environment-agent interaction.""" - -from .on_policy_runner import OnPolicyRunner # isort:skip -from .distillation_runner import DistillationRunner -from .amp_on_policy_runner import AMPOnPolicyRunner - -__all__ = ["OnPolicyRunner", "DistillationRunner", "AMPOnPolicyRunner"] diff --git a/rsl_rl/build/lib/rsl_rl/runners/amp_on_policy_runner.py b/rsl_rl/build/lib/rsl_rl/runners/amp_on_policy_runner.py deleted file mode 100644 index c0b9b9e..0000000 --- a/rsl_rl/build/lib/rsl_rl/runners/amp_on_policy_runner.py +++ /dev/null @@ -1,521 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -from __future__ import annotations - -import os -import statistics -import time -import torch -import warnings -from collections import deque - -import rsl_rl -from rsl_rl.algorithms import AMP_PPO -from rsl_rl.env import VecEnv -from rsl_rl.modules import ActorCritic, ActorCriticRecurrent,DiscriminatorMulti, resolve_rnd_config, resolve_symmetry_config -from rsl_rl.utils import resolve_obs_groups, store_code_state, Normalizer, G1_AMPLoader - - -class AMPOnPolicyRunner: - """On-policy runner for training and evaluation of actor-critic methods.""" - - def __init__(self, env: VecEnv, train_cfg: dict, log_dir: str | None = None, device="cpu"): - self.cfg = train_cfg - self.alg_cfg = train_cfg["algorithm"] - self.policy_cfg = train_cfg["policy"] - self.device = device - self.env = env - - # check if multi-gpu is enabled - self._configure_multi_gpu() - - # store training configuration - self.num_steps_per_env = self.cfg["num_steps_per_env"] - self.save_interval = self.cfg["save_interval"] - - # query observations from environment for algorithm construction - obs = self.env.get_observations() - default_sets = ["critic"] - if "rnd_cfg" in self.alg_cfg and self.alg_cfg["rnd_cfg"] is not None: - default_sets.append("rnd_state") - self.cfg["obs_groups"] = resolve_obs_groups(obs, self.cfg["obs_groups"], default_sets) - - self.amp_data = G1_AMPLoader( - device, - time_between_frames=1/50.0, - preload_transitions=True, - num_preload_transitions=train_cfg["amp_num_preload_transitions"], - motion_files=train_cfg["amp_motion_files"], - num_frames=train_cfg['amp_num_frames'] - ) - - self.amp_observation_dim = self.amp_data.observation_dim if self.cfg["amp_num_obs"] == 0 else self.cfg["amp_num_obs"] - self.amp_num_frames = 0 if self.cfg["amp_num_frames"] == 0 else self.cfg["amp_num_frames"] - self.amp_normalizer = Normalizer(self.amp_observation_dim) - self.discriminator = DiscriminatorMulti( - self.amp_observation_dim, - train_cfg["amp_reward_coef"], - train_cfg["amp_discr_hidden_dims"], - device, - train_cfg["amp_num_frames"], - train_cfg["amp_task_reward_lerp"], - train_cfg['use_lerp'], - ).to(self.device) - - - # create the algorithm - self.alg = self._construct_algorithm(obs) - - # Decide whether to disable logging - # We only log from the process with rank 0 (main process) - self.disable_logs = self.is_distributed and self.gpu_global_rank != 0 - - # Logging - self.log_dir = log_dir - self.writer = None - self.tot_timesteps = 0 - self.tot_time = 0 - self.current_learning_iteration = 0 - self.git_status_repos = [rsl_rl.__file__] - - def learn(self, num_learning_iterations: int, init_at_random_ep_len: bool = False): # noqa: C901 - # initialize writer - self._prepare_logging_writer() - - # randomize initial episode lengths (for exploration) - if init_at_random_ep_len: - self.env.episode_length_buf = torch.randint_like( - self.env.episode_length_buf, high=int(self.env.max_episode_length) - ) - - # start learning - obs = self.env.get_observations().to(self.device) - self.train_mode() # switch to train mode (for dropout for example) - amp_obs = self.env.get_amp_observations() - amp_obs = amp_obs.to(self.device) - if self.amp_num_frames != 0: - self.amp_obs_frames = torch.zeros(size=(self.env.num_envs, self.amp_num_frames, self.amp_observation_dim), device=self.device) - self.amp_obs_frames = torch.concat((self.amp_obs_frames[:, 1:], amp_obs.unsqueeze(1)), dim=1) - - # Book keeping - ep_infos = [] - rewbuffer = deque(maxlen=100) - lenbuffer = deque(maxlen=100) - step_discrewbuffer = deque(maxlen=100) - - cur_reward_sum = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) - cur_episode_length = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) - cur_single_step_disc_rew = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) - # create buffers for logging extrinsic and intrinsic rewards - if self.alg.rnd: - erewbuffer = deque(maxlen=100) - irewbuffer = deque(maxlen=100) - cur_ereward_sum = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) - cur_ireward_sum = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) - - # Ensure all parameters are in-synced - if self.is_distributed: - print(f"Synchronizing parameters for rank {self.gpu_global_rank}...") - self.alg.broadcast_parameters() - - # Start training - start_iter = self.current_learning_iteration - tot_iter = start_iter + num_learning_iterations - for it in range(start_iter, tot_iter): - start = time.time() - # Rollout - with torch.inference_mode(): - for _ in range(self.num_steps_per_env): - # Sample actions - actions = self.alg.act(obs,amp_obs) - # Step the environment - obs, rewards, dones, extras = self.env.step(actions.to(self.env.device)) - # Move to device - obs, rewards, dones = (obs.to(self.device), rewards.to(self.device), dones.to(self.device)) - - next_amp_obs = self.env.get_amp_observations() - next_amp_obs = next_amp_obs.to(self.device) - next_amp_obs_with_term = torch.clone(next_amp_obs) - - reset_env_ids = self.env.reset_env_ids - terminal_amp_states = self.env.get_amp_observations()[reset_env_ids] - next_amp_obs_with_term[reset_env_ids] = terminal_amp_states - self.amp_obs_frames = torch.concat((self.amp_obs_frames[:, 1:], next_amp_obs_with_term.unsqueeze(1)), dim=1) - - amp_reward = torch.zeros(self.env.num_envs, device=obs.device) - - mask = self.env.contact_phase[:, 0] == 1.0 - if mask.any(): - rewards[mask], logit, disc_reward = self.alg.discriminator.predict_amp_reward( - self.amp_obs_frames[mask], rewards[mask], normalizer=self.alg.amp_normalizer - ) - amp_reward[mask] += disc_reward - - # process the step - self.alg.process_env_step(obs, rewards, dones, extras, next_amp_obs_with_term, self.amp_obs_frames) - self.amp_obs_frames[reset_env_ids] = 0 - - amp_obs = torch.clone(next_amp_obs) - # Extract intrinsic rewards (only for logging) - intrinsic_rewards = self.alg.intrinsic_rewards if self.alg.rnd else None - # book keeping - if self.log_dir is not None: - if "episode" in extras: - ep_infos.append(extras["episode"]) - elif "log" in extras: - ep_infos.append(extras["log"]) - # Update rewards - if self.alg.rnd: - cur_ereward_sum += rewards - cur_ireward_sum += intrinsic_rewards # type: ignore - cur_reward_sum += rewards + intrinsic_rewards - else: - cur_reward_sum += rewards - # Update episode length - cur_episode_length += 1 - cur_single_step_disc_rew += amp_reward - # Clear data for completed episodes - # -- common - new_ids = (dones > 0).nonzero(as_tuple=False) - rewbuffer.extend(cur_reward_sum[new_ids][:, 0].cpu().numpy().tolist()) - lenbuffer.extend(cur_episode_length[new_ids][:, 0].cpu().numpy().tolist()) - cur_reward_sum[new_ids] = 0 - cur_episode_length[new_ids] = 0 - to_extend_disc = (cur_single_step_disc_rew[new_ids] / self.env.max_episode_length_s)[:, 0].cpu().numpy() - step_discrewbuffer.extend(to_extend_disc.tolist()) - cur_single_step_disc_rew[new_ids] = 0 - # -- intrinsic and extrinsic rewards - if self.alg.rnd: - erewbuffer.extend(cur_ereward_sum[new_ids][:, 0].cpu().numpy().tolist()) - irewbuffer.extend(cur_ireward_sum[new_ids][:, 0].cpu().numpy().tolist()) - cur_ereward_sum[new_ids] = 0 - cur_ireward_sum[new_ids] = 0 - - stop = time.time() - collection_time = stop - start - start = stop - - # compute returns - self.alg.compute_returns(obs) - - # update policy - loss_dict = self.alg.update() - - stop = time.time() - learn_time = stop - start - self.current_learning_iteration = it - # log info - if self.log_dir is not None and not self.disable_logs: - # Log information - self.log(locals()) - # Save model - if it % self.save_interval == 0: - self.save(os.path.join(self.log_dir, f"model_{it}.pt")) - - # Clear episode infos - ep_infos.clear() - # Save code state - if it == start_iter and not self.disable_logs: - # obtain all the diff files - git_file_paths = store_code_state(self.log_dir, self.git_status_repos) - # if possible store them to wandb - if self.logger_type in ["wandb", "neptune"] and git_file_paths: - for path in git_file_paths: - self.writer.save_file(path) - - # Save the final model after training - if self.log_dir is not None and not self.disable_logs: - self.save(os.path.join(self.log_dir, f"model_{self.current_learning_iteration}.pt")) - - def log(self, locs: dict, width: int = 80, pad: int = 35): - # Compute the collection size - collection_size = self.num_steps_per_env * self.env.num_envs * self.gpu_world_size - # Update total time-steps and time - self.tot_timesteps += collection_size - self.tot_time += locs["collection_time"] + locs["learn_time"] - iteration_time = locs["collection_time"] + locs["learn_time"] - - # -- Episode info - ep_string = "" - if locs["ep_infos"]: - for key in locs["ep_infos"][0]: - infotensor = torch.tensor([], device=self.device) - for ep_info in locs["ep_infos"]: - # handle scalar and zero dimensional tensor infos - if key not in ep_info: - continue - if not isinstance(ep_info[key], torch.Tensor): - ep_info[key] = torch.Tensor([ep_info[key]]) - if len(ep_info[key].shape) == 0: - ep_info[key] = ep_info[key].unsqueeze(0) - infotensor = torch.cat((infotensor, ep_info[key].to(self.device))) - value = torch.mean(infotensor) - # log to logger and terminal - if "/" in key: - self.writer.add_scalar(key, value, locs["it"]) - ep_string += f"""{f'{key}:':>{pad}} {value:.4f}\n""" - else: - self.writer.add_scalar("Episode/" + key, value, locs["it"]) - ep_string += f"""{f'Mean episode {key}:':>{pad}} {value:.4f}\n""" - - mean_std = self.alg.policy.action_std.mean() - fps = int(collection_size / (locs["collection_time"] + locs["learn_time"])) - - # -- Losses - for key, value in locs["loss_dict"].items(): - self.writer.add_scalar(f"Loss/{key}", value, locs["it"]) - self.writer.add_scalar("Loss/learning_rate", self.alg.learning_rate, locs["it"]) - - # -- Policy - self.writer.add_scalar("Policy/mean_noise_std", mean_std.item(), locs["it"]) - - # -- Performance - self.writer.add_scalar("Perf/total_fps", fps, locs["it"]) - self.writer.add_scalar("Perf/collection time", locs["collection_time"], locs["it"]) - self.writer.add_scalar("Perf/learning_time", locs["learn_time"], locs["it"]) - - # -- Training - if len(locs["rewbuffer"]) > 0: - # separate logging for intrinsic and extrinsic rewards - if hasattr(self.alg, "rnd") and self.alg.rnd: - self.writer.add_scalar("Rnd/mean_extrinsic_reward", statistics.mean(locs["erewbuffer"]), locs["it"]) - self.writer.add_scalar("Rnd/mean_intrinsic_reward", statistics.mean(locs["irewbuffer"]), locs["it"]) - self.writer.add_scalar("Rnd/weight", self.alg.rnd.weight, locs["it"]) - # everything else - self.writer.add_scalar("Train/mean_reward", statistics.mean(locs["rewbuffer"]), locs["it"]) - self.writer.add_scalar("Train/mean_episode_length", statistics.mean(locs["lenbuffer"]), locs["it"]) - self.writer.add_scalar('Train/mean_step_disc_reward', statistics.mean(locs['step_discrewbuffer']), locs['it']) - if self.logger_type != "wandb": # wandb does not support non-integer x-axis logging - self.writer.add_scalar("Train/mean_reward/time", statistics.mean(locs["rewbuffer"]), self.tot_time) - self.writer.add_scalar( - "Train/mean_episode_length/time", statistics.mean(locs["lenbuffer"]), self.tot_time - ) - - str = f" \033[1m Learning iteration {locs['it']}/{locs['tot_iter']} \033[0m " - - if len(locs["rewbuffer"]) > 0: - log_string = ( - f"""{'#' * width}\n""" - f"""{str.center(width, ' ')}\n\n""" - f"""{'Computation:':>{pad}} {fps:.0f} steps/s (collection: {locs[ - 'collection_time']:.3f}s, learning {locs['learn_time']:.3f}s)\n""" - f"""{'Mean action noise std:':>{pad}} {mean_std.item():.2f}\n""" - f"""{'Step disc reward:':>{pad}} {statistics.mean(locs['step_discrewbuffer']):.2f}\n""" - ) - # -- Losses - for key, value in locs["loss_dict"].items(): - log_string += f"""{f'Mean {key} loss:':>{pad}} {value:.4f}\n""" - # -- Rewards - if hasattr(self.alg, "rnd") and self.alg.rnd: - log_string += ( - f"""{'Mean extrinsic reward:':>{pad}} {statistics.mean(locs['erewbuffer']):.2f}\n""" - f"""{'Mean intrinsic reward:':>{pad}} {statistics.mean(locs['irewbuffer']):.2f}\n""" - ) - log_string += f"""{'Mean reward:':>{pad}} {statistics.mean(locs['rewbuffer']):.2f}\n""" - # -- episode info - log_string += f"""{'Mean episode length:':>{pad}} {statistics.mean(locs['lenbuffer']):.2f}\n""" - else: - log_string = ( - f"""{'#' * width}\n""" - f"""{str.center(width, ' ')}\n\n""" - f"""{'Computation:':>{pad}} {fps:.0f} steps/s (collection: {locs[ - 'collection_time']:.3f}s, learning {locs['learn_time']:.3f}s)\n""" - f"""{'Mean action noise std:':>{pad}} {mean_std.item():.2f}\n""" - ) - for key, value in locs["loss_dict"].items(): - log_string += f"""{f'{key}:':>{pad}} {value:.4f}\n""" - - log_string += ep_string - log_string += ( - f"""{'-' * width}\n""" - f"""{'Total timesteps:':>{pad}} {self.tot_timesteps}\n""" - f"""{'Iteration time:':>{pad}} {iteration_time:.2f}s\n""" - f"""{'Time elapsed:':>{pad}} {time.strftime("%H:%M:%S", time.gmtime(self.tot_time))}\n""" - f"""{'ETA:':>{pad}} {time.strftime( - "%H:%M:%S", - time.gmtime( - self.tot_time / (locs['it'] - locs['start_iter'] + 1) - * (locs['start_iter'] + locs['num_learning_iterations'] - locs['it']) - ) - )}\n""" - ) - print(log_string) - - def save(self, path: str, infos=None): - # -- Save model - saved_dict = { - "model_state_dict": self.alg.policy.state_dict(), - "optimizer_state_dict": self.alg.optimizer.state_dict(), - "iter": self.current_learning_iteration, - "infos": infos, - } - # -- Save RND model if used - if hasattr(self.alg, "rnd") and self.alg.rnd: - saved_dict["rnd_state_dict"] = self.alg.rnd.state_dict() - saved_dict["rnd_optimizer_state_dict"] = self.alg.rnd_optimizer.state_dict() - torch.save(saved_dict, path) - - # upload model to external logging service - if self.logger_type in ["neptune", "wandb"] and not self.disable_logs: - self.writer.save_model(path, self.current_learning_iteration) - - def load(self, path: str, load_optimizer: bool = True, map_location: str | None = None): - loaded_dict = torch.load(path, weights_only=False, map_location=map_location) - # -- Load model - resumed_training = self.alg.policy.load_state_dict(loaded_dict["model_state_dict"]) - # -- Load RND model if used - if hasattr(self.alg, "rnd") and self.alg.rnd: - self.alg.rnd.load_state_dict(loaded_dict["rnd_state_dict"]) - # -- load optimizer if used - if load_optimizer and resumed_training: - # -- algorithm optimizer - self.alg.optimizer.load_state_dict(loaded_dict["optimizer_state_dict"]) - # -- RND optimizer if used - if hasattr(self.alg, "rnd") and self.alg.rnd: - self.alg.rnd_optimizer.load_state_dict(loaded_dict["rnd_optimizer_state_dict"]) - # -- load current learning iteration - if resumed_training: - self.current_learning_iteration = loaded_dict["iter"] - return loaded_dict["infos"] - - def get_inference_policy(self, device=None): - self.eval_mode() # switch to evaluation mode (dropout for example) - if device is not None: - self.alg.policy.to(device) - return self.alg.policy.act_inference - - def train_mode(self): - # -- PPO - self.alg.policy.train() - self.alg.discriminator.train() - # -- RND - if hasattr(self.alg, "rnd") and self.alg.rnd: - self.alg.rnd.train() - - def eval_mode(self): - # -- PPO - self.alg.policy.eval() - self.alg.discriminator.eval() - # -- RND - if hasattr(self.alg, "rnd") and self.alg.rnd: - self.alg.rnd.eval() - - def add_git_repo_to_log(self, repo_file_path): - self.git_status_repos.append(repo_file_path) - - """ - Helper functions. - """ - - def _configure_multi_gpu(self): - """Configure multi-gpu training.""" - # check if distributed training is enabled - self.gpu_world_size = int(os.getenv("WORLD_SIZE", "1")) - self.is_distributed = self.gpu_world_size > 1 - - # if not distributed training, set local and global rank to 0 and return - if not self.is_distributed: - self.gpu_local_rank = 0 - self.gpu_global_rank = 0 - self.multi_gpu_cfg = None - return - - # get rank and world size - self.gpu_local_rank = int(os.getenv("LOCAL_RANK", "0")) - self.gpu_global_rank = int(os.getenv("RANK", "0")) - - # make a configuration dictionary - self.multi_gpu_cfg = { - "global_rank": self.gpu_global_rank, # rank of the main process - "local_rank": self.gpu_local_rank, # rank of the current process - "world_size": self.gpu_world_size, # total number of processes - } - - # check if user has device specified for local rank - if self.device != f"cuda:{self.gpu_local_rank}": - raise ValueError( - f"Device '{self.device}' does not match expected device for local rank '{self.gpu_local_rank}'." - ) - # validate multi-gpu configuration - if self.gpu_local_rank >= self.gpu_world_size: - raise ValueError( - f"Local rank '{self.gpu_local_rank}' is greater than or equal to world size '{self.gpu_world_size}'." - ) - if self.gpu_global_rank >= self.gpu_world_size: - raise ValueError( - f"Global rank '{self.gpu_global_rank}' is greater than or equal to world size '{self.gpu_world_size}'." - ) - - # initialize torch distributed - torch.distributed.init_process_group(backend="nccl", rank=self.gpu_global_rank, world_size=self.gpu_world_size) - # set device to the local rank - torch.cuda.set_device(self.gpu_local_rank) - - def _construct_algorithm(self, obs) -> AMP_PPO: - """Construct the actor-critic algorithm.""" - # resolve RND config - self.alg_cfg = resolve_rnd_config(self.alg_cfg, obs, self.cfg["obs_groups"], self.env) - - # resolve symmetry config - self.alg_cfg = resolve_symmetry_config(self.alg_cfg, self.env) - - # resolve deprecated normalization config - if self.cfg.get("empirical_normalization") is not None: - warnings.warn( - "The `empirical_normalization` parameter is deprecated. Please set `actor_obs_normalization` and " - "`critic_obs_normalization` as part of the `policy` configuration instead.", - DeprecationWarning, - ) - if self.policy_cfg.get("actor_obs_normalization") is None: - self.policy_cfg["actor_obs_normalization"] = self.cfg["empirical_normalization"] - if self.policy_cfg.get("critic_obs_normalization") is None: - self.policy_cfg["critic_obs_normalization"] = self.cfg["empirical_normalization"] - - # initialize the actor-critic - actor_critic_class = eval(self.policy_cfg.pop("class_name")) - actor_critic: ActorCritic | ActorCriticRecurrent = actor_critic_class( - obs, self.cfg["obs_groups"], self.env.num_actions, **self.policy_cfg - ).to(self.device) - - # initialize the algorithm - alg_class = eval(self.alg_cfg.pop("class_name")) - - alg: AMP_PPO = alg_class(actor_critic, self.discriminator, self.amp_data, self.amp_normalizer, self.amp_num_frames, device=self.device, **self.alg_cfg, multi_gpu_cfg=self.multi_gpu_cfg) - - # initialize the storage - alg.init_storage( - "rl", - self.env.num_envs, - self.num_steps_per_env, - obs, - [self.env.num_actions], - ) - - return alg - - def _prepare_logging_writer(self): - """Prepares the logging writers.""" - if self.log_dir is not None and self.writer is None and not self.disable_logs: - # Launch either Tensorboard or Neptune & Tensorboard summary writer(s), default: Tensorboard. - self.logger_type = self.cfg.get("logger", "tensorboard") - self.logger_type = self.logger_type.lower() - - if self.logger_type == "neptune": - from rsl_rl.utils.neptune_utils import NeptuneSummaryWriter - - self.writer = NeptuneSummaryWriter(log_dir=self.log_dir, flush_secs=10, cfg=self.cfg) - self.writer.log_config(self.env.cfg, self.cfg, self.alg_cfg, self.policy_cfg) - elif self.logger_type == "wandb": - from rsl_rl.utils.wandb_utils import WandbSummaryWriter - - self.writer = WandbSummaryWriter(log_dir=self.log_dir, flush_secs=10, cfg=self.cfg) - self.writer.log_config(self.env.cfg, self.cfg, self.alg_cfg, self.policy_cfg) - elif self.logger_type == "tensorboard": - from torch.utils.tensorboard import SummaryWriter - - self.writer = SummaryWriter(log_dir=self.log_dir, flush_secs=10) - else: - raise ValueError("Logger type not found. Please choose 'neptune', 'wandb' or 'tensorboard'.") diff --git a/rsl_rl/build/lib/rsl_rl/runners/distillation_runner.py b/rsl_rl/build/lib/rsl_rl/runners/distillation_runner.py deleted file mode 100644 index 9cc6a8b..0000000 --- a/rsl_rl/build/lib/rsl_rl/runners/distillation_runner.py +++ /dev/null @@ -1,179 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -from __future__ import annotations - -import os -import time -import torch -from collections import deque - -import rsl_rl -from rsl_rl.algorithms import Distillation -from rsl_rl.env import VecEnv -from rsl_rl.modules import StudentTeacher, StudentTeacherRecurrent -from rsl_rl.runners import OnPolicyRunner -from rsl_rl.utils import resolve_obs_groups, store_code_state - - -class DistillationRunner(OnPolicyRunner): - """On-policy runner for training and evaluation of teacher-student training.""" - - def __init__(self, env: VecEnv, train_cfg: dict, log_dir: str | None = None, device="cpu"): - self.cfg = train_cfg - self.alg_cfg = train_cfg["algorithm"] - self.policy_cfg = train_cfg["policy"] - self.device = device - self.env = env - - # check if multi-gpu is enabled - self._configure_multi_gpu() - - # store training configuration - self.num_steps_per_env = self.cfg["num_steps_per_env"] - self.save_interval = self.cfg["save_interval"] - - # query observations from environment for algorithm construction - obs = self.env.get_observations() - self.cfg["obs_groups"] = resolve_obs_groups(obs, self.cfg["obs_groups"], default_sets=["teacher"]) - - # create the algorithm - self.alg = self._construct_algorithm(obs) - - # Decide whether to disable logging - # We only log from the process with rank 0 (main process) - self.disable_logs = self.is_distributed and self.gpu_global_rank != 0 - - # Logging - self.log_dir = log_dir - self.writer = None - self.tot_timesteps = 0 - self.tot_time = 0 - self.current_learning_iteration = 0 - self.git_status_repos = [rsl_rl.__file__] - - def learn(self, num_learning_iterations: int, init_at_random_ep_len: bool = False): # noqa: C901 - # initialize writer - self._prepare_logging_writer() - # check if teacher is loaded - if not self.alg.policy.loaded_teacher: - raise ValueError("Teacher model parameters not loaded. Please load a teacher model to distill.") - - # randomize initial episode lengths (for exploration) - if init_at_random_ep_len: - self.env.episode_length_buf = torch.randint_like( - self.env.episode_length_buf, high=int(self.env.max_episode_length) - ) - - # start learning - obs = self.env.get_observations().to(self.device) - self.train_mode() # switch to train mode (for dropout for example) - - # Book keeping - ep_infos = [] - rewbuffer = deque(maxlen=100) - lenbuffer = deque(maxlen=100) - cur_reward_sum = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) - cur_episode_length = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) - - # Ensure all parameters are in-synced - if self.is_distributed: - print(f"Synchronizing parameters for rank {self.gpu_global_rank}...") - self.alg.broadcast_parameters() - - # Start training - start_iter = self.current_learning_iteration - tot_iter = start_iter + num_learning_iterations - for it in range(start_iter, tot_iter): - start = time.time() - # Rollout - with torch.inference_mode(): - for _ in range(self.num_steps_per_env): - # Sample actions - actions = self.alg.act(obs) - # Step the environment - obs, rewards, dones, extras = self.env.step(actions.to(self.env.device)) - # Move to device - obs, rewards, dones = (obs.to(self.device), rewards.to(self.device), dones.to(self.device)) - # process the step - self.alg.process_env_step(obs, rewards, dones, extras) - # book keeping - if self.log_dir is not None: - if "episode" in extras: - ep_infos.append(extras["episode"]) - elif "log" in extras: - ep_infos.append(extras["log"]) - # Update rewards - cur_reward_sum += rewards - # Update episode length - cur_episode_length += 1 - # Clear data for completed episodes - new_ids = (dones > 0).nonzero(as_tuple=False) - rewbuffer.extend(cur_reward_sum[new_ids][:, 0].cpu().numpy().tolist()) - lenbuffer.extend(cur_episode_length[new_ids][:, 0].cpu().numpy().tolist()) - cur_reward_sum[new_ids] = 0 - cur_episode_length[new_ids] = 0 - - stop = time.time() - collection_time = stop - start - start = stop - - # update policy - loss_dict = self.alg.update() - - stop = time.time() - learn_time = stop - start - self.current_learning_iteration = it - # log info - if self.log_dir is not None and not self.disable_logs: - # Log information - self.log(locals()) - # Save model - if it % self.save_interval == 0: - self.save(os.path.join(self.log_dir, f"model_{it}.pt")) - - # Clear episode infos - ep_infos.clear() - # Save code state - if it == start_iter and not self.disable_logs: - # obtain all the diff files - git_file_paths = store_code_state(self.log_dir, self.git_status_repos) - # if possible store them to wandb - if self.logger_type in ["wandb", "neptune"] and git_file_paths: - for path in git_file_paths: - self.writer.save_file(path) - - # Save the final model after training - if self.log_dir is not None and not self.disable_logs: - self.save(os.path.join(self.log_dir, f"model_{self.current_learning_iteration}.pt")) - - """ - Helper methods. - """ - - def _construct_algorithm(self, obs) -> Distillation: - """Construct the distillation algorithm.""" - # initialize the actor-critic - student_teacher_class = eval(self.policy_cfg.pop("class_name")) - student_teacher: StudentTeacher | StudentTeacherRecurrent = student_teacher_class( - obs, self.cfg["obs_groups"], self.env.num_actions, **self.policy_cfg - ).to(self.device) - - # initialize the algorithm - alg_class = eval(self.alg_cfg.pop("class_name")) - alg: Distillation = alg_class( - student_teacher, device=self.device, **self.alg_cfg, multi_gpu_cfg=self.multi_gpu_cfg - ) - - # initialize the storage - alg.init_storage( - "distillation", - self.env.num_envs, - self.num_steps_per_env, - obs, - [self.env.num_actions], - ) - - return alg diff --git a/rsl_rl/build/lib/rsl_rl/runners/on_policy_runner.py b/rsl_rl/build/lib/rsl_rl/runners/on_policy_runner.py deleted file mode 100644 index 36f11f3..0000000 --- a/rsl_rl/build/lib/rsl_rl/runners/on_policy_runner.py +++ /dev/null @@ -1,460 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -from __future__ import annotations - -import os -import statistics -import time -import torch -import warnings -from collections import deque - -import rsl_rl -from rsl_rl.algorithms import PPO -from rsl_rl.env import VecEnv -from rsl_rl.modules import ActorCritic, ActorCriticRecurrent, resolve_rnd_config, resolve_symmetry_config -from rsl_rl.utils import resolve_obs_groups, store_code_state - - -class OnPolicyRunner: - """On-policy runner for training and evaluation of actor-critic methods.""" - - def __init__(self, env: VecEnv, train_cfg: dict, log_dir: str | None = None, device="cpu"): - self.cfg = train_cfg - self.alg_cfg = train_cfg["algorithm"] - self.policy_cfg = train_cfg["policy"] - self.device = device - self.env = env - - # check if multi-gpu is enabled - self._configure_multi_gpu() - - # store training configuration - self.num_steps_per_env = self.cfg["num_steps_per_env"] - self.save_interval = self.cfg["save_interval"] - - # query observations from environment for algorithm construction - obs = self.env.get_observations() - default_sets = ["critic"] - if "rnd_cfg" in self.alg_cfg and self.alg_cfg["rnd_cfg"] is not None: - default_sets.append("rnd_state") - self.cfg["obs_groups"] = resolve_obs_groups(obs, self.cfg["obs_groups"], default_sets) - - # create the algorithm - self.alg = self._construct_algorithm(obs) - - # Decide whether to disable logging - # We only log from the process with rank 0 (main process) - self.disable_logs = self.is_distributed and self.gpu_global_rank != 0 - - # Logging - self.log_dir = log_dir - self.writer = None - self.tot_timesteps = 0 - self.tot_time = 0 - self.current_learning_iteration = 0 - self.git_status_repos = [rsl_rl.__file__] - - def learn(self, num_learning_iterations: int, init_at_random_ep_len: bool = False): # noqa: C901 - # initialize writer - self._prepare_logging_writer() - - # randomize initial episode lengths (for exploration) - if init_at_random_ep_len: - self.env.episode_length_buf = torch.randint_like( - self.env.episode_length_buf, high=int(self.env.max_episode_length) - ) - - # start learning - obs = self.env.get_observations().to(self.device) - self.train_mode() # switch to train mode (for dropout for example) - - # Book keeping - ep_infos = [] - rewbuffer = deque(maxlen=100) - lenbuffer = deque(maxlen=100) - cur_reward_sum = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) - cur_episode_length = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) - - # create buffers for logging extrinsic and intrinsic rewards - if self.alg.rnd: - erewbuffer = deque(maxlen=100) - irewbuffer = deque(maxlen=100) - cur_ereward_sum = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) - cur_ireward_sum = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) - - # Ensure all parameters are in-synced - if self.is_distributed: - print(f"Synchronizing parameters for rank {self.gpu_global_rank}...") - self.alg.broadcast_parameters() - - # Start training - start_iter = self.current_learning_iteration - tot_iter = start_iter + num_learning_iterations - for it in range(start_iter, tot_iter): - start = time.time() - # Rollout - with torch.inference_mode(): - for _ in range(self.num_steps_per_env): - # Sample actions - actions = self.alg.act(obs) - # Step the environment - obs, rewards, dones, extras = self.env.step(actions.to(self.env.device)) - # Move to device - obs, rewards, dones = (obs.to(self.device), rewards.to(self.device), dones.to(self.device)) - # process the step - self.alg.process_env_step(obs, rewards, dones, extras) - # Extract intrinsic rewards (only for logging) - intrinsic_rewards = self.alg.intrinsic_rewards if self.alg.rnd else None - # book keeping - if self.log_dir is not None: - if "episode" in extras: - ep_infos.append(extras["episode"]) - elif "log" in extras: - ep_infos.append(extras["log"]) - # Update rewards - if self.alg.rnd: - cur_ereward_sum += rewards - cur_ireward_sum += intrinsic_rewards # type: ignore - cur_reward_sum += rewards + intrinsic_rewards - else: - cur_reward_sum += rewards - # Update episode length - cur_episode_length += 1 - # Clear data for completed episodes - # -- common - new_ids = (dones > 0).nonzero(as_tuple=False) - rewbuffer.extend(cur_reward_sum[new_ids][:, 0].cpu().numpy().tolist()) - lenbuffer.extend(cur_episode_length[new_ids][:, 0].cpu().numpy().tolist()) - cur_reward_sum[new_ids] = 0 - cur_episode_length[new_ids] = 0 - # -- intrinsic and extrinsic rewards - if self.alg.rnd: - erewbuffer.extend(cur_ereward_sum[new_ids][:, 0].cpu().numpy().tolist()) - irewbuffer.extend(cur_ireward_sum[new_ids][:, 0].cpu().numpy().tolist()) - cur_ereward_sum[new_ids] = 0 - cur_ireward_sum[new_ids] = 0 - - stop = time.time() - collection_time = stop - start - start = stop - - # compute returns - self.alg.compute_returns(obs) - - # update policy - loss_dict = self.alg.update() - - stop = time.time() - learn_time = stop - start - self.current_learning_iteration = it - # log info - if self.log_dir is not None and not self.disable_logs: - # Log information - self.log(locals()) - # Save model - if it % self.save_interval == 0: - self.save(os.path.join(self.log_dir, f"model_{it}.pt")) - - # Clear episode infos - ep_infos.clear() - # Save code state - if it == start_iter and not self.disable_logs: - # obtain all the diff files - git_file_paths = store_code_state(self.log_dir, self.git_status_repos) - # if possible store them to wandb - if self.logger_type in ["wandb", "neptune"] and git_file_paths: - for path in git_file_paths: - self.writer.save_file(path) - - # Save the final model after training - if self.log_dir is not None and not self.disable_logs: - self.save(os.path.join(self.log_dir, f"model_{self.current_learning_iteration}.pt")) - - def log(self, locs: dict, width: int = 80, pad: int = 35): - # Compute the collection size - collection_size = self.num_steps_per_env * self.env.num_envs * self.gpu_world_size - # Update total time-steps and time - self.tot_timesteps += collection_size - self.tot_time += locs["collection_time"] + locs["learn_time"] - iteration_time = locs["collection_time"] + locs["learn_time"] - - # -- Episode info - ep_string = "" - if locs["ep_infos"]: - for key in locs["ep_infos"][0]: - infotensor = torch.tensor([], device=self.device) - for ep_info in locs["ep_infos"]: - # handle scalar and zero dimensional tensor infos - if key not in ep_info: - continue - if not isinstance(ep_info[key], torch.Tensor): - ep_info[key] = torch.Tensor([ep_info[key]]) - if len(ep_info[key].shape) == 0: - ep_info[key] = ep_info[key].unsqueeze(0) - infotensor = torch.cat((infotensor, ep_info[key].to(self.device))) - value = torch.mean(infotensor) - # log to logger and terminal - if "/" in key: - self.writer.add_scalar(key, value, locs["it"]) - ep_string += f"""{f'{key}:':>{pad}} {value:.4f}\n""" - else: - self.writer.add_scalar("Episode/" + key, value, locs["it"]) - ep_string += f"""{f'Mean episode {key}:':>{pad}} {value:.4f}\n""" - - mean_std = self.alg.policy.action_std.mean() - fps = int(collection_size / (locs["collection_time"] + locs["learn_time"])) - - # -- Losses - for key, value in locs["loss_dict"].items(): - self.writer.add_scalar(f"Loss/{key}", value, locs["it"]) - self.writer.add_scalar("Loss/learning_rate", self.alg.learning_rate, locs["it"]) - - # -- Policy - self.writer.add_scalar("Policy/mean_noise_std", mean_std.item(), locs["it"]) - - # -- Performance - self.writer.add_scalar("Perf/total_fps", fps, locs["it"]) - self.writer.add_scalar("Perf/collection time", locs["collection_time"], locs["it"]) - self.writer.add_scalar("Perf/learning_time", locs["learn_time"], locs["it"]) - - # -- Training - if len(locs["rewbuffer"]) > 0: - # separate logging for intrinsic and extrinsic rewards - if hasattr(self.alg, "rnd") and self.alg.rnd: - self.writer.add_scalar("Rnd/mean_extrinsic_reward", statistics.mean(locs["erewbuffer"]), locs["it"]) - self.writer.add_scalar("Rnd/mean_intrinsic_reward", statistics.mean(locs["irewbuffer"]), locs["it"]) - self.writer.add_scalar("Rnd/weight", self.alg.rnd.weight, locs["it"]) - # everything else - self.writer.add_scalar("Train/mean_reward", statistics.mean(locs["rewbuffer"]), locs["it"]) - self.writer.add_scalar("Train/mean_episode_length", statistics.mean(locs["lenbuffer"]), locs["it"]) - if self.logger_type != "wandb": # wandb does not support non-integer x-axis logging - self.writer.add_scalar("Train/mean_reward/time", statistics.mean(locs["rewbuffer"]), self.tot_time) - self.writer.add_scalar( - "Train/mean_episode_length/time", statistics.mean(locs["lenbuffer"]), self.tot_time - ) - - str = f" \033[1m Learning iteration {locs['it']}/{locs['tot_iter']} \033[0m " - - if len(locs["rewbuffer"]) > 0: - log_string = ( - f"""{'#' * width}\n""" - f"""{str.center(width, ' ')}\n\n""" - f"""{'Computation:':>{pad}} {fps:.0f} steps/s (collection: {locs[ - 'collection_time']:.3f}s, learning {locs['learn_time']:.3f}s)\n""" - f"""{'Mean action noise std:':>{pad}} {mean_std.item():.2f}\n""" - ) - # -- Losses - for key, value in locs["loss_dict"].items(): - log_string += f"""{f'Mean {key} loss:':>{pad}} {value:.4f}\n""" - # -- Rewards - if hasattr(self.alg, "rnd") and self.alg.rnd: - log_string += ( - f"""{'Mean extrinsic reward:':>{pad}} {statistics.mean(locs['erewbuffer']):.2f}\n""" - f"""{'Mean intrinsic reward:':>{pad}} {statistics.mean(locs['irewbuffer']):.2f}\n""" - ) - log_string += f"""{'Mean reward:':>{pad}} {statistics.mean(locs['rewbuffer']):.2f}\n""" - # -- episode info - log_string += f"""{'Mean episode length:':>{pad}} {statistics.mean(locs['lenbuffer']):.2f}\n""" - else: - log_string = ( - f"""{'#' * width}\n""" - f"""{str.center(width, ' ')}\n\n""" - f"""{'Computation:':>{pad}} {fps:.0f} steps/s (collection: {locs[ - 'collection_time']:.3f}s, learning {locs['learn_time']:.3f}s)\n""" - f"""{'Mean action noise std:':>{pad}} {mean_std.item():.2f}\n""" - ) - for key, value in locs["loss_dict"].items(): - log_string += f"""{f'{key}:':>{pad}} {value:.4f}\n""" - - log_string += ep_string - log_string += ( - f"""{'-' * width}\n""" - f"""{'Total timesteps:':>{pad}} {self.tot_timesteps}\n""" - f"""{'Iteration time:':>{pad}} {iteration_time:.2f}s\n""" - f"""{'Time elapsed:':>{pad}} {time.strftime("%H:%M:%S", time.gmtime(self.tot_time))}\n""" - f"""{'ETA:':>{pad}} {time.strftime( - "%H:%M:%S", - time.gmtime( - self.tot_time / (locs['it'] - locs['start_iter'] + 1) - * (locs['start_iter'] + locs['num_learning_iterations'] - locs['it']) - ) - )}\n""" - ) - print(log_string) - - def save(self, path: str, infos=None): - # -- Save model - saved_dict = { - "model_state_dict": self.alg.policy.state_dict(), - "optimizer_state_dict": self.alg.optimizer.state_dict(), - "iter": self.current_learning_iteration, - "infos": infos, - } - # -- Save RND model if used - if hasattr(self.alg, "rnd") and self.alg.rnd: - saved_dict["rnd_state_dict"] = self.alg.rnd.state_dict() - saved_dict["rnd_optimizer_state_dict"] = self.alg.rnd_optimizer.state_dict() - torch.save(saved_dict, path) - - # upload model to external logging service - if self.logger_type in ["neptune", "wandb"] and not self.disable_logs: - self.writer.save_model(path, self.current_learning_iteration) - - def load(self, path: str, load_optimizer: bool = True, map_location: str | None = None): - loaded_dict = torch.load(path, weights_only=False, map_location=map_location) - # -- Load model - resumed_training = self.alg.policy.load_state_dict(loaded_dict["model_state_dict"]) - # -- Load RND model if used - if hasattr(self.alg, "rnd") and self.alg.rnd: - self.alg.rnd.load_state_dict(loaded_dict["rnd_state_dict"]) - # -- load optimizer if used - if load_optimizer and resumed_training: - # -- algorithm optimizer - self.alg.optimizer.load_state_dict(loaded_dict["optimizer_state_dict"]) - # -- RND optimizer if used - if hasattr(self.alg, "rnd") and self.alg.rnd: - self.alg.rnd_optimizer.load_state_dict(loaded_dict["rnd_optimizer_state_dict"]) - # -- load current learning iteration - if resumed_training: - self.current_learning_iteration = loaded_dict["iter"] - return loaded_dict["infos"] - - def get_inference_policy(self, device=None): - self.eval_mode() # switch to evaluation mode (dropout for example) - if device is not None: - self.alg.policy.to(device) - return self.alg.policy.act_inference - - def train_mode(self): - # -- PPO - self.alg.policy.train() - # -- RND - if hasattr(self.alg, "rnd") and self.alg.rnd: - self.alg.rnd.train() - - def eval_mode(self): - # -- PPO - self.alg.policy.eval() - # -- RND - if hasattr(self.alg, "rnd") and self.alg.rnd: - self.alg.rnd.eval() - - def add_git_repo_to_log(self, repo_file_path): - self.git_status_repos.append(repo_file_path) - - """ - Helper functions. - """ - - def _configure_multi_gpu(self): - """Configure multi-gpu training.""" - # check if distributed training is enabled - self.gpu_world_size = int(os.getenv("WORLD_SIZE", "1")) - self.is_distributed = self.gpu_world_size > 1 - - # if not distributed training, set local and global rank to 0 and return - if not self.is_distributed: - self.gpu_local_rank = 0 - self.gpu_global_rank = 0 - self.multi_gpu_cfg = None - return - - # get rank and world size - self.gpu_local_rank = int(os.getenv("LOCAL_RANK", "0")) - self.gpu_global_rank = int(os.getenv("RANK", "0")) - - # make a configuration dictionary - self.multi_gpu_cfg = { - "global_rank": self.gpu_global_rank, # rank of the main process - "local_rank": self.gpu_local_rank, # rank of the current process - "world_size": self.gpu_world_size, # total number of processes - } - - # check if user has device specified for local rank - if self.device != f"cuda:{self.gpu_local_rank}": - raise ValueError( - f"Device '{self.device}' does not match expected device for local rank '{self.gpu_local_rank}'." - ) - # validate multi-gpu configuration - if self.gpu_local_rank >= self.gpu_world_size: - raise ValueError( - f"Local rank '{self.gpu_local_rank}' is greater than or equal to world size '{self.gpu_world_size}'." - ) - if self.gpu_global_rank >= self.gpu_world_size: - raise ValueError( - f"Global rank '{self.gpu_global_rank}' is greater than or equal to world size '{self.gpu_world_size}'." - ) - - # initialize torch distributed - torch.distributed.init_process_group(backend="nccl", rank=self.gpu_global_rank, world_size=self.gpu_world_size) - # set device to the local rank - torch.cuda.set_device(self.gpu_local_rank) - - def _construct_algorithm(self, obs) -> PPO: - """Construct the actor-critic algorithm.""" - # resolve RND config - self.alg_cfg = resolve_rnd_config(self.alg_cfg, obs, self.cfg["obs_groups"], self.env) - - # resolve symmetry config - self.alg_cfg = resolve_symmetry_config(self.alg_cfg, self.env) - - # resolve deprecated normalization config - if self.cfg.get("empirical_normalization") is not None: - warnings.warn( - "The `empirical_normalization` parameter is deprecated. Please set `actor_obs_normalization` and " - "`critic_obs_normalization` as part of the `policy` configuration instead.", - DeprecationWarning, - ) - if self.policy_cfg.get("actor_obs_normalization") is None: - self.policy_cfg["actor_obs_normalization"] = self.cfg["empirical_normalization"] - if self.policy_cfg.get("critic_obs_normalization") is None: - self.policy_cfg["critic_obs_normalization"] = self.cfg["empirical_normalization"] - - # initialize the actor-critic - actor_critic_class = eval(self.policy_cfg.pop("class_name")) - actor_critic: ActorCritic | ActorCriticRecurrent = actor_critic_class( - obs, self.cfg["obs_groups"], self.env.num_actions, **self.policy_cfg - ).to(self.device) - - # initialize the algorithm - alg_class = eval(self.alg_cfg.pop("class_name")) - alg: PPO = alg_class(actor_critic, device=self.device, **self.alg_cfg, multi_gpu_cfg=self.multi_gpu_cfg) - - # initialize the storage - alg.init_storage( - "rl", - self.env.num_envs, - self.num_steps_per_env, - obs, - [self.env.num_actions], - ) - - return alg - - def _prepare_logging_writer(self): - """Prepares the logging writers.""" - if self.log_dir is not None and self.writer is None and not self.disable_logs: - # Launch either Tensorboard or Neptune & Tensorboard summary writer(s), default: Tensorboard. - self.logger_type = self.cfg.get("logger", "tensorboard") - self.logger_type = self.logger_type.lower() - - if self.logger_type == "neptune": - from rsl_rl.utils.neptune_utils import NeptuneSummaryWriter - - self.writer = NeptuneSummaryWriter(log_dir=self.log_dir, flush_secs=10, cfg=self.cfg) - self.writer.log_config(self.env.cfg, self.cfg, self.alg_cfg, self.policy_cfg) - elif self.logger_type == "wandb": - from rsl_rl.utils.wandb_utils import WandbSummaryWriter - - self.writer = WandbSummaryWriter(log_dir=self.log_dir, flush_secs=10, cfg=self.cfg) - self.writer.log_config(self.env.cfg, self.cfg, self.alg_cfg, self.policy_cfg) - elif self.logger_type == "tensorboard": - from torch.utils.tensorboard import SummaryWriter - - self.writer = SummaryWriter(log_dir=self.log_dir, flush_secs=10) - else: - raise ValueError("Logger type not found. Please choose 'neptune', 'wandb' or 'tensorboard'.") diff --git a/rsl_rl/build/lib/rsl_rl/storage/__init__.py b/rsl_rl/build/lib/rsl_rl/storage/__init__.py deleted file mode 100644 index 1624330..0000000 --- a/rsl_rl/build/lib/rsl_rl/storage/__init__.py +++ /dev/null @@ -1,10 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -"""Implementation of transitions storage for RL-agent.""" - -from .rollout_storage import RolloutStorage -from .replay_buffer_multi import ReplayBufferMulti -__all__ = ["RolloutStorage", "ReplayBufferMulti"] diff --git a/rsl_rl/build/lib/rsl_rl/storage/replay_buffer_multi.py b/rsl_rl/build/lib/rsl_rl/storage/replay_buffer_multi.py deleted file mode 100644 index 6462b8e..0000000 --- a/rsl_rl/build/lib/rsl_rl/storage/replay_buffer_multi.py +++ /dev/null @@ -1,38 +0,0 @@ -import torch -import numpy as np - - -class ReplayBufferMulti: - """Fixed-size buffer to store experience tuples.""" - - def __init__(self, obs_dim, buffer_size, num_amp_frames, device): - """Initialize a ReplayBuffer object. - Arguments: - buffer_size (int): maximum size of buffer - """ - self.states = torch.zeros(buffer_size, num_amp_frames, obs_dim).to(device) - self.num_amp_frames = num_amp_frames - self.buffer_size = buffer_size - self.device = device - - self.step = 0 - self.num_samples = 0 - - def insert(self, states): - """Add new states to memory.""" - num_states = states.shape[0] - start_idx = self.step - end_idx = self.step + num_states - if end_idx > self.buffer_size: - self.states[self.step:self.buffer_size] = states[:self.buffer_size - self.step] - self.states[:end_idx - self.buffer_size] = states[self.buffer_size - self.step:] - else: - self.states[start_idx:end_idx] = states - - self.num_samples = min(self.buffer_size, max(end_idx, self.num_samples)) - self.step = (self.step + num_states) % self.buffer_size - - def feed_forward_generator(self, num_mini_batch, mini_batch_size): - for _ in range(num_mini_batch): - sample_idxs = np.random.choice(self.num_samples, size=mini_batch_size) - yield (self.states[sample_idxs].to(self.device)) diff --git a/rsl_rl/build/lib/rsl_rl/storage/rollout_storage.py b/rsl_rl/build/lib/rsl_rl/storage/rollout_storage.py deleted file mode 100644 index e9309b3..0000000 --- a/rsl_rl/build/lib/rsl_rl/storage/rollout_storage.py +++ /dev/null @@ -1,260 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -from __future__ import annotations - -import torch -from tensordict import TensorDict - -from rsl_rl.utils import split_and_pad_trajectories - - -class RolloutStorage: - class Transition: - def __init__(self): - self.observations = None - self.actions = None - self.privileged_actions = None - self.rewards = None - self.dones = None - self.values = None - self.actions_log_prob = None - self.action_mean = None - self.action_sigma = None - self.hidden_states = None - - def clear(self): - self.__init__() - - def __init__( - self, - training_type, - num_envs, - num_transitions_per_env, - obs, - actions_shape, - device="cpu", - ): - # store inputs - self.training_type = training_type - self.device = device - self.num_transitions_per_env = num_transitions_per_env - self.num_envs = num_envs - self.actions_shape = actions_shape - - # Core - self.observations = TensorDict( - {key: torch.zeros(num_transitions_per_env, *value.shape, device=device) for key, value in obs.items()}, - batch_size=[num_transitions_per_env, num_envs], - device=self.device, - ) - self.rewards = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device) - self.actions = torch.zeros(num_transitions_per_env, num_envs, *actions_shape, device=self.device) - self.dones = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device).byte() - - # for distillation - if training_type == "distillation": - self.privileged_actions = torch.zeros(num_transitions_per_env, num_envs, *actions_shape, device=self.device) - - # for reinforcement learning - if training_type == "rl": - self.values = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device) - self.actions_log_prob = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device) - self.mu = torch.zeros(num_transitions_per_env, num_envs, *actions_shape, device=self.device) - self.sigma = torch.zeros(num_transitions_per_env, num_envs, *actions_shape, device=self.device) - self.returns = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device) - self.advantages = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device) - - # For RNN networks - self.saved_hidden_states_a = None - self.saved_hidden_states_c = None - - # counter for the number of transitions stored - self.step = 0 - - def add_transitions(self, transition: Transition): - # check if the transition is valid - if self.step >= self.num_transitions_per_env: - raise OverflowError("Rollout buffer overflow! You should call clear() before adding new transitions.") - - # Core - self.observations[self.step].copy_(transition.observations) - self.actions[self.step].copy_(transition.actions) - self.rewards[self.step].copy_(transition.rewards.view(-1, 1)) - self.dones[self.step].copy_(transition.dones.view(-1, 1)) - - # for distillation - if self.training_type == "distillation": - self.privileged_actions[self.step].copy_(transition.privileged_actions) - - # for reinforcement learning - if self.training_type == "rl": - self.values[self.step].copy_(transition.values) - self.actions_log_prob[self.step].copy_(transition.actions_log_prob.view(-1, 1)) - self.mu[self.step].copy_(transition.action_mean) - self.sigma[self.step].copy_(transition.action_sigma) - - # For RNN networks - self._save_hidden_states(transition.hidden_states) - - # increment the counter - self.step += 1 - - def _save_hidden_states(self, hidden_states): - if hidden_states is None or hidden_states == (None, None): - return - # make a tuple out of GRU hidden state sto match the LSTM format - hid_a = hidden_states[0] if isinstance(hidden_states[0], tuple) else (hidden_states[0],) - hid_c = hidden_states[1] if isinstance(hidden_states[1], tuple) else (hidden_states[1],) - # initialize if needed - if self.saved_hidden_states_a is None: - self.saved_hidden_states_a = [ - torch.zeros(self.observations.shape[0], *hid_a[i].shape, device=self.device) for i in range(len(hid_a)) - ] - self.saved_hidden_states_c = [ - torch.zeros(self.observations.shape[0], *hid_c[i].shape, device=self.device) for i in range(len(hid_c)) - ] - # copy the states - for i in range(len(hid_a)): - self.saved_hidden_states_a[i][self.step].copy_(hid_a[i]) - self.saved_hidden_states_c[i][self.step].copy_(hid_c[i]) - - def clear(self): - self.step = 0 - - def compute_returns(self, last_values, gamma, lam, normalize_advantage: bool = True): - advantage = 0 - for step in reversed(range(self.num_transitions_per_env)): - # if we are at the last step, bootstrap the return value - if step == self.num_transitions_per_env - 1: - next_values = last_values - else: - next_values = self.values[step + 1] - # 1 if we are not in a terminal state, 0 otherwise - next_is_not_terminal = 1.0 - self.dones[step].float() - # TD error: r_t + gamma * V(s_{t+1}) - V(s_t) - delta = self.rewards[step] + next_is_not_terminal * gamma * next_values - self.values[step] - # Advantage: A(s_t, a_t) = delta_t + gamma * lambda * A(s_{t+1}, a_{t+1}) - advantage = delta + next_is_not_terminal * gamma * lam * advantage - # Return: R_t = A(s_t, a_t) + V(s_t) - self.returns[step] = advantage + self.values[step] - - # Compute the advantages - self.advantages = self.returns - self.values - # Normalize the advantages if flag is set - # This is to prevent double normalization (i.e. if per minibatch normalization is used) - if normalize_advantage: - self.advantages = (self.advantages - self.advantages.mean()) / (self.advantages.std() + 1e-8) - - # for distillation - def generator(self): - if self.training_type != "distillation": - raise ValueError("This function is only available for distillation training.") - - for i in range(self.num_transitions_per_env): - yield self.observations[i], self.actions[i], self.privileged_actions[i], self.dones[i] - - # for reinforcement learning with feedforward networks - def mini_batch_generator(self, num_mini_batches, num_epochs=8): - if self.training_type != "rl": - raise ValueError("This function is only available for reinforcement learning training.") - batch_size = self.num_envs * self.num_transitions_per_env - mini_batch_size = batch_size // num_mini_batches - indices = torch.randperm(num_mini_batches * mini_batch_size, requires_grad=False, device=self.device) - - # Core - observations = self.observations.flatten(0, 1) - actions = self.actions.flatten(0, 1) - values = self.values.flatten(0, 1) - returns = self.returns.flatten(0, 1) - - # For PPO - old_actions_log_prob = self.actions_log_prob.flatten(0, 1) - advantages = self.advantages.flatten(0, 1) - old_mu = self.mu.flatten(0, 1) - old_sigma = self.sigma.flatten(0, 1) - - for epoch in range(num_epochs): - for i in range(num_mini_batches): - # Select the indices for the mini-batch - start = i * mini_batch_size - end = (i + 1) * mini_batch_size - batch_idx = indices[start:end] - - # Create the mini-batch - # -- Core - obs_batch = observations[batch_idx] - actions_batch = actions[batch_idx] - - # -- For PPO - target_values_batch = values[batch_idx] - returns_batch = returns[batch_idx] - old_actions_log_prob_batch = old_actions_log_prob[batch_idx] - advantages_batch = advantages[batch_idx] - old_mu_batch = old_mu[batch_idx] - old_sigma_batch = old_sigma[batch_idx] - - # yield the mini-batch - yield obs_batch, actions_batch, target_values_batch, advantages_batch, returns_batch, old_actions_log_prob_batch, old_mu_batch, old_sigma_batch, ( - None, - None, - ), None - - # for reinfrocement learning with recurrent networks - def recurrent_mini_batch_generator(self, num_mini_batches, num_epochs=8): - if self.training_type != "rl": - raise ValueError("This function is only available for reinforcement learning training.") - padded_obs_trajectories, trajectory_masks = split_and_pad_trajectories(self.observations, self.dones) - - mini_batch_size = self.num_envs // num_mini_batches - for ep in range(num_epochs): - first_traj = 0 - for i in range(num_mini_batches): - start = i * mini_batch_size - stop = (i + 1) * mini_batch_size - - dones = self.dones.squeeze(-1) - last_was_done = torch.zeros_like(dones, dtype=torch.bool) - last_was_done[1:] = dones[:-1] - last_was_done[0] = True - trajectories_batch_size = torch.sum(last_was_done[:, start:stop]) - last_traj = first_traj + trajectories_batch_size - - masks_batch = trajectory_masks[:, first_traj:last_traj] - obs_batch = padded_obs_trajectories[:, first_traj:last_traj] - actions_batch = self.actions[:, start:stop] - old_mu_batch = self.mu[:, start:stop] - old_sigma_batch = self.sigma[:, start:stop] - returns_batch = self.returns[:, start:stop] - advantages_batch = self.advantages[:, start:stop] - values_batch = self.values[:, start:stop] - old_actions_log_prob_batch = self.actions_log_prob[:, start:stop] - - # reshape to [num_envs, time, num layers, hidden dim] (original shape: [time, num_layers, num_envs, hidden_dim]) - # then take only time steps after dones (flattens num envs and time dimensions), - # take a batch of trajectories and finally reshape back to [num_layers, batch, hidden_dim] - last_was_done = last_was_done.permute(1, 0) - hid_a_batch = [ - saved_hidden_states.permute(2, 0, 1, 3)[last_was_done][first_traj:last_traj] - .transpose(1, 0) - .contiguous() - for saved_hidden_states in self.saved_hidden_states_a - ] - hid_c_batch = [ - saved_hidden_states.permute(2, 0, 1, 3)[last_was_done][first_traj:last_traj] - .transpose(1, 0) - .contiguous() - for saved_hidden_states in self.saved_hidden_states_c - ] - # remove the tuple for GRU - hid_a_batch = hid_a_batch[0] if len(hid_a_batch) == 1 else hid_a_batch - hid_c_batch = hid_c_batch[0] if len(hid_c_batch) == 1 else hid_c_batch - - yield obs_batch, actions_batch, values_batch, advantages_batch, returns_batch, old_actions_log_prob_batch, old_mu_batch, old_sigma_batch, ( - hid_a_batch, - hid_c_batch, - ), masks_batch - - first_traj = last_traj diff --git a/rsl_rl/build/lib/rsl_rl/utils/__init__.py b/rsl_rl/build/lib/rsl_rl/utils/__init__.py deleted file mode 100644 index f5781f1..0000000 --- a/rsl_rl/build/lib/rsl_rl/utils/__init__.py +++ /dev/null @@ -1,13 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -"""Helper functions.""" - -from .utils import * -from .motion_loader_g1 import G1_AMPLoader - -__all__ = [ - "G1_AMPLoader", -] \ No newline at end of file diff --git a/rsl_rl/build/lib/rsl_rl/utils/motion_loader_g1.py b/rsl_rl/build/lib/rsl_rl/utils/motion_loader_g1.py deleted file mode 100644 index 677c630..0000000 --- a/rsl_rl/build/lib/rsl_rl/utils/motion_loader_g1.py +++ /dev/null @@ -1,388 +0,0 @@ -import os -from os.path import join as pjoin -import glob -import json -import logging - -import torch -import numpy as np -from pybullet_utils import transformations - -from rsl_rl.utils import motion_util - -_EPS = np.finfo(float).eps * 4.0 -def quaternion_slerp(q0, q1, fraction, spin=0, shortestpath=True): - """Batch quaternion spherical linear interpolation.""" - - out = torch.zeros_like(q0) - - zero_mask = torch.isclose(fraction, torch.zeros_like(fraction)).squeeze() - ones_mask = torch.isclose(fraction, torch.ones_like(fraction)).squeeze() - out[zero_mask] = q0[zero_mask] - out[ones_mask] = q1[ones_mask] - - d = torch.sum(q0 * q1, dim=-1, keepdim=True) - dist_mask = (torch.abs(torch.abs(d) - 1.0) < _EPS).squeeze() - out[dist_mask] = q0[dist_mask] - - if shortestpath: - d_old = torch.clone(d) - d = torch.where(d_old < 0, -d, d) - q1 = torch.where(d_old < 0, -q1, q1) - - angle = torch.acos(d) + spin * torch.pi - angle_mask = (torch.abs(angle) < _EPS).squeeze() - out[angle_mask] = q0[angle_mask] - - final_mask = torch.logical_or(zero_mask, ones_mask) - final_mask = torch.logical_or(final_mask, dist_mask) - final_mask = torch.logical_or(final_mask, angle_mask) - final_mask = torch.logical_not(final_mask) - - isin = 1.0 / angle - q0 *= torch.sin((1.0 - fraction) * angle) * isin - q1 *= torch.sin(fraction * angle) * isin - q0 += q1 - out[final_mask] = q0[final_mask] - return out - - -class G1_AMPLoader: - - def __init__( - self, - device, - time_between_frames, - motion_files, - preload_transitions=False, - num_preload_transitions=1000000, - num_frames=5, - ): - """Expert dataset provides AMP observations from Dog mocap dataset. - - time_between_frames: Amount of time in seconds between transition. - """ - self.device = device - self.time_between_frames = time_between_frames - self.num_frames = num_frames - - # Values to store for each trajectory. - self.trajectories = [] - self.trajectories_full = [] - self.trajectory_names = [] - self.trajectory_idxs = [] - self.trajectory_lens = [] # Traj length in seconds. - self.trajectory_weights = [] - self.trajectory_frame_durations = [] - self.trajectory_num_frames = [] - self.motion_dir = motion_files - # import ipdb; ipdb.set_trace() - for i, motion_file in enumerate(os.listdir(motion_files)): - self.trajectory_names.append(motion_file) - motion_path = pjoin(motion_files, motion_file) - motion_data = np.load(motion_path, allow_pickle=True) - motion_data_processed = np.zeros((motion_data.shape[0],36)) - - for f_i in range(motion_data.shape[0]): - motion_data_processed[f_i, :3] = motion_data[f_i, :3] # base pos - motion_data_processed[f_i, 3:7] = motion_data[f_i, 3:7] # base quat (wxyz) - motion_data_processed[f_i, 7:35] = motion_data[f_i, 7:35] # base vel - ''' - NOTE The order of motion_data_processed is - base pos 0:3, - base quat 3:7, wxyz - dof pos 7:36, (mujoco joint order) - ''' - self.trajectories.append(torch.tensor( - motion_data_processed[:, 7:], - dtype=torch.float32, - device=self.device - )) - - self.trajectories_full.append(torch.tensor( - motion_data_processed, - dtype=torch.float32, - device=self.device - )) - - self.trajectory_idxs.append(i) - self.trajectory_weights.append(1 / len(os.listdir(motion_files))) - frame_duration = 1 / 50 - - self.trajectory_frame_durations.append(frame_duration) - traj_len = (motion_data_processed.shape[0] - 1) * frame_duration # seconds - self.trajectory_lens.append(traj_len) - self.trajectory_num_frames.append(float(motion_data_processed.shape[0])) - print(f"Loaded {traj_len}s. motion from {motion_file}.") - - # Trajectory weights are used to sample some trajectories more than others. - self.trajectory_weights = np.array(self.trajectory_weights) / np.sum(self.trajectory_weights) - self.trajectory_frame_durations = np.array(self.trajectory_frame_durations) - self.trajectory_lens = np.array(self.trajectory_lens) - self.trajectory_num_frames = np.array(self.trajectory_num_frames) - - # Preload transitions. - self.preload_transitions = preload_transitions - if self.preload_transitions: - print(f'Preloading {num_preload_transitions} transitions') - - traj_idxs = self.weighted_traj_idx_sample_batch(num_preload_transitions) - times = self.traj_time_sample_batch(traj_idxs) - self.preloaded_s_prior = self.get_full_frame_at_time_batch(traj_idxs, times - self.time_between_frames) - self.preloaded_s = self.get_full_frame_at_time_batch(traj_idxs, times) - self.preloaded_s_next = self.get_full_frame_at_time_batch(traj_idxs, times + self.time_between_frames) - print(f'Finished preloading') - - # 预加载多帧数据 - self.preloaded_frames = [] - for i in range(self.num_frames): - frame_time = times + (i - (self.num_frames - 2)) * self.time_between_frames - full_frame = self.get_full_frame_at_time_batch(traj_idxs, frame_time) - # 预处理:提前提取并连接需要的列(7:26 和 29:33),避免每次生成时重复切片 - processed_frame = torch.cat([ - full_frame[:, 7:26], - full_frame[:, 29:33] - ], dim=-1) - self.preloaded_frames.append(processed_frame) - print(f'Finished preloading multiple frames') - - self.all_trajectories_full = torch.vstack(self.trajectories_full) - - def weighted_traj_idx_sample(self): - """Get traj idx via weighted sampling.""" - return np.random.choice( - self.trajectory_idxs, p=self.trajectory_weights) - - def weighted_traj_idx_sample_batch(self, size): - """Batch sample traj idxs.""" - return np.random.choice( - self.trajectory_idxs, size=size, p=self.trajectory_weights, - replace=True) - - def traj_time_sample(self, traj_idx): - """Sample random time for traj.""" - subst = self.time_between_frames + self.trajectory_frame_durations[traj_idx] - return max( - 0, (self.trajectory_lens[traj_idx] * np.random.uniform() - subst)) - - def traj_time_sample_batch(self, traj_idxs): - """Sample random time for multiple trajectories.""" - subst = self.time_between_frames + self.trajectory_frame_durations[traj_idxs] - time_samples = self.trajectory_lens[traj_idxs] * np.random.uniform(size=len(traj_idxs)) - subst - return np.maximum(np.zeros_like(time_samples), time_samples) - - def slerp(self, val0, val1, blend): - return (1.0 - blend) * val0 + blend * val1 - - def get_trajectory(self, traj_idx): - """Returns trajectory of AMP observations.""" - return self.trajectories_full[traj_idx] - - def get_frame_at_time(self, traj_idx, time): - """Returns frame for the given trajectory at the specified time.""" - p = float(time) / self.trajectory_lens[traj_idx] - n = self.trajectories[traj_idx].shape[0] - idx_low, idx_high = int(np.floor(p * n)), int(np.ceil(p * n)) - frame_start = self.trajectories[traj_idx][idx_low] - frame_end = self.trajectories[traj_idx][idx_high] - blend = p * n - idx_low - return self.slerp(frame_start, frame_end, blend) - - def get_frame_at_time_batch(self, traj_idxs, times): - """Returns frame for the given trajectory at the specified time.""" - p = times / self.trajectory_lens[traj_idxs] - n = self.trajectory_num_frames[traj_idxs] - idx_low, idx_high = np.floor(p * n).astype(np.int32), np.ceil(p * n).astype(np.int32) - all_frame_starts = torch.zeros(len(traj_idxs), self.observation_dim, device=self.device) - all_frame_ends = torch.zeros(len(traj_idxs), self.observation_dim, device=self.device) - for traj_idx in set(traj_idxs): - trajectory = self.trajectories[traj_idx] - traj_mask = traj_idxs == traj_idx - all_frame_starts[traj_mask] = trajectory[idx_low[traj_mask]] - all_frame_ends[traj_mask] = trajectory[idx_high[traj_mask]] - blend = torch.tensor(p * n - idx_low, device=self.device, dtype=torch.float32).unsqueeze(-1) - return self.slerp(all_frame_starts, all_frame_ends, blend) - - def get_full_frame_at_time(self, traj_idx, time): - """Returns full frame for the given trajectory at the specified time.""" - p = float(time) / self.trajectory_lens[traj_idx] - n = self.trajectories_full[traj_idx].shape[0] - idx_low, idx_high = int(np.floor(p * n)), int(np.ceil(p * n)) - frame_start = self.trajectories_full[traj_idx][idx_low] - frame_end = self.trajectories_full[traj_idx][idx_high] - blend = p * n - idx_low - print(idx_low, idx_high) - return self.blend_frame_pose(frame_start, frame_end, blend) - - def get_full_frame_at_time_batch(self, traj_idxs, times): - p = times / self.trajectory_lens[traj_idxs] - n = self.trajectory_num_frames[traj_idxs] - idx_low, idx_high = np.floor(p * n).astype(np.int32), np.ceil(p * n).astype(np.int32) - all_frame_pos_starts = torch.zeros(len(traj_idxs), 3, device=self.device) - all_frame_pos_ends = torch.zeros(len(traj_idxs), 3, device=self.device) - all_frame_rot_starts = torch.zeros(len(traj_idxs), 4, device=self.device) - all_frame_rot_ends = torch.zeros(len(traj_idxs), 4, device=self.device) - all_frame_amp_starts = torch.zeros(len(traj_idxs), 29, device=self.device) - all_frame_amp_ends = torch.zeros(len(traj_idxs), 29, device=self.device) - for traj_idx in set(traj_idxs): - trajectory = self.trajectories_full[traj_idx] - traj_mask = traj_idxs == traj_idx - all_frame_pos_starts[traj_mask] = G1_AMPLoader.get_root_pos_batch(trajectory[idx_low[traj_mask]]) - all_frame_pos_ends[traj_mask] = G1_AMPLoader.get_root_pos_batch(trajectory[idx_high[traj_mask]]) - all_frame_rot_starts[traj_mask] = G1_AMPLoader.get_root_rot_batch(trajectory[idx_low[traj_mask]]) - all_frame_rot_ends[traj_mask] = G1_AMPLoader.get_root_rot_batch(trajectory[idx_high[traj_mask]]) - all_frame_amp_starts[traj_mask] = trajectory[idx_low[traj_mask]][:, 7:36] # base vel3+ang3, dof vel23+ang23 - all_frame_amp_ends[traj_mask] = trajectory[idx_high[traj_mask]][:, 7:36] # base vel3+ang3, dof vel23+ang23 - blend = torch.tensor(p * n - idx_low, device=self.device, dtype=torch.float32).unsqueeze(-1) - pos_blend = self.slerp(all_frame_pos_starts, all_frame_pos_ends, blend) - rot_blend = quaternion_slerp(all_frame_rot_starts, all_frame_rot_ends, blend) - amp_blend = self.slerp(all_frame_amp_starts, all_frame_amp_ends, blend) - return torch.cat([pos_blend, rot_blend, amp_blend], dim=-1) - - def get_frame(self): - """Returns random frame.""" - traj_idx = self.weighted_traj_idx_sample() - sampled_time = self.traj_time_sample(traj_idx) - return self.get_frame_at_time(traj_idx, sampled_time) - - def get_full_frame(self): - """Returns random full frame.""" - traj_idx = self.weighted_traj_idx_sample() - sampled_time = self.traj_time_sample(traj_idx) - return self.get_full_frame_at_time(traj_idx, sampled_time) - - def get_full_frame_batch(self, num_frames): - if self.preload_transitions: - idxs = np.random.choice( - self.preloaded_s.shape[0], size=num_frames) - return self.preloaded_s[idxs] - else: - traj_idxs = self.weighted_traj_idx_sample_batch(num_frames) - times = self.traj_time_sample_batch(traj_idxs) - return self.get_full_frame_at_time_batch(traj_idxs, times) - - def blend_frame_pose(self, frame0, frame1, blend): - """Linearly interpolate between two frames, including orientation. - - Args: - frame0: First frame to be blended corresponds to (blend = 0). - frame1: Second frame to be blended corresponds to (blend = 1). - blend: Float between [0, 1], specifying the interpolation between - the two frames. - Returns: - An interpolation of the two frames. - """ - root_pos0, root_pos1 = G1_AMPLoader.get_root_pos(frame0), G1_AMPLoader.get_root_pos(frame1) - root_rot0, root_rot1 = G1_AMPLoader.get_root_rot(frame0), G1_AMPLoader.get_root_rot(frame1) - joints0, joints1 = G1_AMPLoader.get_joint_pose(frame0), G1_AMPLoader.get_joint_pose(frame1) - # tar_toe_pos_0, tar_toe_pos_1 = G1_AMPLoader.get_tar_toe_pos_local(frame0), G1_AMPLoader.get_tar_toe_pos_local(frame1) - linear_vel_0, linear_vel_1 = G1_AMPLoader.get_linear_vel(frame0), G1_AMPLoader.get_linear_vel(frame1) - angular_vel_0, angular_vel_1 = G1_AMPLoader.get_angular_vel(frame0), G1_AMPLoader.get_angular_vel(frame1) - joint_vel_0, joint_vel_1 = G1_AMPLoader.get_joint_vel(frame0), G1_AMPLoader.get_joint_vel(frame1) - - blend_root_pos = self.slerp(root_pos0, root_pos1, blend) - blend_root_rot = transformations.quaternion_slerp(root_rot0.cpu().numpy(), root_rot1.cpu().numpy(), blend) - blend_root_rot = torch.tensor(motion_util.standardize_quaternion(blend_root_rot),dtype=torch.float32, device=self.device) - blend_joints = self.slerp(joints0, joints1, blend) - # blend_tar_toe_pos = self.slerp(tar_toe_pos_0, tar_toe_pos_1, blend) - blend_linear_vel = self.slerp(linear_vel_0, linear_vel_1, blend) - blend_angular_vel = self.slerp(angular_vel_0, angular_vel_1, blend) - blend_joints_vel = self.slerp(joint_vel_0, joint_vel_1, blend) - - # return - # torch.cat([ - # blend_root_pos, blend_root_rot, blend_linear_vel, blend_angular_vel, blend_joints, blend_joints_vel]) - return torch.cat([blend_root_pos, blend_root_rot, blend_linear_vel, blend_angular_vel, blend_joints]) - - def feed_forward_generator_23dof_multi(self, num_mini_batch, mini_batch_size): - """Generates a batch of AMP transitions.""" - # import ipdb; ipdb.set_trace() - for _ in range(num_mini_batch): - if self.preload_transitions: - idxs = np.random.choice(self.preloaded_s.shape[0], size=mini_batch_size) - - frames = [] - for i in range(self.num_frames): - # 数据已在预加载时预处理,直接索引即可 - s = self.preloaded_frames[i][idxs] - frames.append(s) - else: - NotImplementedError('preload transition') - yield torch.stack(frames, dim=1) # [batch, num_frames, 16] - - - - - def quaternion_to_euler_array(self, quat): - # Ensure quaternion is in the correct format [x, y, z, w] - x, y, z, w =quat - - # Roll (x-axis rotation) - t0 = +2.0 * (w * x + y * z) - t1 = +1.0 - 2.0 * (x * x + y * y) - roll_x = np.arctan2(t0, t1) - - # Pitch (y-axis rotation) - t2 = +2.0 * (w * y - z * x) - t2 = np.clip(t2, -1.0, 1.0) - pitch_y = np.arcsin(t2) - - # Yaw (z-axis rotation) - t3 = +2.0 * (w * z + x * y) - t4 = +1.0 - 2.0 * (y * y + z * z) - yaw_z = np.arctan2(t3, t4) - - # Returns roll, pitch, yaw in a NumPy array in radians - return np.array([roll_x, pitch_y, yaw_z]) - - def euler_to_quaternion(self, root_rot): - roll, pitch, yaw = root_rot[0], root_rot[1], root_rot[2] - cy = np.cos(yaw * 0.5) - sy = np.sin(yaw * 0.5) - cp = np.cos(pitch * 0.5) - sp = np.sin(pitch * 0.5) - cr = np.cos(roll * 0.5) - sr = np.sin(roll * 0.5) - - qw = cy * cp * cr + sy * sp * sr - qx = cy * cp * sr - sy * sp * cr - qy = sy * cp * sr + cy * sp * cr - qz = sy * cp * cr - cy * sp * sr - - return np.array([qx, qy, qz, qw]) - - @property - def observation_dim(self): - """Size of AMP observations.""" - return self.trajectories[0].shape[1] + 1 - - @property - def num_motions(self): - return len(self.trajectory_names) - @staticmethod - def get_root_pos(pose): - return pose[0:3] - - @staticmethod - def get_root_pos_batch(poses): - return poses[:, 0:3] - - @staticmethod - def get_root_rot(pose): - return pose[3:7] - - @staticmethod - def get_root_rot_batch(poses): - return poses[:, 3:7] - - @staticmethod - def get_joint_pose_batch_12dof(poses): - return poses[:, 13:25] - - @staticmethod - def get_tar_toe_pos_local(pose): - return pose[G1_AMPLoader.TAR_TOE_POS_LOCAL_START_IDX:G1_AMPLoader.TAR_TOE_POS_LOCAL_END_IDX] - - @staticmethod - def get_tar_toe_pos_local_batch(poses): - return poses[:, G1_AMPLoader.TAR_TOE_POS_LOCAL_START_IDX:G1_AMPLoader.TAR_TOE_POS_LOCAL_END_IDX] diff --git a/rsl_rl/build/lib/rsl_rl/utils/motion_util.py b/rsl_rl/build/lib/rsl_rl/utils/motion_util.py deleted file mode 100644 index 3d49bfa..0000000 --- a/rsl_rl/build/lib/rsl_rl/utils/motion_util.py +++ /dev/null @@ -1,97 +0,0 @@ -# coding=utf-8 -# Copyright 2020 The Google Research Authors. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Utility functions for processing motion clips.""" - -import os -import inspect -currentdir = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe()))) -parentdir = os.path.dirname(os.path.dirname(currentdir)) -os.sys.path.insert(0, parentdir) - -import numpy as np - -from rsl_rl.utils import pose3d -# from pybullet_utils import transformations - - -def standardize_quaternion(q): - """Returns a quaternion where q.w >= 0 to remove redundancy due to q = -q. - - Args: - q: A quaternion to be standardized. - - Returns: - A quaternion with q.w >= 0. - - """ - if q[-1] < 0: - q = -q - return q - - -def normalize_rotation_angle(theta): - """Returns a rotation angle normalized between [-pi, pi]. - - Args: - theta: angle of rotation (radians). - - Returns: - An angle of rotation normalized between [-pi, pi]. - - """ - norm_theta = theta - if np.abs(norm_theta) > np.pi: - norm_theta = np.fmod(norm_theta, 2 * np.pi) - if norm_theta >= 0: - norm_theta += -2 * np.pi - else: - norm_theta += 2 * np.pi - - return norm_theta - - -def calc_heading(q): - """Returns the heading of a rotation q, specified as a quaternion. - - The heading represents the rotational component of q along the vertical - axis (z axis). - - Args: - q: A quaternion that the heading is to be computed from. - - Returns: - An angle representing the rotation about the z axis. - - """ - ref_dir = np.array([1, 0, 0]) - rot_dir = pose3d.QuaternionRotatePoint(ref_dir, q) - heading = np.arctan2(rot_dir[1], rot_dir[0]) - return heading - - -# def calc_heading_rot(q): -# """Return a quaternion representing the heading rotation of q along the vertical axis (z axis). - -# Args: -# q: A quaternion that the heading is to be computed from. - -# Returns: -# A quaternion representing the rotation about the z axis. - -# """ -# heading = calc_heading(q) -# q_heading = transformations.quaternion_about_axis(heading, [0, 0, 1]) -# return q_heading diff --git a/rsl_rl/build/lib/rsl_rl/utils/neptune_utils.py b/rsl_rl/build/lib/rsl_rl/utils/neptune_utils.py deleted file mode 100644 index 3796ec8..0000000 --- a/rsl_rl/build/lib/rsl_rl/utils/neptune_utils.py +++ /dev/null @@ -1,94 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -from __future__ import annotations - -import os -from dataclasses import asdict -from torch.utils.tensorboard import SummaryWriter - -try: - import neptune -except ModuleNotFoundError: - raise ModuleNotFoundError("neptune-client is required to log to Neptune.") - - -class NeptuneLogger: - def __init__(self, project, token): - self.run = neptune.init_run(project=project, api_token=token) - - def store_config(self, env_cfg, runner_cfg, alg_cfg, policy_cfg): - self.run["runner_cfg"] = runner_cfg - self.run["policy_cfg"] = policy_cfg - self.run["alg_cfg"] = alg_cfg - self.run["env_cfg"] = asdict(env_cfg) - - -class NeptuneSummaryWriter(SummaryWriter): - """Summary writer for Neptune.""" - - def __init__(self, log_dir: str, flush_secs: int, cfg): - super().__init__(log_dir, flush_secs) - - try: - project = cfg["neptune_project"] - except KeyError: - raise KeyError("Please specify neptune_project in the runner config, e.g. legged_gym.") - - try: - token = os.environ["NEPTUNE_API_TOKEN"] - except KeyError: - raise KeyError( - "Neptune api token not found. Please run or add to ~/.bashrc: export NEPTUNE_API_TOKEN=YOUR_API_TOKEN" - ) - - try: - entity = os.environ["NEPTUNE_USERNAME"] - except KeyError: - raise KeyError( - "Neptune username not found. Please run or add to ~/.bashrc: export NEPTUNE_USERNAME=YOUR_USERNAME" - ) - - neptune_project = entity + "/" + project - - self.neptune_logger = NeptuneLogger(neptune_project, token) - - self.name_map = { - "Train/mean_reward/time": "Train/mean_reward_time", - "Train/mean_episode_length/time": "Train/mean_episode_length_time", - } - - run_name = os.path.split(log_dir)[-1] - - self.neptune_logger.run["log_dir"].log(run_name) - - def _map_path(self, path): - if path in self.name_map: - return self.name_map[path] - else: - return path - - def add_scalar(self, tag, scalar_value, global_step=None, walltime=None, new_style=False): - super().add_scalar( - tag, - scalar_value, - global_step=global_step, - walltime=walltime, - new_style=new_style, - ) - self.neptune_logger.run[self._map_path(tag)].log(scalar_value, step=global_step) - - def stop(self): - self.neptune_logger.run.stop() - - def log_config(self, env_cfg, runner_cfg, alg_cfg, policy_cfg): - self.neptune_logger.store_config(env_cfg, runner_cfg, alg_cfg, policy_cfg) - - def save_model(self, model_path, iter): - self.neptune_logger.run["model/saved_model_" + str(iter)].upload(model_path) - - def save_file(self, path, iter=None): - name = path.rsplit("/", 1)[-1].split(".")[0] - self.neptune_logger.run["git_diff/" + name].upload(path) diff --git a/rsl_rl/build/lib/rsl_rl/utils/pose3d.py b/rsl_rl/build/lib/rsl_rl/utils/pose3d.py deleted file mode 100644 index a4cba1e..0000000 --- a/rsl_rl/build/lib/rsl_rl/utils/pose3d.py +++ /dev/null @@ -1,283 +0,0 @@ -# coding=utf-8 -# Copyright 2020 The Google Research Authors. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Utilities for 3D pose conversion.""" -import math -import numpy as np - -# from pybullet_utils import transformations - -VECTOR3_0 = np.zeros(3, dtype=np.float64) -VECTOR3_1 = np.ones(3, dtype=np.float64) -VECTOR3_X = np.array([1, 0, 0], dtype=np.float64) -VECTOR3_Y = np.array([0, 1, 0], dtype=np.float64) -VECTOR3_Z = np.array([0, 0, 1], dtype=np.float64) - -# QUATERNION_IDENTITY is the multiplicative identity 1.0 + 0i + 0j + 0k. -# When interpreted as a rotation, it is the identity rotation. -QUATERNION_IDENTITY = np.array([0.0, 0.0, 0.0, 1.0], dtype=np.float64) - - -def Vector3RandomNormal(sigma, mu=VECTOR3_0): - """Returns a random 3D vector from a normal distribution. - - Each component is selected independently from a normal distribution. - - Args: - sigma: Scale (or stddev) of distribution for all variables. - mu: Mean of distribution for each variable. - - Returns: - A 3D vector in a numpy array. - """ - - random_v3 = np.random.normal(scale=sigma, size=3) + mu - return random_v3 - - -def Vector3RandomUniform(low=VECTOR3_0, high=VECTOR3_1): - """Returns a 3D vector selected uniformly from the input box. - - Args: - low: The min-value corner of the box. - high: The max-value corner of the box. - - Returns: - A 3D vector in a numpy array. - """ - - random_x = np.random.uniform(low=low[0], high=high[0]) - random_y = np.random.uniform(low=low[1], high=high[1]) - random_z = np.random.uniform(low=low[2], high=high[2]) - return np.array([random_x, random_y, random_z]) - - -def Vector3RandomUnit(): - """Returns a random 3D vector with unit length. - - Generates a 3D vector selected uniformly from the unit sphere. - - Returns: - A normalized 3D vector in a numpy array. - """ - longitude = np.random.uniform(low=-math.pi, high=math.pi) - sin_latitude = np.random.uniform(low=-1.0, high=1.0) - cos_latitude = math.sqrt(1.0 - sin_latitude * sin_latitude) - x = math.cos(longitude) * cos_latitude - y = math.sin(longitude) * cos_latitude - z = sin_latitude - return np.array([x, y, z], dtype=np.float64) - - -def QuaternionNormalize(q): - """Normalizes the quaternion to length 1. - - Divides the quaternion by its magnitude. If the magnitude is too - small, returns the quaternion identity value (1.0). - - Args: - q: A quaternion to be normalized. - - Raises: - ValueError: If input quaternion has length near zero. - - Returns: - A quaternion with magnitude 1 in a numpy array [x, y, z, w]. - - """ - q_norm = np.linalg.norm(q) - if np.isclose(q_norm, 0.0): - raise ValueError( - 'Quaternion may not be zero in QuaternionNormalize: |q| = %f, q = %s' % - (q_norm, q)) - return q / q_norm - - -def QuaternionFromAxisAngle(axis, angle): - """Returns a quaternion that generates the given axis-angle rotation. - - Returns the quaternion: sin(angle/2) * axis + cos(angle/2). - - Args: - axis: Axis of rotation, a 3D vector in a numpy array. - angle: The angle of rotation (radians). - - Raises: - ValueError: If input axis is not a normalizable 3D vector. - - Returns: - A unit quaternion in a numpy array. - - """ - if len(axis) != 3: - raise ValueError('Axis vector should have three components: %s' % axis) - axis_norm = np.linalg.norm(axis) - if np.isclose(axis_norm, 0.0): - raise ValueError('Axis vector may not have zero length: |v| = %f, v = %s' % - (axis_norm, axis)) - half_angle = angle * 0.5 - q = np.zeros(4, dtype=np.float64) - q[0:3] = axis - q[0:3] *= math.sin(half_angle) / axis_norm - q[3] = math.cos(half_angle) - return q - - -def QuaternionToAxisAngle(quat, default_axis=VECTOR3_Z, direction_axis=None): - """Calculates axis and angle of rotation performed by a quaternion. - - Calculates the axis and angle of the rotation performed by the quaternion. - The quaternion should have four values and be normalized. - - Args: - quat: Unit quaternion in a numpy array. - default_axis: 3D vector axis used if the rotation is near to zero. Without - this default, small rotations would result in an exception. It is - reasonable to use a default axis for tiny rotations, because zero angle - rotations about any axis are equivalent. - direction_axis: Used to disambiguate rotation directions. If the - direction_axis is specified, the axis of the rotation will be chosen such - that its inner product with the direction_axis is non-negative. - - Raises: - ValueError: If quat is not a normalized quaternion. - - Returns: - axis: Axis of rotation. - angle: Angle in radians. - """ - if len(quat) != 4: - raise ValueError( - 'Quaternion should have four components [x, y, z, w]: %s' % quat) - if not np.isclose(1.0, np.linalg.norm(quat)): - raise ValueError('Quaternion should have unit length: |q| = %f, q = %s' % - (np.linalg.norm(quat), quat)) - axis = quat[:3].copy() - axis_norm = np.linalg.norm(axis) - min_axis_norm = 1e-8 - if axis_norm < min_axis_norm: - axis = default_axis - if len(default_axis) != 3: - raise ValueError('Axis vector should have three components: %s' % axis) - if not np.isclose(np.linalg.norm(axis), 1.0): - raise ValueError('Axis vector should have unit length: |v| = %f, v = %s' % - (np.linalg.norm(axis), axis)) - else: - axis /= axis_norm - sin_half_angle = axis_norm - if direction_axis is not None and np.inner(axis, direction_axis) < 0: - sin_half_angle = -sin_half_angle - axis = -axis - cos_half_angle = quat[3] - half_angle = math.atan2(sin_half_angle, cos_half_angle) - angle = half_angle * 2 - return axis, angle - - -def QuaternionRandomRotation(max_angle=math.pi): - """Creates a random small rotation around a random axis. - - Generates a small rotation with the axis vector selected uniformly - from the unit sphere and an angle selected from a uniform - distribution over [0, max_angle]. - - If the max_angle is not specified, the rotation should be selected - uniformly over all possible rotation angles. - - Args: - max_angle: The maximum angle of rotation (radians). - - Returns: - A unit quaternion in a numpy array. - - """ - - angle = np.random.uniform(low=0, high=max_angle) - axis = Vector3RandomUnit() - return QuaternionFromAxisAngle(axis, angle) - - -# def QuaternionRotatePoint(point, quat): -# """Performs a rotation by quaternion. - -# Rotate the point by the quaternion using quaternion multiplication, -# (q * p * q^-1), without constructing the rotation matrix. - -# Args: -# point: The point to be rotated. -# quat: The rotation represented as a quaternion [x, y, z, w]. - -# Returns: -# A 3D vector in a numpy array. -# """ - -# q_point = np.array([point[0], point[1], point[2], 0.0]) -# quat_inverse = transformations.quaternion_inverse(quat) -# q_point_rotated = transformations.quaternion_multiply( -# transformations.quaternion_multiply(quat, q_point), quat_inverse) -# return q_point_rotated[:3] - - -def IsRotationMatrix(m): - """Returns true if the 3x3 submatrix represents a rotation. - - Args: - m: A transformation matrix. - - Raises: - ValueError: If input is not a matrix of size at least 3x3. - - Returns: - True if the 3x3 submatrix is a rotation (orthogonal). - """ - if len(m.shape) != 2 or m.shape[0] < 3 or m.shape[1] < 3: - raise ValueError('Matrix should be 3x3 or 4x4: %s\n %s' % (m.shape, m)) - rot = m[:3, :3] - eye = np.matmul(rot, np.transpose(rot)) - return np.isclose(eye, np.identity(3), atol=1e-4).all() - -# def ZAxisAlignedRobotPoseTool(robot_pose_tool): -# """Returns the current gripper pose rotated for alignment with the z-axis. - -# Args: -# robot_pose_tool: a pose3d.Pose3d() instance. - -# Returns: -# An instance of pose.Transform representing the current gripper pose -# rotated for alignment with the z-axis. -# """ -# # Align the current pose to the z-axis. -# robot_pose_tool.quaternion = transformations.quaternion_multiply( -# RotationBetween( -# robot_pose_tool.matrix4x4[0:3, 0:3].dot(np.array([0, 0, 1])), -# np.array([0.0, 0.0, -1.0])), robot_pose_tool.quaternion) -# return robot_pose_tool - -# def RotationBetween(a_translation_b, a_translation_c): -# """Computes the rotation from one vector to another. - -# The computed rotation has the property that: - -# a_translation_c = a_rotation_b_to_c * a_translation_b - -# Args: -# a_translation_b: vec3, vector to rotate from -# a_translation_c: vec3, vector to rotate to - -# Returns: -# a_rotation_b_to_c: new Orientation -# """ -# rotation = rotation3.Rotation3.rotation_between( -# a_translation_b, a_translation_c, err_msg='RotationBetween') -# return rotation.quaternion.xyzw diff --git a/rsl_rl/build/lib/rsl_rl/utils/utils.py b/rsl_rl/build/lib/rsl_rl/utils/utils.py deleted file mode 100644 index 3605622..0000000 --- a/rsl_rl/build/lib/rsl_rl/utils/utils.py +++ /dev/null @@ -1,360 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -from __future__ import annotations - -import git -import importlib -import os -import pathlib -import torch -import warnings -from tensordict import TensorDict -from typing import Callable -import numpy as np -class RunningMeanStd: - def __init__(self, epsilon: float = 1e-4, shape: Tuple[int, ...] = ()): - """ - Calculates the running mean and std of a data stream - https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Parallel_algorithm - :param epsilon: helps with arithmetic issues - :param shape: the shape of the data stream's output - """ - self.mean = np.zeros(shape, np.float64) - self.var = np.ones(shape, np.float64) - self.count = epsilon - - def update(self, arr: np.ndarray) -> None: - batch_mean = np.mean(arr, axis=0) - batch_var = np.var(arr, axis=0) - batch_count = arr.shape[0] - self.update_from_moments(batch_mean, batch_var, batch_count) - - def update_from_moments(self, batch_mean: np.ndarray, batch_var: np.ndarray, batch_count: int) -> None: - delta = batch_mean - self.mean - tot_count = self.count + batch_count - - new_mean = self.mean + delta * batch_count / tot_count - m_a = self.var * self.count - m_b = batch_var * batch_count - m_2 = m_a + m_b + np.square(delta) * self.count * batch_count / (self.count + batch_count) - new_var = m_2 / (self.count + batch_count) - - new_count = batch_count + self.count - - self.mean = new_mean - self.var = new_var - self.count = new_count - - -class Normalizer(RunningMeanStd): - def __init__(self, input_dim, epsilon=1e-4, clip_obs=10.0): - super().__init__(shape=input_dim) - self.epsilon = epsilon - self.clip_obs = clip_obs - - def normalize(self, input): - return np.clip((input - self.mean) / np.sqrt(self.var + self.epsilon), -self.clip_obs, self.clip_obs) - - def normalize_torch(self, input, device): - mean_torch = torch.tensor(self.mean, device=device, dtype=torch.float32) - std_torch = torch.sqrt(torch.tensor(self.var + self.epsilon, device=device, dtype=torch.float32)) - return torch.clamp((input - mean_torch) / std_torch, -self.clip_obs, self.clip_obs) - - def update_normalizer(self, rollouts, expert_loader): - policy_data_generator = rollouts.feed_forward_generator_amp(None, mini_batch_size=expert_loader.batch_size) - expert_data_generator = expert_loader.dataset.feed_forward_generator_amp(expert_loader.batch_size) - - for expert_batch, policy_batch in zip(expert_data_generator, policy_data_generator): - self.update(torch.vstack(tuple(policy_batch) + tuple(expert_batch)).cpu().numpy()) - - -def resolve_nn_activation(act_name: str) -> torch.nn.Module: - """Resolves the activation function from the name. - - Args: - act_name: The name of the activation function. - - Returns: - The activation function. - - Raises: - ValueError: If the activation function is not found. - """ - act_dict = { - "elu": torch.nn.ELU(), - "selu": torch.nn.SELU(), - "relu": torch.nn.ReLU(), - "crelu": torch.nn.CELU(), - "lrelu": torch.nn.LeakyReLU(), - "tanh": torch.nn.Tanh(), - "sigmoid": torch.nn.Sigmoid(), - "softplus": torch.nn.Softplus(), - "gelu": torch.nn.GELU(), - "swish": torch.nn.SiLU(), - "mish": torch.nn.Mish(), - "identity": torch.nn.Identity(), - } - - act_name = act_name.lower() - if act_name in act_dict: - return act_dict[act_name] - else: - raise ValueError(f"Invalid activation function '{act_name}'. Valid activations are: {list(act_dict.keys())}") - - -def resolve_optimizer(optimizer_name: str) -> torch.optim.Optimizer: - """Resolves the optimizer from the name. - - Args: - optimizer_name: The name of the optimizer. - - Returns: - The optimizer. - - Raises: - ValueError: If the optimizer is not found. - """ - optimizer_dict = { - "adam": torch.optim.Adam, - "adamw": torch.optim.AdamW, - "sgd": torch.optim.SGD, - "rmsprop": torch.optim.RMSprop, - } - - optimizer_name = optimizer_name.lower() - if optimizer_name in optimizer_dict: - return optimizer_dict[optimizer_name] - else: - raise ValueError(f"Invalid optimizer '{optimizer_name}'. Valid optimizers are: {list(optimizer_dict.keys())}") - - -def split_and_pad_trajectories( - tensor: torch.Tensor | TensorDict, dones: torch.Tensor -) -> tuple[torch.Tensor | TensorDict, torch.Tensor]: - """Splits trajectories at done indices. Then concatenates them and pads with zeros up to the length of the longest - trajectory. Returns masks corresponding to valid parts of the trajectories. - - Example: - Input: [[a1, a2, a3, a4 | a5, a6], - [b1, b2 | b3, b4, b5 | b6]] - - Output:[[a1, a2, a3, a4], | [[True, True, True, True], - [a5, a6, 0, 0], | [True, True, False, False], - [b1, b2, 0, 0], | [True, True, False, False], - [b3, b4, b5, 0], | [True, True, True, False], - [b6, 0, 0, 0]] | [True, False, False, False]] - - Assumes that the input has the following order of dimensions: [time, number of envs, additional dimensions] - """ - - dones = dones.clone() - dones[-1] = 1 - # Permute the buffers to have order (num_envs, num_transitions_per_env, ...), for correct reshaping - flat_dones = dones.transpose(1, 0).reshape(-1, 1) - # Get length of trajectory by counting the number of successive not done elements - done_indices = torch.cat((flat_dones.new_tensor([-1], dtype=torch.int64), flat_dones.nonzero()[:, 0])) - trajectory_lengths = done_indices[1:] - done_indices[:-1] - trajectory_lengths_list = trajectory_lengths.tolist() - # Extract the individual trajectories - if isinstance(tensor, TensorDict): - padded_trajectories = {} - for k, v in tensor.items(): - # split the tensor into trajectories - trajectories = torch.split(v.transpose(1, 0).flatten(0, 1), trajectory_lengths_list) - # add at least one full length trajectory - trajectories = trajectories + (torch.zeros(v.shape[0], *v.shape[2:], device=v.device),) - # pad the trajectories to the length of the longest trajectory - padded_trajectories[k] = torch.nn.utils.rnn.pad_sequence(trajectories) - # remove the added tensor - padded_trajectories[k] = padded_trajectories[k][:, :-1] - padded_trajectories = TensorDict( - padded_trajectories, batch_size=[tensor.batch_size[0], len(trajectory_lengths_list)] - ) - else: - # split the tensor into trajectories - trajectories = torch.split(tensor.transpose(1, 0).flatten(0, 1), trajectory_lengths_list) - # add at least one full length trajectory - trajectories = trajectories + (torch.zeros(tensor.shape[0], *tensor.shape[2:], device=tensor.device),) - # pad the trajectories to the length of the longest trajectory - padded_trajectories = torch.nn.utils.rnn.pad_sequence(trajectories) - # remove the added tensor - padded_trajectories = padded_trajectories[:, :-1] - # create masks for the valid parts of the trajectories - trajectory_masks = trajectory_lengths > torch.arange(0, tensor.shape[0], device=tensor.device).unsqueeze(1) - return padded_trajectories, trajectory_masks - - -def unpad_trajectories(trajectories, masks): - """Does the inverse operation of split_and_pad_trajectories()""" - # Need to transpose before and after the masking to have proper reshaping - return ( - trajectories.transpose(1, 0)[masks.transpose(1, 0)] - .view(-1, trajectories.shape[0], trajectories.shape[-1]) - .transpose(1, 0) - ) - - -def store_code_state(logdir, repositories) -> list: - git_log_dir = os.path.join(logdir, "git") - os.makedirs(git_log_dir, exist_ok=True) - file_paths = [] - for repository_file_path in repositories: - try: - repo = git.Repo(repository_file_path, search_parent_directories=True) - t = repo.head.commit.tree - except Exception: - print(f"Could not find git repository in {repository_file_path}. Skipping.") - # skip if not a git repository - continue - # get the name of the repository - repo_name = pathlib.Path(repo.working_dir).name - diff_file_name = os.path.join(git_log_dir, f"{repo_name}.diff") - # check if the diff file already exists - if os.path.isfile(diff_file_name): - continue - # write the diff file - print(f"Storing git diff for '{repo_name}' in: {diff_file_name}") - with open(diff_file_name, "x", encoding="utf-8") as f: - content = f"--- git status ---\n{repo.git.status()} \n\n\n--- git diff ---\n{repo.git.diff(t)}" - f.write(content) - # add the file path to the list of files to be uploaded - file_paths.append(diff_file_name) - return file_paths - - -def string_to_callable(name: str) -> Callable: - """Resolves the module and function names to return the function. - - Args: - name: The function name. The format should be 'module:attribute_name'. - - Raises: - ValueError: When the resolved attribute is not a function. - ValueError: When unable to resolve the attribute. - - Returns: - The function loaded from the module. - """ - try: - mod_name, attr_name = name.split(":") - mod = importlib.import_module(mod_name) - callable_object = getattr(mod, attr_name) - # check if attribute is callable - if callable(callable_object): - return callable_object - else: - raise ValueError(f"The imported object is not callable: '{name}'") - except AttributeError as e: - msg = ( - "We could not interpret the entry as a callable object. The format of input should be" - f" 'module:attribute_name'\nWhile processing input '{name}', received the error:\n {e}." - ) - raise ValueError(msg) - - -def resolve_obs_groups( - obs: TensorDict, obs_groups: dict[str, list[str]], default_sets: list[str] -) -> dict[str, list[str]]: - """Validates the observation configuration and defaults missing observation sets. - - The input is an observation dictionary `obs` containing observation groups and a configuration dictionary - `obs_groups` where the keys are the observation sets and the values are lists of observation groups. - - The configuration dictionary could for example look like: - { - "policy": ["group_1", "group_2"], - "critic": ["group_1", "group_3"] - } - - This means that the 'policy' observation set will contain the observations "group_1" and "group_2" and the - 'critic' observation set will contain the observations "group_1" and "group_3". This function will check that all - the observations in the 'policy' and 'critic' observation sets are present in the observation dictionary from the - environment. - - Additionally, if one of the `default_sets`, e.g. "critic", is not present in the configuration dictionary, - this function will: - - 1. Check if a group with the same name exists in the observations and assign this group to the observation set. - 2. If 1. fails, it will assign the observations from the 'policy' observation set to the default observation set. - - Args: - obs: Observations from the environment in the form of a dictionary. - obs_groups: Observation sets configuration. - default_sets: Reserved observation set names used by the algorithm (besides 'policy'). - If not provided in 'obs_groups', a default behavior gets triggered. - - Returns: - The resolved observation groups. - - Raises: - ValueError: If any observation set is an empty list. - ValueError: If any observation set contains an observation term that is not present in the observations. - """ - # check if policy observation set exists - if "policy" not in obs_groups.keys(): - if "policy" in obs: - obs_groups["policy"] = ["policy"] - warnings.warn( - "The observation configuration dictionary 'obs_groups' must contain the 'policy' key." - " As an observation group with the name 'policy' was found, this is assumed to be the observation set." - " Consider adding the 'policy' key to the 'obs_groups' dictionary for clarity." - " This behavior will be removed in a future version." - ) - else: - raise ValueError( - "The observation configuration dictionary 'obs_groups' must contain the 'policy' key." - f" Found keys: {list(obs_groups.keys())}" - ) - - # check all observation sets for valid observation groups - for set_name, groups in obs_groups.items(): - # check if the list is empty - if len(groups) == 0: - msg = f"The '{set_name}' key in the 'obs_groups' dictionary can not be an empty list." - if set_name in default_sets: - if set_name not in obs: - msg += " Consider removing the key to default to the observations used for the 'policy' set." - else: - msg += ( - f" Consider removing the key to default to the observation '{set_name}' from the environment." - ) - raise ValueError(msg) - # check groups exist inside the observations from the environment - for group in groups: - if group not in obs: - raise ValueError( - f"Observation '{group}' in observation set '{set_name}' not found in the observations from the" - f" environment. Available observations from the environment: {list(obs.keys())}" - ) - - # fill missing observation sets - for default_set_name in default_sets: - if default_set_name not in obs_groups.keys(): - if default_set_name in obs: - obs_groups[default_set_name] = [default_set_name] - warnings.warn( - f"The observation configuration dictionary 'obs_groups' must contain the '{default_set_name}' key." - f" As an observation group with the name '{default_set_name}' was found, this is assumed to be the" - f" observation set. Consider adding the '{default_set_name}' key to the 'obs_groups' dictionary for" - " clarity. This behavior will be removed in a future version." - ) - else: - obs_groups[default_set_name] = obs_groups["policy"].copy() - warnings.warn( - f"The observation configuration dictionary 'obs_groups' must contain the '{default_set_name}' key." - f" As the configuration for '{default_set_name}' is missing, the observations from the 'policy' set" - f" are used. Consider adding the '{default_set_name}' key to the 'obs_groups' dictionary for" - " clarity. This behavior will be removed in a future version." - ) - - # print the final parsed observation sets - print("-" * 80) - print("Resolved observation sets: ") - for set_name, groups in obs_groups.items(): - print("\t", set_name, ": ", groups) - print("-" * 80) - - return obs_groups diff --git a/rsl_rl/build/lib/rsl_rl/utils/wandb_utils.py b/rsl_rl/build/lib/rsl_rl/utils/wandb_utils.py deleted file mode 100644 index 243e82d..0000000 --- a/rsl_rl/build/lib/rsl_rl/utils/wandb_utils.py +++ /dev/null @@ -1,87 +0,0 @@ -# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause - -from __future__ import annotations - -import os -from dataclasses import asdict -from torch.utils.tensorboard import SummaryWriter - -try: - import wandb -except ModuleNotFoundError: - raise ModuleNotFoundError("Wandb is required to log to Weights and Biases.") - - -class WandbSummaryWriter(SummaryWriter): - """Summary writer for Weights and Biases.""" - - def __init__(self, log_dir: str, flush_secs: int, cfg): - super().__init__(log_dir, flush_secs) - - # Get the run name - run_name = os.path.split(log_dir)[-1] - - try: - project = cfg["wandb_project"] - except KeyError: - raise KeyError("Please specify wandb_project in the runner config, e.g. legged_gym.") - - try: - entity = os.environ["WANDB_USERNAME"] - except KeyError: - entity = None - - # Initialize wandb - wandb.init(project=project, entity=entity, name=run_name) - - # Add log directory to wandb - wandb.config.update({"log_dir": log_dir}) - - self.name_map = { - "Train/mean_reward/time": "Train/mean_reward_time", - "Train/mean_episode_length/time": "Train/mean_episode_length_time", - } - - def store_config(self, env_cfg, runner_cfg, alg_cfg, policy_cfg): - wandb.config.update({"runner_cfg": runner_cfg}) - wandb.config.update({"policy_cfg": policy_cfg}) - wandb.config.update({"alg_cfg": alg_cfg}) - try: - wandb.config.update({"env_cfg": env_cfg.to_dict()}) - except Exception: - wandb.config.update({"env_cfg": asdict(env_cfg)}) - - def add_scalar(self, tag, scalar_value, global_step=None, walltime=None, new_style=False): - super().add_scalar( - tag, - scalar_value, - global_step=global_step, - walltime=walltime, - new_style=new_style, - ) - wandb.log({self._map_path(tag): scalar_value}, step=global_step) - - def stop(self): - wandb.finish() - - def log_config(self, env_cfg, runner_cfg, alg_cfg, policy_cfg): - self.store_config(env_cfg, runner_cfg, alg_cfg, policy_cfg) - - def save_model(self, model_path, iter): - wandb.save(model_path, base_path=os.path.dirname(model_path)) - - def save_file(self, path, iter=None): - wandb.save(path, base_path=os.path.dirname(path)) - - """ - Private methods. - """ - - def _map_path(self, path): - if path in self.name_map: - return self.name_map[path] - else: - return path diff --git a/rsl_rl/env/__pycache__/__init__.cpython-312.pyc b/rsl_rl/env/__pycache__/__init__.cpython-312.pyc deleted file mode 100644 index 4ed9b78..0000000 Binary files a/rsl_rl/env/__pycache__/__init__.cpython-312.pyc and /dev/null differ diff --git a/rsl_rl/env/__pycache__/vec_env.cpython-312.pyc b/rsl_rl/env/__pycache__/vec_env.cpython-312.pyc deleted file mode 100644 index 1ded7ac..0000000 Binary files a/rsl_rl/env/__pycache__/vec_env.cpython-312.pyc and /dev/null differ diff --git a/rsl_rl/modules/__pycache__/__init__.cpython-312.pyc b/rsl_rl/modules/__pycache__/__init__.cpython-312.pyc deleted file mode 100644 index 0571d1c..0000000 Binary files a/rsl_rl/modules/__pycache__/__init__.cpython-312.pyc and /dev/null differ diff --git a/rsl_rl/modules/__pycache__/actor_critic.cpython-312.pyc b/rsl_rl/modules/__pycache__/actor_critic.cpython-312.pyc deleted file mode 100644 index a75e11b..0000000 Binary files a/rsl_rl/modules/__pycache__/actor_critic.cpython-312.pyc and /dev/null differ diff --git a/rsl_rl/modules/__pycache__/actor_critic_recurrent.cpython-312.pyc b/rsl_rl/modules/__pycache__/actor_critic_recurrent.cpython-312.pyc deleted file mode 100644 index da53671..0000000 Binary files a/rsl_rl/modules/__pycache__/actor_critic_recurrent.cpython-312.pyc and /dev/null differ diff --git a/rsl_rl/modules/__pycache__/discriminator_multi.cpython-312.pyc b/rsl_rl/modules/__pycache__/discriminator_multi.cpython-312.pyc deleted file mode 100644 index c25a93e..0000000 Binary files a/rsl_rl/modules/__pycache__/discriminator_multi.cpython-312.pyc and /dev/null differ diff --git a/rsl_rl/modules/__pycache__/rnd.cpython-312.pyc b/rsl_rl/modules/__pycache__/rnd.cpython-312.pyc deleted file mode 100644 index 96eef25..0000000 Binary files a/rsl_rl/modules/__pycache__/rnd.cpython-312.pyc and /dev/null differ diff --git a/rsl_rl/modules/__pycache__/student_teacher.cpython-312.pyc b/rsl_rl/modules/__pycache__/student_teacher.cpython-312.pyc deleted file mode 100644 index 48ba5bf..0000000 Binary files a/rsl_rl/modules/__pycache__/student_teacher.cpython-312.pyc and /dev/null differ diff --git a/rsl_rl/modules/__pycache__/student_teacher_recurrent.cpython-312.pyc b/rsl_rl/modules/__pycache__/student_teacher_recurrent.cpython-312.pyc deleted file mode 100644 index 4f74bc3..0000000 Binary files a/rsl_rl/modules/__pycache__/student_teacher_recurrent.cpython-312.pyc and /dev/null differ diff --git a/rsl_rl/modules/__pycache__/symmetry.cpython-312.pyc b/rsl_rl/modules/__pycache__/symmetry.cpython-312.pyc deleted file mode 100644 index 6e226f1..0000000 Binary files a/rsl_rl/modules/__pycache__/symmetry.cpython-312.pyc and /dev/null differ diff --git a/rsl_rl/networks/__pycache__/__init__.cpython-312.pyc b/rsl_rl/networks/__pycache__/__init__.cpython-312.pyc deleted file mode 100644 index d21790c..0000000 Binary files a/rsl_rl/networks/__pycache__/__init__.cpython-312.pyc and /dev/null differ diff --git a/rsl_rl/networks/__pycache__/memory.cpython-312.pyc b/rsl_rl/networks/__pycache__/memory.cpython-312.pyc deleted file mode 100644 index 1ea714a..0000000 Binary files a/rsl_rl/networks/__pycache__/memory.cpython-312.pyc and /dev/null differ diff --git a/rsl_rl/networks/__pycache__/mlp.cpython-312.pyc b/rsl_rl/networks/__pycache__/mlp.cpython-312.pyc deleted file mode 100644 index 941e41d..0000000 Binary files a/rsl_rl/networks/__pycache__/mlp.cpython-312.pyc and /dev/null differ diff --git a/rsl_rl/networks/__pycache__/normalization.cpython-312.pyc b/rsl_rl/networks/__pycache__/normalization.cpython-312.pyc deleted file mode 100644 index 5c9811b..0000000 Binary files a/rsl_rl/networks/__pycache__/normalization.cpython-312.pyc and /dev/null differ diff --git a/rsl_rl/runners/__pycache__/__init__.cpython-312.pyc b/rsl_rl/runners/__pycache__/__init__.cpython-312.pyc deleted file mode 100644 index f3c0656..0000000 Binary files a/rsl_rl/runners/__pycache__/__init__.cpython-312.pyc and /dev/null differ diff --git a/rsl_rl/runners/__pycache__/amp_on_policy_runner.cpython-312.pyc b/rsl_rl/runners/__pycache__/amp_on_policy_runner.cpython-312.pyc deleted file mode 100644 index 7aa2f80..0000000 Binary files a/rsl_rl/runners/__pycache__/amp_on_policy_runner.cpython-312.pyc and /dev/null differ diff --git a/rsl_rl/runners/__pycache__/distillation_runner.cpython-312.pyc b/rsl_rl/runners/__pycache__/distillation_runner.cpython-312.pyc deleted file mode 100644 index bd26702..0000000 Binary files a/rsl_rl/runners/__pycache__/distillation_runner.cpython-312.pyc and /dev/null differ diff --git a/rsl_rl/runners/__pycache__/on_policy_runner.cpython-312.pyc b/rsl_rl/runners/__pycache__/on_policy_runner.cpython-312.pyc deleted file mode 100644 index dada129..0000000 Binary files a/rsl_rl/runners/__pycache__/on_policy_runner.cpython-312.pyc and /dev/null differ diff --git a/rsl_rl/storage/__pycache__/__init__.cpython-312.pyc b/rsl_rl/storage/__pycache__/__init__.cpython-312.pyc deleted file mode 100644 index a8bbe7c..0000000 Binary files a/rsl_rl/storage/__pycache__/__init__.cpython-312.pyc and /dev/null differ diff --git a/rsl_rl/storage/__pycache__/replay_buffer_multi.cpython-312.pyc b/rsl_rl/storage/__pycache__/replay_buffer_multi.cpython-312.pyc deleted file mode 100644 index 5451aaf..0000000 Binary files a/rsl_rl/storage/__pycache__/replay_buffer_multi.cpython-312.pyc and /dev/null differ diff --git a/rsl_rl/storage/__pycache__/rollout_storage.cpython-312.pyc b/rsl_rl/storage/__pycache__/rollout_storage.cpython-312.pyc deleted file mode 100644 index 7274cec..0000000 Binary files a/rsl_rl/storage/__pycache__/rollout_storage.cpython-312.pyc and /dev/null differ diff --git a/rsl_rl/utils/__pycache__/__init__.cpython-312.pyc b/rsl_rl/utils/__pycache__/__init__.cpython-312.pyc deleted file mode 100644 index e483828..0000000 Binary files a/rsl_rl/utils/__pycache__/__init__.cpython-312.pyc and /dev/null differ diff --git a/rsl_rl/utils/__pycache__/motion_loader_g1.cpython-312.pyc b/rsl_rl/utils/__pycache__/motion_loader_g1.cpython-312.pyc deleted file mode 100644 index fd5328a..0000000 Binary files a/rsl_rl/utils/__pycache__/motion_loader_g1.cpython-312.pyc and /dev/null differ diff --git a/rsl_rl/utils/__pycache__/motion_util.cpython-312.pyc b/rsl_rl/utils/__pycache__/motion_util.cpython-312.pyc deleted file mode 100644 index 078ba2c..0000000 Binary files a/rsl_rl/utils/__pycache__/motion_util.cpython-312.pyc and /dev/null differ diff --git a/rsl_rl/utils/__pycache__/pose3d.cpython-312.pyc b/rsl_rl/utils/__pycache__/pose3d.cpython-312.pyc deleted file mode 100644 index 12646b2..0000000 Binary files a/rsl_rl/utils/__pycache__/pose3d.cpython-312.pyc and /dev/null differ diff --git a/rsl_rl/utils/__pycache__/utils.cpython-312.pyc b/rsl_rl/utils/__pycache__/utils.cpython-312.pyc deleted file mode 100644 index be27767..0000000 Binary files a/rsl_rl/utils/__pycache__/utils.cpython-312.pyc and /dev/null differ diff --git a/src/.DS_Store b/src/.DS_Store new file mode 100644 index 0000000..a611596 Binary files /dev/null and b/src/.DS_Store differ diff --git a/src/mjlab_husky/.DS_Store b/src/mjlab_husky/.DS_Store new file mode 100644 index 0000000..00f5a6c Binary files /dev/null and b/src/mjlab_husky/.DS_Store differ diff --git a/src/mjlab_husky/__pycache__/__init__.cpython-312.pyc b/src/mjlab_husky/__pycache__/__init__.cpython-312.pyc deleted file mode 100644 index 9717d47..0000000 Binary files a/src/mjlab_husky/__pycache__/__init__.cpython-312.pyc and /dev/null differ diff --git a/src/mjlab_husky/__pycache__/lerobot_numpy.cpython-312.pyc b/src/mjlab_husky/__pycache__/lerobot_numpy.cpython-312.pyc deleted file mode 100644 index 6c00803..0000000 Binary files a/src/mjlab_husky/__pycache__/lerobot_numpy.cpython-312.pyc and /dev/null differ diff --git a/src/mjlab_husky/__pycache__/mujoco_gl.cpython-312.pyc b/src/mjlab_husky/__pycache__/mujoco_gl.cpython-312.pyc deleted file mode 100644 index 0771451..0000000 Binary files a/src/mjlab_husky/__pycache__/mujoco_gl.cpython-312.pyc and /dev/null differ diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/__pycache__/__init__.cpython-312.pyc b/src/mjlab_husky/asset_zoo/robots/skateboard/__pycache__/__init__.cpython-312.pyc deleted file mode 100644 index 101653e..0000000 Binary files a/src/mjlab_husky/asset_zoo/robots/skateboard/__pycache__/__init__.cpython-312.pyc and /dev/null differ diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/__pycache__/g1_skater_constants.cpython-312.pyc b/src/mjlab_husky/asset_zoo/robots/skateboard/__pycache__/g1_skater_constants.cpython-312.pyc deleted file mode 100644 index 75bac0a..0000000 Binary files a/src/mjlab_husky/asset_zoo/robots/skateboard/__pycache__/g1_skater_constants.cpython-312.pyc and /dev/null differ diff --git a/src/mjlab_husky/envs/__pycache__/__init__.cpython-312.pyc b/src/mjlab_husky/envs/__pycache__/__init__.cpython-312.pyc deleted file mode 100644 index 1324f5b..0000000 Binary files a/src/mjlab_husky/envs/__pycache__/__init__.cpython-312.pyc and /dev/null differ diff --git a/src/mjlab_husky/envs/__pycache__/g1_skate_rl_env.cpython-312.pyc b/src/mjlab_husky/envs/__pycache__/g1_skate_rl_env.cpython-312.pyc deleted file mode 100644 index 6ffe39d..0000000 Binary files a/src/mjlab_husky/envs/__pycache__/g1_skate_rl_env.cpython-312.pyc and /dev/null differ diff --git a/src/mjlab_husky/lerobot_numpy.py b/src/mjlab_husky/lerobot_numpy.py index 361e8eb..d43b87e 100644 --- a/src/mjlab_husky/lerobot_numpy.py +++ b/src/mjlab_husky/lerobot_numpy.py @@ -8,6 +8,20 @@ import numpy as np import torch +def import_lerobot_dataset_class(): + """Return ``LeRobotDataset`` or raise with install hints (optional ``lerobot`` extra).""" + try: + from lerobot.datasets.lerobot_dataset import LeRobotDataset # type: ignore + except Exception as e: # pragma: no cover + raise RuntimeError( + "未安装 `lerobot`(LeRobot v3 写数据需要)。\n" + "在项目目录执行:\n" + " uv pip install lerobot\n" + f"原始导入错误: {e}" + ) from e + return LeRobotDataset + + def batched_vector_to_numpy1d(x: Any, env_i: int = 0) -> np.ndarray: """Take env index `env_i`, flatten to 1D float32 (LeRobot parquet-friendly).""" if x is None: diff --git a/src/mjlab_husky/mujoco_gl.py b/src/mjlab_husky/mujoco_gl.py index 41b958c..b682dc1 100644 --- a/src/mjlab_husky/mujoco_gl.py +++ b/src/mjlab_husky/mujoco_gl.py @@ -7,14 +7,24 @@ import sys def ensure_mujoco_headless_gl() -> None: - """If no display server, default to EGL so ``mujoco.Renderer`` works on headless GPU hosts. + """If no display server, set GL backends *before* ``import mujoco`` / ``import OpenGL``. - Override explicitly with ``MUJOCO_GL`` (e.g. ``osmesa`` on CPU-only machines). + On headless Linux, **PyOpenGL** would still match ``sys.platform == "linux"`` to the + **GLX** plugin, so ``PLATFORM.GL`` stays ``None`` and you get errors like + ``'NoneType' has no attribute 'glGetError'`` even when ``MUJOCO_GL=osmesa``. + + We set ``PYOPENGL_PLATFORM`` to match ``MUJOCO_GL`` (``osmesa`` or ``egl``). """ if sys.platform != "linux": return if os.environ.get("DISPLAY") or os.environ.get("WAYLAND_DISPLAY"): return - if os.environ.get("MUJOCO_GL"): - return - os.environ["MUJOCO_GL"] = "egl" + + if not os.environ.get("MUJOCO_GL"): + os.environ.setdefault("MUJOCO_GL", "osmesa") + + gl = (os.environ.get("MUJOCO_GL") or "osmesa").lower() + if gl == "osmesa": + os.environ.setdefault("PYOPENGL_PLATFORM", "osmesa") + elif gl == "egl": + os.environ.setdefault("PYOPENGL_PLATFORM", "egl") diff --git a/src/mjlab_husky/rl/__pycache__/__init__.cpython-312.pyc b/src/mjlab_husky/rl/__pycache__/__init__.cpython-312.pyc deleted file mode 100644 index e44b5e9..0000000 Binary files a/src/mjlab_husky/rl/__pycache__/__init__.cpython-312.pyc and /dev/null differ diff --git a/src/mjlab_husky/rl/__pycache__/config.cpython-312.pyc b/src/mjlab_husky/rl/__pycache__/config.cpython-312.pyc deleted file mode 100644 index f556d57..0000000 Binary files a/src/mjlab_husky/rl/__pycache__/config.cpython-312.pyc and /dev/null differ diff --git a/src/mjlab_husky/rl/__pycache__/vecenv_wrapper.cpython-312.pyc b/src/mjlab_husky/rl/__pycache__/vecenv_wrapper.cpython-312.pyc deleted file mode 100644 index 76c1752..0000000 Binary files a/src/mjlab_husky/rl/__pycache__/vecenv_wrapper.cpython-312.pyc and /dev/null differ diff --git a/src/mjlab_husky/scripts/__pycache__/play.cpython-312.pyc b/src/mjlab_husky/scripts/__pycache__/play.cpython-312.pyc deleted file mode 100644 index 66db666..0000000 Binary files a/src/mjlab_husky/scripts/__pycache__/play.cpython-312.pyc and /dev/null differ diff --git a/src/mjlab_husky/scripts/export_lerobot_qpos.py b/src/mjlab_husky/scripts/export_lerobot_qpos.py index 7cbe208..21ed84c 100644 --- a/src/mjlab_husky/scripts/export_lerobot_qpos.py +++ b/src/mjlab_husky/scripts/export_lerobot_qpos.py @@ -14,7 +14,7 @@ import numpy as np import torch import tyro -from mjlab_husky.lerobot_numpy import batched_vector_to_numpy1d +from mjlab_husky.lerobot_numpy import batched_vector_to_numpy1d, import_lerobot_dataset_class from mjlab_husky.mujoco_gl import ensure_mujoco_headless_gl ensure_mujoco_headless_gl() @@ -25,19 +25,6 @@ from mjlab_husky.rl import RslRlVecEnvWrapper from mjlab_husky.tasks.registry import load_env_cfg, load_rl_cfg, load_runner_cls -def _import_lerobot(): - try: - from lerobot.datasets.lerobot_dataset import LeRobotDataset # type: ignore - except Exception as e: # pragma: no cover - raise RuntimeError( - "未安装 `lerobot`(LeRobot v3 写数据需要)。\n" - "请先在项目目录执行:\n" - " uv pip install lerobot\n" - f"原始导入错误: {e}" - ) from e - return LeRobotDataset - - @dataclass(frozen=True) class ExportConfig: task_id: str @@ -56,7 +43,7 @@ class ExportConfig: def main(cfg: ExportConfig) -> None: configure_torch_backends() - LeRobotDataset = _import_lerobot() + LeRobotDataset = import_lerobot_dataset_class() device = cfg.device or ("cuda:0" if torch.cuda.is_available() else "cpu") env_cfg = load_env_cfg(cfg.task_id, play=True) diff --git a/src/mjlab_husky/scripts/play.py b/src/mjlab_husky/scripts/play.py index c98272c..cad1ae7 100644 --- a/src/mjlab_husky/scripts/play.py +++ b/src/mjlab_husky/scripts/play.py @@ -23,7 +23,9 @@ from mjlab.utils.os import get_wandb_checkpoint_path from mjlab.utils.torch import configure_torch_backends from mjlab.utils.wrappers import VideoRecorder from mjlab.viewer import NativeMujocoViewer, ViserPlayViewer +from mjlab_husky.viewer.rerun_native_play_viewer import RerunNativePlayViewer from mjlab_husky.viewer.rerun_play_viewer import RerunPlayViewer, RerunPlayViewerCfg +from mjlab_husky.viewer.rerun_viser_play_viewer import RerunViserPlayViewer @dataclass(frozen=True) @@ -40,9 +42,11 @@ class PlayConfig: video_height: int = 1080 video_width: int = 1920 camera: int | str | None = None - viewer: Literal["auto", "native", "viser", "rerun"] = "auto" + viewer: Literal["auto", "native", "viser", "rerun", "rerun_native", "rerun_viser"] = "auto" rerun_web_port: int = 8080 rerun_grpc_port: int | None = None + # Viser 独立端口(与 rerun_web_port 分离,便于 RoboHub 双 iframe / 避免默认都与 8080 冲突) + viser_port: int = 8081 # 浏览器直连云主机时填公网/LAN IP(与 --rerun-web-port 对应);否则用 SSH 转发保持默认 127.0.0.1。 rerun_connect_host: str | None = None # If False, open the printed `?url=...` link yourself. Use `--no-rerun-open-browser`. @@ -126,7 +130,7 @@ def run_play(task_id: str, cfg: PlayConfig): env_cfg.scene.num_envs = cfg.num_envs resolved_viewer = _resolve_play_viewer(cfg) - if resolved_viewer == "rerun" and not cfg.video: + if resolved_viewer in ("rerun", "rerun_native", "rerun_viser") and not cfg.video: env_cfg.viewer.height = cfg.rerun_viewer_height env_cfg.viewer.width = cfg.rerun_viewer_width else: @@ -136,7 +140,10 @@ def run_play(task_id: str, cfg: PlayConfig): # Rerun needs rgb_array so we can log the same offscreen camera as mjlab (tracking robot). render_mode = ( "rgb_array" - if ((TRAINED_MODE and cfg.video) or resolved_viewer == "rerun") + if ( + (TRAINED_MODE and cfg.video) + or resolved_viewer in ("rerun", "rerun_native", "rerun_viser") + ) else None ) if cfg.video and DUMMY_MODE: @@ -181,6 +188,22 @@ def run_play(task_id: str, cfg: PlayConfig): runner.load(str(resume_path), map_location=device) policy = runner.get_inference_policy(device=device) + if cfg.lerobot_record and resolved_viewer in ("rerun_native", "rerun_viser"): + raise RuntimeError( + "`--lerobot-record` 当前仅支持 `--viewer rerun`(独立 Rerun 循环)。" + "请去掉 `--lerobot-record` 或改用 `--viewer rerun`。" + ) + + if resolved_viewer == "rerun_viser": + if cfg.viser_port == cfg.rerun_web_port: + raise ValueError( + "`--viser-port` 与 `--rerun-web-port` 不能相同(Viser 与 Rerun Web 需各占一个端口)。" + ) + if cfg.rerun_grpc_port is not None and cfg.viser_port == cfg.rerun_grpc_port: + raise ValueError( + "`--viser-port` 与 `--rerun-grpc-port` 不能相同。" + ) + if resolved_viewer == "native": NativeMujocoViewer(env, policy).run() elif resolved_viewer == "viser": @@ -214,6 +237,65 @@ def run_play(task_id: str, cfg: PlayConfig): lerobot_save_on_env_reset=cfg.lerobot_save_on_env_reset, ), ).run() + elif resolved_viewer == "rerun_native": + RerunNativePlayViewer( + env, + policy, + rerun_cfg=RerunPlayViewerCfg( + web_port=cfg.rerun_web_port, + open_browser=cfg.rerun_open_browser, + grpc_port=cfg.rerun_grpc_port, + connect_host=cfg.rerun_connect_host, + log_camera=True, + camera_log_stride=max(1, cfg.rerun_camera_log_stride), + qpos_log_stride=max(1, cfg.rerun_qpos_log_stride), + camera_max_side=max(64, cfg.rerun_camera_max_side), + grpc_newest_first=cfg.rerun_grpc_newest_first, + log_qpos_joint_scalars=cfg.rerun_log_qpos_joint_scalars, + log_qpos_text=cfg.rerun_log_qpos_text, + qpos_text_decimals=max(0, min(8, cfg.rerun_qpos_text_decimals)), + log_qpos_tensor=cfg.rerun_log_qpos_tensor, + qpos_env_idx=max(0, cfg.rerun_qpos_env_idx), + lerobot_record=False, + lerobot_out_dir=cfg.lerobot_out_dir, + lerobot_dataset_name=cfg.lerobot_dataset_name, + lerobot_steps_per_episode=max(1, cfg.lerobot_steps_per_episode), + lerobot_task=cfg.lerobot_task, + lerobot_overwrite=cfg.lerobot_overwrite, + lerobot_record_stride=max(1, cfg.lerobot_record_stride), + lerobot_save_on_env_reset=cfg.lerobot_save_on_env_reset, + ), + ).run() + elif resolved_viewer == "rerun_viser": + RerunViserPlayViewer( + env, + policy, + rerun_cfg=RerunPlayViewerCfg( + web_port=cfg.rerun_web_port, + open_browser=cfg.rerun_open_browser, + grpc_port=cfg.rerun_grpc_port, + connect_host=cfg.rerun_connect_host, + log_camera=True, + camera_log_stride=max(1, cfg.rerun_camera_log_stride), + qpos_log_stride=max(1, cfg.rerun_qpos_log_stride), + camera_max_side=max(64, cfg.rerun_camera_max_side), + grpc_newest_first=cfg.rerun_grpc_newest_first, + log_qpos_joint_scalars=cfg.rerun_log_qpos_joint_scalars, + log_qpos_text=cfg.rerun_log_qpos_text, + qpos_text_decimals=max(0, min(8, cfg.rerun_qpos_text_decimals)), + log_qpos_tensor=cfg.rerun_log_qpos_tensor, + qpos_env_idx=max(0, cfg.rerun_qpos_env_idx), + lerobot_record=False, + lerobot_out_dir=cfg.lerobot_out_dir, + lerobot_dataset_name=cfg.lerobot_dataset_name, + lerobot_steps_per_episode=max(1, cfg.lerobot_steps_per_episode), + lerobot_task=cfg.lerobot_task, + lerobot_overwrite=cfg.lerobot_overwrite, + lerobot_record_stride=max(1, cfg.lerobot_record_stride), + lerobot_save_on_env_reset=cfg.lerobot_save_on_env_reset, + ), + viser_port=cfg.viser_port, + ).run() else: raise RuntimeError(f"Unsupported viewer backend: {resolved_viewer}") diff --git a/src/mjlab_husky/tasks/__pycache__/__init__.cpython-312.pyc b/src/mjlab_husky/tasks/__pycache__/__init__.cpython-312.pyc deleted file mode 100644 index 70eedb7..0000000 Binary files a/src/mjlab_husky/tasks/__pycache__/__init__.cpython-312.pyc and /dev/null differ diff --git a/src/mjlab_husky/tasks/__pycache__/registry.cpython-312.pyc b/src/mjlab_husky/tasks/__pycache__/registry.cpython-312.pyc deleted file mode 100644 index fcd5d88..0000000 Binary files a/src/mjlab_husky/tasks/__pycache__/registry.cpython-312.pyc and /dev/null differ diff --git a/src/mjlab_husky/tasks/skater/__pycache__/__init__.cpython-312.pyc b/src/mjlab_husky/tasks/skater/__pycache__/__init__.cpython-312.pyc deleted file mode 100644 index 67ba59a..0000000 Binary files a/src/mjlab_husky/tasks/skater/__pycache__/__init__.cpython-312.pyc and /dev/null differ diff --git a/src/mjlab_husky/tasks/skater/__pycache__/skater_env_cfg.cpython-312.pyc b/src/mjlab_husky/tasks/skater/__pycache__/skater_env_cfg.cpython-312.pyc deleted file mode 100644 index 43485df..0000000 Binary files a/src/mjlab_husky/tasks/skater/__pycache__/skater_env_cfg.cpython-312.pyc and /dev/null differ diff --git a/src/mjlab_husky/tasks/skater/config/__pycache__/__init__.cpython-312.pyc b/src/mjlab_husky/tasks/skater/config/__pycache__/__init__.cpython-312.pyc deleted file mode 100644 index 8548d2e..0000000 Binary files a/src/mjlab_husky/tasks/skater/config/__pycache__/__init__.cpython-312.pyc and /dev/null differ diff --git a/src/mjlab_husky/tasks/skater/config/g1/__pycache__/__init__.cpython-312.pyc b/src/mjlab_husky/tasks/skater/config/g1/__pycache__/__init__.cpython-312.pyc deleted file mode 100644 index f49fadd..0000000 Binary files a/src/mjlab_husky/tasks/skater/config/g1/__pycache__/__init__.cpython-312.pyc and /dev/null differ diff --git a/src/mjlab_husky/tasks/skater/config/g1/__pycache__/env_cfgs.cpython-312.pyc b/src/mjlab_husky/tasks/skater/config/g1/__pycache__/env_cfgs.cpython-312.pyc deleted file mode 100644 index 3c11465..0000000 Binary files a/src/mjlab_husky/tasks/skater/config/g1/__pycache__/env_cfgs.cpython-312.pyc and /dev/null differ diff --git a/src/mjlab_husky/tasks/skater/config/g1/__pycache__/rl_cfg.cpython-312.pyc b/src/mjlab_husky/tasks/skater/config/g1/__pycache__/rl_cfg.cpython-312.pyc deleted file mode 100644 index 8f09e0f..0000000 Binary files a/src/mjlab_husky/tasks/skater/config/g1/__pycache__/rl_cfg.cpython-312.pyc and /dev/null differ diff --git a/src/mjlab_husky/tasks/skater/mdp/__pycache__/__init__.cpython-312.pyc b/src/mjlab_husky/tasks/skater/mdp/__pycache__/__init__.cpython-312.pyc deleted file mode 100644 index 3fbbffc..0000000 Binary files a/src/mjlab_husky/tasks/skater/mdp/__pycache__/__init__.cpython-312.pyc and /dev/null differ diff --git a/src/mjlab_husky/tasks/skater/mdp/__pycache__/observations.cpython-312.pyc b/src/mjlab_husky/tasks/skater/mdp/__pycache__/observations.cpython-312.pyc deleted file mode 100644 index 626ce58..0000000 Binary files a/src/mjlab_husky/tasks/skater/mdp/__pycache__/observations.cpython-312.pyc and /dev/null differ diff --git a/src/mjlab_husky/tasks/skater/mdp/__pycache__/rewards.cpython-312.pyc b/src/mjlab_husky/tasks/skater/mdp/__pycache__/rewards.cpython-312.pyc deleted file mode 100644 index a6db09a..0000000 Binary files a/src/mjlab_husky/tasks/skater/mdp/__pycache__/rewards.cpython-312.pyc and /dev/null differ diff --git a/src/mjlab_husky/tasks/skater/mdp/__pycache__/terminations.cpython-312.pyc b/src/mjlab_husky/tasks/skater/mdp/__pycache__/terminations.cpython-312.pyc deleted file mode 100644 index 84ee485..0000000 Binary files a/src/mjlab_husky/tasks/skater/mdp/__pycache__/terminations.cpython-312.pyc and /dev/null differ diff --git a/src/mjlab_husky/tasks/skater/mdp/__pycache__/velocity_command.cpython-312.pyc b/src/mjlab_husky/tasks/skater/mdp/__pycache__/velocity_command.cpython-312.pyc deleted file mode 100644 index 6deb649..0000000 Binary files a/src/mjlab_husky/tasks/skater/mdp/__pycache__/velocity_command.cpython-312.pyc and /dev/null differ diff --git a/src/mjlab_husky/tasks/skater/rl/__pycache__/__init__.cpython-312.pyc b/src/mjlab_husky/tasks/skater/rl/__pycache__/__init__.cpython-312.pyc deleted file mode 100644 index dac96fb..0000000 Binary files a/src/mjlab_husky/tasks/skater/rl/__pycache__/__init__.cpython-312.pyc and /dev/null differ diff --git a/src/mjlab_husky/tasks/skater/rl/__pycache__/exporter.cpython-312.pyc b/src/mjlab_husky/tasks/skater/rl/__pycache__/exporter.cpython-312.pyc deleted file mode 100644 index 28cf1aa..0000000 Binary files a/src/mjlab_husky/tasks/skater/rl/__pycache__/exporter.cpython-312.pyc and /dev/null differ diff --git a/src/mjlab_husky/tasks/skater/rl/__pycache__/runner.cpython-312.pyc b/src/mjlab_husky/tasks/skater/rl/__pycache__/runner.cpython-312.pyc deleted file mode 100644 index ca95cbd..0000000 Binary files a/src/mjlab_husky/tasks/skater/rl/__pycache__/runner.cpython-312.pyc and /dev/null differ diff --git a/src/mjlab_husky/viewer/__pycache__/__init__.cpython-312.pyc b/src/mjlab_husky/viewer/__pycache__/__init__.cpython-312.pyc deleted file mode 100644 index 32299af..0000000 Binary files a/src/mjlab_husky/viewer/__pycache__/__init__.cpython-312.pyc and /dev/null differ diff --git a/src/mjlab_husky/viewer/__pycache__/rerun_play_viewer.cpython-312.pyc b/src/mjlab_husky/viewer/__pycache__/rerun_play_viewer.cpython-312.pyc deleted file mode 100644 index 9f574c4..0000000 Binary files a/src/mjlab_husky/viewer/__pycache__/rerun_play_viewer.cpython-312.pyc and /dev/null differ diff --git a/src/mjlab_husky/viewer/rerun_native_play_viewer.py b/src/mjlab_husky/viewer/rerun_native_play_viewer.py new file mode 100644 index 0000000..2c60be9 --- /dev/null +++ b/src/mjlab_husky/viewer/rerun_native_play_viewer.py @@ -0,0 +1,54 @@ +"""MuJoCo 原生 passive 窗口 + Rerun Web:同一套 mjlab 仿真、同一策略步进。""" + +from __future__ import annotations + +import os +from typing import Any + +from mjlab.viewer.base import VerbosityLevel +from mjlab.viewer.native import NativeMujocoViewer + +from mjlab_husky.viewer.rerun_play_viewer import RerunPlayViewer, RerunPlayViewerCfg + + +class RerunNativePlayViewer(NativeMujocoViewer): + """在 `NativeMujocoViewer` 的 `BaseViewer` 主循环上叠加 Rerun 时间轴与相机/qpos 日志。 + + 要求有图形会话(``DISPLAY`` / ``WAYLAND_DISPLAY``),否则无法启动 ``launch_passive``。 + 离屏 ``rgb_array`` 给 Rerun 的相机与 MuJoCo 窗口使用同一份 ``env.unwrapped.sim`` 状态。 + """ + + def __init__( + self, + env: Any, + policy: Any, + rerun_cfg: RerunPlayViewerCfg, + frame_rate: float = 60.0, + verbosity: int = VerbosityLevel.INFO, + ) -> None: + super().__init__(env, policy, frame_rate=frame_rate, verbosity=verbosity) + self._rerun = RerunPlayViewer(env, policy, rerun_cfg) + + def setup(self) -> None: + if not (os.environ.get("DISPLAY") or os.environ.get("WAYLAND_DISPLAY")): + raise RuntimeError( + "viewer=rerun_native 需要 **MuJoCo 自带的 GLFW 原生窗口**(仓库 `NativeMujocoViewer`)," + "因此进程里必须有 DISPLAY / WAYLAND_DISPLAY。\n" + "纯 SSH / RoboHub 技能容器里若未注入显示环境变量,此处为空就会失败;可选:\n" + " • 只要 Rerun、不要 GLFW 面板: --viewer rerun\n" + " • 虚拟屏幕(常见无桌面服务器): sudo apt-get install -y xvfb\n" + " xvfb-run -a uv run play ... --viewer rerun_native ...\n" + " • 本机有桌面 / 远程桌面:先 export DISPLAY=:0(或实际编号),再运行。\n" + "当前 DISPLAY/WAYLAND_DISPLAY 均为未设置。" + ) + print( + "[Rerun+Native] 单后端:策略与物理与 `play --viewer native` 相同," + "并额外把离屏相机与 qpos 写入 Rerun(见终端里的 http 链接)。" + ) + self._rerun.start_rerun_servers(app_id="mjlab_husky_play_rerun_native") + super().setup() + self.reset_environment() + + def sync_env_to_viewer(self) -> None: + super().sync_env_to_viewer() + self._rerun.log_post_step_to_rerun(self._step_count) diff --git a/src/mjlab_husky/viewer/rerun_play_viewer.py b/src/mjlab_husky/viewer/rerun_play_viewer.py index 1aa9a02..d264393 100644 --- a/src/mjlab_husky/viewer/rerun_play_viewer.py +++ b/src/mjlab_husky/viewer/rerun_play_viewer.py @@ -15,7 +15,7 @@ import torch import rerun as rr -from mjlab_husky.lerobot_numpy import batched_vector_to_numpy1d +from mjlab_husky.lerobot_numpy import batched_vector_to_numpy1d, import_lerobot_dataset_class def _downsample_max_side(img: np.ndarray, max_side: int) -> np.ndarray: @@ -164,8 +164,9 @@ class RerunPlayViewer: img = _downsample_max_side(img, self._cfg.camera_max_side) rr.log("camera/robot", rr.Image(img)) - def run(self): - rr.init("mjlab_husky_play_rerun", spawn=False) + def start_rerun_servers(self, *, app_id: str = "mjlab_husky_play_rerun") -> None: + """Initialize Rerun + gRPC + web viewer (no sim loop). Used by ``run()`` and ``RerunNativePlayViewer``.""" + rr.init(app_id, spawn=False) sg_kw: dict[str, Any] = {"newest_first": self._cfg.grpc_newest_first} if self._cfg.grpc_port is not None: @@ -228,12 +229,21 @@ class RerunPlayViewer: connect_to=connect_uri, ) + def log_post_step_to_rerun(self, step: int) -> None: + """After ``env.step``, stream qpos / offscreen camera to Rerun (same timeline as standalone ``run``).""" + rr.set_time("step", sequence=step) + self._try_log_mujoco_state(step) + self._try_log_camera(step) + + def run(self): + self.start_rerun_servers() + reset_out = self._env.reset() obs = reset_out[0] if isinstance(reset_out, tuple) and len(reset_out) == 2 else reset_out step = 0 if self._cfg.lerobot_record: - from lerobot.datasets.lerobot_dataset import LeRobotDataset # type: ignore + LeRobotDataset = import_lerobot_dataset_class() unwrapped = getattr(self._env, "unwrapped", self._env) sim_data = getattr(getattr(unwrapped, "sim", None), "data", None) diff --git a/src/mjlab_husky/viewer/rerun_viser_play_viewer.py b/src/mjlab_husky/viewer/rerun_viser_play_viewer.py new file mode 100644 index 0000000..3e0c442 --- /dev/null +++ b/src/mjlab_husky/viewer/rerun_viser_play_viewer.py @@ -0,0 +1,65 @@ +"""Viser(浏览器里的 mjlab 三维面板)+ Rerun:单套仿真,两个服务各占不同 TCP 端口。""" + +from __future__ import annotations + +from typing import Any + +import viser +from typing_extensions import override + +from mjlab.viewer.base import VerbosityLevel +from mjlab.viewer.viser import ViserPlayViewer + +from mjlab_husky.viewer.rerun_play_viewer import RerunPlayViewer, RerunPlayViewerCfg + + +class RerunViserPlayViewer(ViserPlayViewer): + """与 ``ViserPlayViewer`` 相同的主循环,额外启动 Rerun Web+gRPC,并在每步写入时间轴。 + + - **Viser** 占用 ``--viser-port``(默认 `8081`,避免与 Rerun Web 默认 `8080` 冲突)。 + - **Rerun** 占用 ``--rerun-web-port`` 与 ``--rerun-grpc-port``(与 ``--viewer rerun`` 一致)。 + + GLFW ``NativeMujocoViewer`` 不监听端口;若要在浏览器里并排嵌入「三维 + Rerun」,请用本 viewer 而非 ``rerun_native``。 + """ + + def __init__( + self, + env: Any, + policy: Any, + rerun_cfg: RerunPlayViewerCfg, + viser_port: int = 8081, + frame_rate: float = 60.0, + verbosity: int = VerbosityLevel.SILENT, + ) -> None: + super().__init__(env, policy, frame_rate=frame_rate, verbosity=verbosity) + self._rerun = RerunPlayViewer(env, policy, rerun_cfg) + self._viser_port = int(viser_port) + self._url_host = (rerun_cfg.connect_host or "127.0.0.1").strip() + + @override + def setup(self) -> None: + self._rerun.start_rerun_servers(app_id="mjlab_husky_play_rerun_viser") + + _real = viser.ViserServer + port = self._viser_port + + def _force_port_ViserServer(*args: Any, **kwargs: Any): + kwargs = dict(kwargs) + kwargs["port"] = port + return _real(*args, **kwargs) + + viser.ViserServer = _force_port_ViserServer # type: ignore[misc, assignment] + try: + super().setup() + finally: + viser.ViserServer = _real + + print( + f"[Rerun+Viser] 单后端:Rerun 使用 --rerun-web-port / --rerun-grpc-port;" + f"Viser 三维面板: http://{self._url_host}:{port}/" + ) + + @override + def sync_env_to_viewer(self) -> None: + super().sync_env_to_viewer() + self._rerun.log_post_step_to_rerun(self._step_count)