updated
This commit is contained in:
Bug
parent
7e3da40be3
commit
fe683c1e1b
61
README.md
61
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`，以终端打印为准）。
|
||||
|
|
|
|||
Binary file not shown.
|
|
@ -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" }
|
||||
mjlab = { git = "https://github.com/mujocolab/mjlab.git", rev = "13212ad" }
|
||||
rsl-rl-lib = { path = "rsl_rl" }
|
||||
|
|
|
|||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
|
|
@ -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."""
|
||||
|
|
@ -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"]
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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."""
|
||||
|
|
@ -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"]
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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."""
|
||||
|
|
@ -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"]
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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"]
|
||||
|
|
@ -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
|
||||
|
||||
|
|
@ -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",
|
||||
]
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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")
|
||||
|
|
@ -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")
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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()
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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"]
|
||||
|
|
@ -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'.")
|
||||
|
|
@ -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
|
||||
|
|
@ -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'.")
|
||||
|
|
@ -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"]
|
||||
|
|
@ -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))
|
||||
|
|
@ -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
|
||||
|
|
@ -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",
|
||||
]
|
||||
|
|
@ -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]
|
||||
|
|
@ -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
|
||||
|
|
@ -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)
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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"]
|
||||
|
|
@ -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
|
||||
|
||||
|
|
@ -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",
|
||||
]
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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")
|
||||
|
|
@ -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")
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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()
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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"]
|
||||
|
|
@ -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'.")
|
||||
|
|
@ -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
|
||||
|
|
@ -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'.")
|
||||
|
|
@ -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"]
|
||||
|
|
@ -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))
|
||||
|
|
@ -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
|
||||
|
|
@ -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",
|
||||
]
|
||||
|
|
@ -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]
|
||||
|
|
@ -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
|
||||
|
|
@ -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)
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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"]
|
||||
|
|
@ -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
|
||||
|
||||
|
|
@ -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",
|
||||
]
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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")
|
||||
|
|
@ -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")
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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()
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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"]
|
||||
|
|
@ -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'.")
|
||||
|
|
@ -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
|
||||
|
|
@ -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'.")
|
||||
|
|
@ -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"]
|
||||
Some files were not shown because too many files have changed in this diff Show More
Loading…
Reference in New Issue