init commit.
This commit is contained in:
wertyuilife
24
.flake8
Normal file
24
.flake8
Normal file
@@ -0,0 +1,24 @@
|
||||
# copied from https://github.com/isaac-sim/IsaacLab/blob/main/.flake8
|
||||
|
||||
[flake8]
|
||||
show-source=True
|
||||
statistics=True
|
||||
per-file-ignores=*/__init__.py:F401
|
||||
# E402: Module level import not at top of file
|
||||
# E501: Line too long
|
||||
# W503: Line break before binary operator
|
||||
# E203: Whitespace before ':' -> conflicts with black
|
||||
# D401: First line should be in imperative mood
|
||||
# R504: Unnecessary variable assignment before return statement.
|
||||
# R505: Unnecessary elif after return statement
|
||||
# SIM102: Use a single if-statement instead of nested if-statements
|
||||
# SIM117: Merge with statements for context managers that have same scope.
|
||||
ignore=E402,E501,W503,E203,D401,R504,R505,SIM102,SIM117
|
||||
max-line-length = 120
|
||||
max-complexity = 30
|
||||
exclude=_*,.vscode,.git,docs/**
|
||||
# docstrings
|
||||
docstring-convention=google
|
||||
# annotations
|
||||
suppress-none-returning=True
|
||||
allow-star-arg-any=True
|
||||
89
.gitignore
vendored
Normal file
89
.gitignore
vendored
Normal file
@@ -0,0 +1,89 @@
|
||||
# Omniverse
|
||||
**/*.dmp
|
||||
**/.thumbs
|
||||
|
||||
# Python
|
||||
.DS_Store
|
||||
**/*.egg-info/
|
||||
**/__pycache__/
|
||||
**/.pytest_cache/
|
||||
**/*.pyc
|
||||
**/*.pb
|
||||
|
||||
# IDE
|
||||
**/.idea/
|
||||
**/.vscode/
|
||||
# Don't ignore the top-level .vscode directory as it is
|
||||
# used to configure VS Code settings
|
||||
!.vscode
|
||||
|
||||
# Outputs
|
||||
**/runs/*
|
||||
**/logs/*
|
||||
**/recordings/*
|
||||
**/output/*
|
||||
**/outputs/*
|
||||
**/videos/*
|
||||
**/wandb/*
|
||||
**/.neptune/*
|
||||
docker/artifacts/
|
||||
*.tmp
|
||||
*.npz
|
||||
|
||||
# Isaac-Sim packman
|
||||
_isaac_sim*
|
||||
_repo
|
||||
_build
|
||||
.lastformat
|
||||
|
||||
# LSP
|
||||
pyrightconfig.json
|
||||
build/
|
||||
**/build/
|
||||
*.o
|
||||
*.so
|
||||
*.a
|
||||
*.log
|
||||
__pycache__/
|
||||
.DS_Store
|
||||
.vscode/
|
||||
.idea/
|
||||
# logs
|
||||
logs/
|
||||
**/logs/
|
||||
|
||||
# build
|
||||
build/
|
||||
**/build/
|
||||
|
||||
# binaries
|
||||
*.o
|
||||
*.so
|
||||
*.a
|
||||
|
||||
# cache
|
||||
__pycache__/
|
||||
*.pyc
|
||||
|
||||
# editor
|
||||
.vscode/
|
||||
.idea/
|
||||
.DS_Store
|
||||
|
||||
# ML checkpoints / datasets
|
||||
*.pt
|
||||
*.pth
|
||||
*.onnx
|
||||
*.npz
|
||||
*.npy
|
||||
|
||||
# Large ML artifacts (do not commit)
|
||||
*.pt
|
||||
*.pth
|
||||
*.ckpt
|
||||
*.onnx
|
||||
|
||||
# logs
|
||||
logs/
|
||||
**/logs/
|
||||
outputs/
|
||||
201
LICENSE
Normal file
201
LICENSE
Normal file
@@ -0,0 +1,201 @@
|
||||
Apache License
|
||||
Version 2.0, January 2004
|
||||
http://www.apache.org/licenses/
|
||||
|
||||
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
|
||||
|
||||
1. Definitions.
|
||||
|
||||
"License" shall mean the terms and conditions for use, reproduction,
|
||||
and distribution as defined by Sections 1 through 9 of this document.
|
||||
|
||||
"Licensor" shall mean the copyright owner or entity authorized by
|
||||
the copyright owner that is granting the License.
|
||||
|
||||
"Legal Entity" shall mean the union of the acting entity and all
|
||||
other entities that control, are controlled by, or are under common
|
||||
control with that entity. For the purposes of this definition,
|
||||
"control" means (i) the power, direct or indirect, to cause the
|
||||
direction or management of such entity, whether by contract or
|
||||
otherwise, or (ii) ownership of fifty percent (50%) or more of the
|
||||
outstanding shares, or (iii) beneficial ownership of such entity.
|
||||
|
||||
"You" (or "Your") shall mean an individual or Legal Entity
|
||||
exercising permissions granted by this License.
|
||||
|
||||
"Source" form shall mean the preferred form for making modifications,
|
||||
including but not limited to software source code, documentation
|
||||
source, and configuration files.
|
||||
|
||||
"Object" form shall mean any form resulting from mechanical
|
||||
transformation or translation of a Source form, including but
|
||||
not limited to compiled object code, generated documentation,
|
||||
and conversions to other media types.
|
||||
|
||||
"Work" shall mean the work of authorship, whether in Source or
|
||||
Object form, made available under the License, as indicated by a
|
||||
copyright notice that is included in or attached to the work
|
||||
(an example is provided in the Appendix below).
|
||||
|
||||
"Derivative Works" shall mean any work, whether in Source or Object
|
||||
form, that is based on (or derived from) the Work and for which the
|
||||
editorial revisions, annotations, elaborations, or other modifications
|
||||
represent, as a whole, an original work of authorship. For the purposes
|
||||
of this License, Derivative Works shall not include works that remain
|
||||
separable from, or merely link (or bind by name) to the interfaces of,
|
||||
the Work and Derivative Works thereof.
|
||||
|
||||
"Contribution" shall mean any work of authorship, including
|
||||
the original version of the Work and any modifications or additions
|
||||
to that Work or Derivative Works thereof, that is intentionally
|
||||
submitted to Licensor for inclusion in the Work by the copyright owner
|
||||
or by an individual or Legal Entity authorized to submit on behalf of
|
||||
the copyright owner. For the purposes of this definition, "submitted"
|
||||
means any form of electronic, verbal, or written communication sent
|
||||
to the Licensor or its representatives, including but not limited to
|
||||
communication on electronic mailing lists, source code control systems,
|
||||
and issue tracking systems that are managed by, or on behalf of, the
|
||||
Licensor for the purpose of discussing and improving the Work, but
|
||||
excluding communication that is conspicuously marked or otherwise
|
||||
designated in writing by the copyright owner as "Not a Contribution."
|
||||
|
||||
"Contributor" shall mean Licensor and any individual or Legal Entity
|
||||
on behalf of whom a Contribution has been received by Licensor and
|
||||
subsequently incorporated within the Work.
|
||||
|
||||
2. Grant of Copyright License. Subject to the terms and conditions of
|
||||
this License, each Contributor hereby grants to You a perpetual,
|
||||
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||
copyright license to reproduce, prepare Derivative Works of,
|
||||
publicly display, publicly perform, sublicense, and distribute the
|
||||
Work and such Derivative Works in Source or Object form.
|
||||
|
||||
3. Grant of Patent License. Subject to the terms and conditions of
|
||||
this License, each Contributor hereby grants to You a perpetual,
|
||||
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||
(except as stated in this section) patent license to make, have made,
|
||||
use, offer to sell, sell, import, and otherwise transfer the Work,
|
||||
where such license applies only to those patent claims licensable
|
||||
by such Contributor that are necessarily infringed by their
|
||||
Contribution(s) alone or by combination of their Contribution(s)
|
||||
with the Work to which such Contribution(s) was submitted. If You
|
||||
institute patent litigation against any entity (including a
|
||||
cross-claim or counterclaim in a lawsuit) alleging that the Work
|
||||
or a Contribution incorporated within the Work constitutes direct
|
||||
or contributory patent infringement, then any patent licenses
|
||||
granted to You under this License for that Work shall terminate
|
||||
as of the date such litigation is filed.
|
||||
|
||||
4. Redistribution. You may reproduce and distribute copies of the
|
||||
Work or Derivative Works thereof in any medium, with or without
|
||||
modifications, and in Source or Object form, provided that You
|
||||
meet the following conditions:
|
||||
|
||||
(a) You must give any other recipients of the Work or
|
||||
Derivative Works a copy of this License; and
|
||||
|
||||
(b) You must cause any modified files to carry prominent notices
|
||||
stating that You changed the files; and
|
||||
|
||||
(c) You must retain, in the Source form of any Derivative Works
|
||||
that You distribute, all copyright, patent, trademark, and
|
||||
attribution notices from the Source form of the Work,
|
||||
excluding those notices that do not pertain to any part of
|
||||
the Derivative Works; and
|
||||
|
||||
(d) If the Work includes a "NOTICE" text file as part of its
|
||||
distribution, then any Derivative Works that You distribute must
|
||||
include a readable copy of the attribution notices contained
|
||||
within such NOTICE file, excluding those notices that do not
|
||||
pertain to any part of the Derivative Works, in at least one
|
||||
of the following places: within a NOTICE text file distributed
|
||||
as part of the Derivative Works; within the Source form or
|
||||
documentation, if provided along with the Derivative Works; or,
|
||||
within a display generated by the Derivative Works, if and
|
||||
wherever such third-party notices normally appear. The contents
|
||||
of the NOTICE file are for informational purposes only and
|
||||
do not modify the License. You may add Your own attribution
|
||||
notices within Derivative Works that You distribute, alongside
|
||||
or as an addendum to the NOTICE text from the Work, provided
|
||||
that such additional attribution notices cannot be construed
|
||||
as modifying the License.
|
||||
|
||||
You may add Your own copyright statement to Your modifications and
|
||||
may provide additional or different license terms and conditions
|
||||
for use, reproduction, or distribution of Your modifications, or
|
||||
for any such Derivative Works as a whole, provided Your use,
|
||||
reproduction, and distribution of the Work otherwise complies with
|
||||
the conditions stated in this License.
|
||||
|
||||
5. Submission of Contributions. Unless You explicitly state otherwise,
|
||||
any Contribution intentionally submitted for inclusion in the Work
|
||||
by You to the Licensor shall be under the terms and conditions of
|
||||
this License, without any additional terms or conditions.
|
||||
Notwithstanding the above, nothing herein shall supersede or modify
|
||||
the terms of any separate license agreement you may have executed
|
||||
with Licensor regarding such Contributions.
|
||||
|
||||
6. Trademarks. This License does not grant permission to use the trade
|
||||
names, trademarks, service marks, or product names of the Licensor,
|
||||
except as required for reasonable and customary use in describing the
|
||||
origin of the Work and reproducing the content of the NOTICE file.
|
||||
|
||||
7. Disclaimer of Warranty. Unless required by applicable law or
|
||||
agreed to in writing, Licensor provides the Work (and each
|
||||
Contributor provides its Contributions) on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
|
||||
implied, including, without limitation, any warranties or conditions
|
||||
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
|
||||
PARTICULAR PURPOSE. You are solely responsible for determining the
|
||||
appropriateness of using or redistributing the Work and assume any
|
||||
risks associated with Your exercise of permissions under this License.
|
||||
|
||||
8. Limitation of Liability. In no event and under no legal theory,
|
||||
whether in tort (including negligence), contract, or otherwise,
|
||||
unless required by applicable law (such as deliberate and grossly
|
||||
negligent acts) or agreed to in writing, shall any Contributor be
|
||||
liable to You for damages, including any direct, indirect, special,
|
||||
incidental, or consequential damages of any character arising as a
|
||||
result of this License or out of the use or inability to use the
|
||||
Work (including but not limited to damages for loss of goodwill,
|
||||
work stoppage, computer failure or malfunction, or any and all
|
||||
other commercial damages or losses), even if such Contributor
|
||||
has been advised of the possibility of such damages.
|
||||
|
||||
9. Accepting Warranty or Additional Liability. While redistributing
|
||||
the Work or Derivative Works thereof, You may choose to offer,
|
||||
and charge a fee for, acceptance of support, warranty, indemnity,
|
||||
or other liability obligations and/or rights consistent with this
|
||||
License. However, in accepting such obligations, You may act only
|
||||
on Your own behalf and on Your sole responsibility, not on behalf
|
||||
of any other Contributor, and only if You agree to indemnify,
|
||||
defend, and hold each Contributor harmless for any liability
|
||||
incurred by, or claims asserted against, such Contributor by reason
|
||||
of your accepting any such warranty or additional liability.
|
||||
|
||||
END OF TERMS AND CONDITIONS
|
||||
|
||||
APPENDIX: How to apply the Apache License to your work.
|
||||
|
||||
To apply the Apache License to your work, attach the following
|
||||
boilerplate notice, with the fields enclosed by brackets "[]"
|
||||
replaced with your own identifying information. (Don't include
|
||||
the brackets!) The text should be enclosed in the appropriate
|
||||
comment syntax for the file format. We also recommend that a
|
||||
file or class name and description of purpose be included on the
|
||||
same "printed page" as the copyright notice for easier
|
||||
identification within third-party archives.
|
||||
|
||||
Copyright 2024-2025 Ziqi Fan
|
||||
|
||||
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.
|
||||
41
README.md
Normal file
41
README.md
Normal file
@@ -0,0 +1,41 @@
|
||||
# go2_rl_robotlab
|
||||
|
||||
## Overview
|
||||
|
||||
Train Unitree Go2 with MoE-CTS.
|
||||
|
||||
This is a reproduction version of [go2_rl_gym](https://github.com/wty-yy/go2_rl_gym) on robotlab/isaaclab.
|
||||
|
||||
## Installation Guide
|
||||
|
||||
### 1. Base Installation
|
||||
Install Isaac Lab by following the [installation guide](https://isaac-sim.github.io/IsaacLab/main/source/setup/installation/index.html).
|
||||
|
||||
Notice that we use certain version of IsaacLab packages, make sure:
|
||||
```
|
||||
isaacsim <= 5.1.0.0 # tested on 5.1.0.0
|
||||
isaaclab <= 0.53.1 # tested on 0.53.1
|
||||
isaaclab-rl <= 0.4.7 # tested on 0.4.7
|
||||
```
|
||||
|
||||
### 2. Modified Library Setup
|
||||
This branch uses a customized version of `rsl_rl` and `robot_lab`. To install it, run the following commands in your terminal:
|
||||
|
||||
```bash
|
||||
python -m pip install -e source/robot_lab
|
||||
python -m pip install -e source/rsl_rl
|
||||
```
|
||||
|
||||
## Try examples
|
||||
|
||||
You can use the following commands to run all environments:
|
||||
|
||||
RSL-RL:
|
||||
|
||||
```bash
|
||||
# Train
|
||||
python scripts/reinforcement_learning/rsl_rl/train.py --task=<ENV_NAME> --headless
|
||||
|
||||
# Play
|
||||
python scripts/reinforcement_learning/rsl_rl/play.py --task=<ENV_NAME>
|
||||
```
|
||||
7
cmd.md
Normal file
7
cmd.md
Normal file
@@ -0,0 +1,7 @@
|
||||
## train
|
||||
python scripts/reinforcement_learning/rsl_rl/train.py \
|
||||
--task=Robotlab-Go2-v0 --headless --num_envs=1024
|
||||
|
||||
## eval
|
||||
python scripts/reinforcement_learning/rsl_rl/play.py \
|
||||
--task=Robotlab-Go2-v0 --num_envs=64
|
||||
96
scripts/reinforcement_learning/rsl_rl/cli_args.py
Normal file
96
scripts/reinforcement_learning/rsl_rl/cli_args.py
Normal file
@@ -0,0 +1,96 @@
|
||||
# Copyright (c) 2024-2025 Ziqi Fan
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# Copyright (c) 2024-2025, The Isaac Lab Project Developers.
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import random
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from isaaclab_rl.rsl_rl import RslRlBaseRunnerCfg
|
||||
|
||||
|
||||
def add_rsl_rl_args(parser: argparse.ArgumentParser):
|
||||
"""Add RSL-RL arguments to the parser.
|
||||
|
||||
Args:
|
||||
parser: The parser to add the arguments to.
|
||||
"""
|
||||
# create a new argument group
|
||||
arg_group = parser.add_argument_group("rsl_rl", description="Arguments for RSL-RL agent.")
|
||||
# -- experiment arguments
|
||||
arg_group.add_argument(
|
||||
"--experiment_name", type=str, default=None, help="Name of the experiment folder where logs will be stored."
|
||||
)
|
||||
arg_group.add_argument("--run_name", type=str, default=None, help="Run name suffix to the log directory.")
|
||||
# -- load arguments
|
||||
arg_group.add_argument("--resume", action="store_true", default=False, help="Whether to resume from a checkpoint.")
|
||||
arg_group.add_argument("--load_run", type=str, default=None, help="Name of the run folder to resume from.")
|
||||
arg_group.add_argument("--checkpoint", type=str, default=None, help="Checkpoint file to resume from.")
|
||||
# -- logger arguments
|
||||
arg_group.add_argument(
|
||||
"--logger", type=str, default=None, choices={"wandb", "tensorboard", "neptune"}, help="Logger module to use."
|
||||
)
|
||||
arg_group.add_argument(
|
||||
"--log_project_name", type=str, default=None, help="Name of the logging project when using wandb or neptune."
|
||||
)
|
||||
|
||||
|
||||
def parse_rsl_rl_cfg(task_name: str, args_cli: argparse.Namespace) -> RslRlBaseRunnerCfg:
|
||||
"""Parse configuration for RSL-RL agent based on inputs.
|
||||
|
||||
Args:
|
||||
task_name: The name of the environment.
|
||||
args_cli: The command line arguments.
|
||||
|
||||
Returns:
|
||||
The parsed configuration for RSL-RL agent based on inputs.
|
||||
"""
|
||||
from isaaclab_tasks.utils.parse_cfg import load_cfg_from_registry
|
||||
|
||||
# load the default configuration
|
||||
rslrl_cfg: RslRlBaseRunnerCfg = load_cfg_from_registry(task_name, "rsl_rl_cfg_entry_point")
|
||||
rslrl_cfg = update_rsl_rl_cfg(rslrl_cfg, args_cli)
|
||||
return rslrl_cfg
|
||||
|
||||
|
||||
def update_rsl_rl_cfg(agent_cfg: RslRlBaseRunnerCfg, args_cli: argparse.Namespace):
|
||||
"""Update configuration for RSL-RL agent based on inputs.
|
||||
|
||||
Args:
|
||||
agent_cfg: The configuration for RSL-RL agent.
|
||||
args_cli: The command line arguments.
|
||||
|
||||
Returns:
|
||||
The updated configuration for RSL-RL agent based on inputs.
|
||||
"""
|
||||
# override the default configuration with CLI arguments
|
||||
if hasattr(args_cli, "seed") and args_cli.seed is not None:
|
||||
# randomly sample a seed if seed = -1
|
||||
if args_cli.seed == -1:
|
||||
args_cli.seed = random.randint(0, 10000)
|
||||
agent_cfg.seed = args_cli.seed
|
||||
if args_cli.experiment_name is not None:
|
||||
agent_cfg.experiment_name = args_cli.experiment_name
|
||||
if args_cli.resume is not None:
|
||||
agent_cfg.resume = args_cli.resume
|
||||
if args_cli.load_run is not None:
|
||||
agent_cfg.load_run = args_cli.load_run
|
||||
if args_cli.checkpoint is not None:
|
||||
agent_cfg.load_checkpoint = args_cli.checkpoint
|
||||
if args_cli.run_name is not None:
|
||||
agent_cfg.run_name = args_cli.run_name
|
||||
if args_cli.logger is not None:
|
||||
agent_cfg.logger = args_cli.logger
|
||||
# set the project name for wandb and neptune
|
||||
if agent_cfg.logger in {"wandb", "neptune"} and args_cli.log_project_name:
|
||||
agent_cfg.wandb_project = args_cli.log_project_name
|
||||
agent_cfg.neptune_project = args_cli.log_project_name
|
||||
|
||||
return agent_cfg
|
||||
228
scripts/reinforcement_learning/rsl_rl/play.py
Normal file
228
scripts/reinforcement_learning/rsl_rl/play.py
Normal file
@@ -0,0 +1,228 @@
|
||||
# Copyright (c) 2024-2025 Ziqi Fan
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# Copyright (c) 2024-2025, The Isaac Lab Project Developers.
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
"""Script to play a checkpoint if an RL agent from RSL-RL."""
|
||||
|
||||
"""Launch Isaac Sim Simulator first."""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import sys
|
||||
|
||||
from isaaclab.app import AppLauncher
|
||||
|
||||
# local imports
|
||||
import cli_args # isort: skip
|
||||
from utils import export_cts_policy_as_jit, export_cts_policy_as_onnx
|
||||
|
||||
# add argparse arguments
|
||||
parser = argparse.ArgumentParser(description="Train an RL agent with RSL-RL.")
|
||||
parser.add_argument("--video", action="store_true", default=False, help="Record videos during training.")
|
||||
parser.add_argument("--video_length", type=int, default=200, help="Length of the recorded video (in steps).")
|
||||
parser.add_argument(
|
||||
"--disable_fabric", action="store_true", default=False, help="Disable fabric and use USD I/O operations."
|
||||
)
|
||||
parser.add_argument("--num_envs", type=int, default=None, help="Number of environments to simulate.")
|
||||
parser.add_argument("--task", type=str, default=None, help="Name of the task.")
|
||||
parser.add_argument(
|
||||
"--agent", type=str, default="rsl_rl_cfg_entry_point", help="Name of the RL agent configuration entry point."
|
||||
)
|
||||
parser.add_argument("--seed", type=int, default=None, help="Seed used for the environment")
|
||||
parser.add_argument(
|
||||
"--use_pretrained_checkpoint",
|
||||
action="store_true",
|
||||
help="Use the pre-trained checkpoint from Nucleus.",
|
||||
)
|
||||
parser.add_argument("--real-time", action="store_true", default=False, help="Run in real-time, if possible.")
|
||||
parser.add_argument("--keyboard", action="store_true", default=False, help="Whether to use keyboard.")
|
||||
# append RSL-RL cli arguments
|
||||
cli_args.add_rsl_rl_args(parser)
|
||||
# append AppLauncher cli args
|
||||
AppLauncher.add_app_launcher_args(parser)
|
||||
# parse the arguments
|
||||
args_cli, hydra_args = parser.parse_known_args()
|
||||
# always enable cameras to record video
|
||||
if args_cli.video:
|
||||
args_cli.enable_cameras = True
|
||||
|
||||
# clear out sys.argv for Hydra
|
||||
sys.argv = [sys.argv[0]] + hydra_args
|
||||
|
||||
# launch omniverse app
|
||||
app_launcher = AppLauncher(args_cli)
|
||||
simulation_app = app_launcher.app
|
||||
|
||||
"""Rest everything follows."""
|
||||
|
||||
import gymnasium as gym
|
||||
import time
|
||||
import torch
|
||||
# from scripts.reinforcement_learning.utils import camera_follow
|
||||
from rsl_rl.runners import DistillationRunner, OnPolicyRunner, OnPolicyRunnerCTS
|
||||
|
||||
from isaaclab.devices import Se2Keyboard, Se2KeyboardCfg
|
||||
from isaaclab.envs import (
|
||||
DirectMARLEnv,
|
||||
DirectMARLEnvCfg,
|
||||
DirectRLEnvCfg,
|
||||
ManagerBasedRLEnvCfg,
|
||||
multi_agent_to_single_agent,
|
||||
)
|
||||
from isaaclab.utils.assets import retrieve_file_path
|
||||
from isaaclab.utils.dict import print_dict
|
||||
# from isaaclab.utils.pretrained_checkpoint import get_published_pretrained_checkpoint
|
||||
from isaaclab_rl.rsl_rl import RslRlBaseRunnerCfg, RslRlVecEnvWrapper, export_policy_as_jit, export_policy_as_onnx
|
||||
from isaaclab_tasks.utils import get_checkpoint_path
|
||||
from isaaclab_tasks.utils.hydra import hydra_task_config
|
||||
import robot_lab.tasks # noqa: F401
|
||||
|
||||
|
||||
@hydra_task_config(args_cli.task, args_cli.agent)
|
||||
def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agent_cfg: RslRlBaseRunnerCfg):
|
||||
"""Play with RSL-RL agent."""
|
||||
# grab task name for checkpoint path
|
||||
task_name = args_cli.task.split(":")[-1]
|
||||
|
||||
# override configurations with non-hydra CLI arguments
|
||||
agent_cfg: RslRlBaseRunnerCfg = cli_args.update_rsl_rl_cfg(agent_cfg, args_cli)
|
||||
env_cfg.scene.num_envs = args_cli.num_envs if args_cli.num_envs is not None else 64
|
||||
|
||||
# set the environment seed
|
||||
# note: certain randomizations occur in the environment initialization so we set the seed here
|
||||
env_cfg.seed = agent_cfg.seed
|
||||
env_cfg.sim.device = args_cli.device if args_cli.device is not None else env_cfg.sim.device
|
||||
|
||||
# disable randomization for play
|
||||
env_cfg.observations.policy.enable_corruption = False
|
||||
# remove random pushing
|
||||
env_cfg.events.randomize_apply_external_force_torque = None
|
||||
env_cfg.events.push_robot = None
|
||||
env_cfg.curriculum.command_levels_lin_vel = None
|
||||
env_cfg.curriculum.command_levels_ang_vel = None
|
||||
|
||||
# specify directory for logging experiments
|
||||
log_root_path = os.path.join("logs", "rsl_rl", agent_cfg.experiment_name)
|
||||
log_root_path = os.path.abspath(log_root_path)
|
||||
print(f"[INFO] Loading experiment from directory: {log_root_path}")
|
||||
if args_cli.use_pretrained_checkpoint:
|
||||
# resume_path = get_published_pretrained_checkpoint("rsl_rl", task_name)
|
||||
# if not resume_path:
|
||||
# print("[INFO] Unfortunately a pre-trained checkpoint is currently unavailable for this task.")
|
||||
# return
|
||||
raise NotImplementedError("Pre-trained checkpoint retrieval is disabled temporarily.")
|
||||
elif args_cli.checkpoint:
|
||||
resume_path = retrieve_file_path(args_cli.checkpoint)
|
||||
else:
|
||||
resume_path = get_checkpoint_path(log_root_path, agent_cfg.load_run, agent_cfg.load_checkpoint)
|
||||
|
||||
log_dir = os.path.dirname(resume_path)
|
||||
|
||||
# set the log directory for the environment (works for all environment types)
|
||||
env_cfg.log_dir = log_dir
|
||||
|
||||
# create isaac environment
|
||||
env = gym.make(args_cli.task, cfg=env_cfg, render_mode="rgb_array" if args_cli.video else None)
|
||||
|
||||
# convert to single-agent instance if required by the RL algorithm
|
||||
if isinstance(env.unwrapped, DirectMARLEnv):
|
||||
env = multi_agent_to_single_agent(env)
|
||||
|
||||
# wrap for video recording
|
||||
if args_cli.video:
|
||||
video_kwargs = {
|
||||
"video_folder": os.path.join(log_dir, "videos", "play"),
|
||||
"step_trigger": lambda step: step == 0,
|
||||
"video_length": args_cli.video_length,
|
||||
"disable_logger": True,
|
||||
}
|
||||
print("[INFO] Recording videos during training.")
|
||||
print_dict(video_kwargs, nesting=4)
|
||||
env = gym.wrappers.RecordVideo(env, **video_kwargs)
|
||||
|
||||
# wrap around environment for rsl-rl
|
||||
env = RslRlVecEnvWrapper(env, clip_actions=agent_cfg.clip_actions)
|
||||
|
||||
print(f"[INFO]: Loading model checkpoint from: {resume_path}")
|
||||
# load previously trained model
|
||||
if agent_cfg.class_name == "OnPolicyRunner":
|
||||
runner = OnPolicyRunner(env, agent_cfg.to_dict(), log_dir=None, device=agent_cfg.device)
|
||||
elif agent_cfg.class_name == "DistillationRunner":
|
||||
runner = DistillationRunner(env, agent_cfg.to_dict(), log_dir=None, device=agent_cfg.device)
|
||||
elif agent_cfg.class_name == "OnPolicyRunnerCTS":
|
||||
runner = OnPolicyRunnerCTS(env, agent_cfg.to_dict(), log_dir=log_dir, device=agent_cfg.device)
|
||||
else:
|
||||
raise ValueError(f"Unsupported runner class: {agent_cfg.class_name}")
|
||||
runner.load(resume_path)
|
||||
|
||||
# obtain the trained policy for inference
|
||||
policy = runner.get_inference_policy(device=env.unwrapped.device)
|
||||
|
||||
# extract the neural network module
|
||||
# we do this in a try-except to maintain backwards compatibility.
|
||||
try:
|
||||
# version 2.3 onwards
|
||||
policy_nn = runner.alg.policy
|
||||
except AttributeError:
|
||||
# version 2.2 and below
|
||||
policy_nn = runner.alg.actor_critic
|
||||
|
||||
# extract the normalizer
|
||||
if hasattr(policy_nn, "actor_obs_normalizer"):
|
||||
normalizer = policy_nn.actor_obs_normalizer
|
||||
elif hasattr(policy_nn, "student_obs_normalizer"):
|
||||
normalizer = policy_nn.student_obs_normalizer
|
||||
else:
|
||||
normalizer = None
|
||||
|
||||
# export policy to onnx/jit
|
||||
export_model_dir = os.path.join(os.path.dirname(resume_path), "exported")
|
||||
if agent_cfg.class_name == "OnPolicyRunnerCTS":
|
||||
export_cts_policy_as_jit(policy_nn, actor_obs_normalizer=policy_nn.actor_obs_normalizer, single_obs_normalizer=policy_nn.single_obs_normalizer, path=export_model_dir, filename="policy.pt")
|
||||
export_cts_policy_as_onnx(policy_nn, actor_obs_normalizer=policy_nn.actor_obs_normalizer, single_obs_normalizer=policy_nn.single_obs_normalizer, path=export_model_dir, filename="policy.onnx")
|
||||
else:
|
||||
export_policy_as_jit(policy_nn, normalizer=normalizer, path=export_model_dir, filename="policy.pt")
|
||||
export_policy_as_onnx(policy_nn, normalizer=normalizer, path=export_model_dir, filename="policy.onnx")
|
||||
|
||||
dt = env.unwrapped.step_dt
|
||||
|
||||
# env.unwrapped.eye = (1.1, 3.3, 0.9)
|
||||
# reset environment
|
||||
obs = env.get_observations()
|
||||
timestep = 0
|
||||
# simulate environment
|
||||
while simulation_app.is_running():
|
||||
start_time = time.time()
|
||||
# run everything in inference mode
|
||||
with torch.inference_mode():
|
||||
# agent stepping
|
||||
actions = policy(obs)
|
||||
# env stepping
|
||||
obs, _, dones, _ = env.step(actions)
|
||||
# reset recurrent states for episodes that have terminated
|
||||
policy_nn.reset(dones)
|
||||
if args_cli.video:
|
||||
timestep += 1
|
||||
# Exit the play loop after recording one video
|
||||
if timestep == args_cli.video_length:
|
||||
break
|
||||
# camera_follow(env)
|
||||
|
||||
# time delay for real-time evaluation
|
||||
sleep_time = dt - (time.time() - start_time)
|
||||
if args_cli.real_time and sleep_time > 0:
|
||||
time.sleep(sleep_time)
|
||||
|
||||
# close the simulator
|
||||
env.close()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# run the main function
|
||||
main()
|
||||
# close sim app
|
||||
simulation_app.close()
|
||||
218
scripts/reinforcement_learning/rsl_rl/train.py
Normal file
218
scripts/reinforcement_learning/rsl_rl/train.py
Normal file
@@ -0,0 +1,218 @@
|
||||
# Copyright (c) 2024-2025 Ziqi Fan
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# Copyright (c) 2024-2025, The Isaac Lab Project Developers.
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
"""Script to train RL agent with RSL-RL."""
|
||||
|
||||
"""Launch Isaac Sim Simulator first."""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import sys
|
||||
|
||||
from isaaclab.app import AppLauncher
|
||||
|
||||
# local imports
|
||||
import cli_args # isort: skip
|
||||
|
||||
# add argparse arguments
|
||||
parser = argparse.ArgumentParser(description="Train an RL agent with RSL-RL.")
|
||||
parser.add_argument("--video", action="store_true", default=False, help="Record videos during training.")
|
||||
parser.add_argument("--video_length", type=int, default=200, help="Length of the recorded video (in steps).")
|
||||
parser.add_argument("--video_interval", type=int, default=2000, help="Interval between video recordings (in steps).")
|
||||
parser.add_argument("--num_envs", type=int, default=None, help="Number of environments to simulate.")
|
||||
parser.add_argument("--task", type=str, default=None, help="Name of the task.")
|
||||
parser.add_argument(
|
||||
"--agent", type=str, default="rsl_rl_cfg_entry_point", help="Name of the RL agent configuration entry point."
|
||||
)
|
||||
parser.add_argument("--seed", type=int, default=None, help="Seed used for the environment")
|
||||
parser.add_argument("--max_iterations", type=int, default=None, help="RL Policy training iterations.")
|
||||
parser.add_argument(
|
||||
"--distributed", action="store_true", default=False, help="Run training with multiple GPUs or nodes."
|
||||
)
|
||||
parser.add_argument("--export_io_descriptors", action="store_true", default=False, help="Export IO descriptors.")
|
||||
# append RSL-RL cli arguments
|
||||
cli_args.add_rsl_rl_args(parser)
|
||||
# append AppLauncher cli args
|
||||
AppLauncher.add_app_launcher_args(parser)
|
||||
args_cli, hydra_args = parser.parse_known_args()
|
||||
|
||||
# always enable cameras to record video
|
||||
if args_cli.video:
|
||||
args_cli.enable_cameras = True
|
||||
|
||||
# clear out sys.argv for Hydra
|
||||
sys.argv = [sys.argv[0]] + hydra_args
|
||||
|
||||
# launch omniverse app
|
||||
app_launcher = AppLauncher(args_cli)
|
||||
simulation_app = app_launcher.app
|
||||
|
||||
"""Check for minimum supported RSL-RL version."""
|
||||
|
||||
import importlib.metadata as metadata
|
||||
|
||||
from packaging import version
|
||||
|
||||
# check minimum supported rsl-rl version
|
||||
RSL_RL_VERSION = "3.0.1"
|
||||
installed_version = metadata.version("rsl-rl-lib")
|
||||
if version.parse(installed_version) < version.parse(RSL_RL_VERSION):
|
||||
cmd = [r"python", "-m", "pip", "install", f"rsl-rl-lib=={RSL_RL_VERSION}"]
|
||||
print(
|
||||
f"Please install the correct version of RSL-RL.\nExisting version is: '{installed_version}'"
|
||||
f" and required version is: '{RSL_RL_VERSION}'.\nTo install the correct version, run:"
|
||||
f"\n\n\t{' '.join(cmd)}\n"
|
||||
)
|
||||
exit(1)
|
||||
|
||||
"""Rest everything follows."""
|
||||
|
||||
import gymnasium as gym
|
||||
import torch
|
||||
from datetime import datetime
|
||||
|
||||
# local imports
|
||||
from utils import Logger
|
||||
|
||||
import omni
|
||||
from rsl_rl.runners import DistillationRunner, OnPolicyRunner, OnPolicyRunnerCTS
|
||||
|
||||
from isaaclab.envs import (
|
||||
DirectMARLEnv,
|
||||
DirectMARLEnvCfg,
|
||||
DirectRLEnvCfg,
|
||||
ManagerBasedRLEnvCfg,
|
||||
multi_agent_to_single_agent,
|
||||
)
|
||||
from isaaclab.utils.dict import print_dict
|
||||
from isaaclab.utils.io import dump_yaml
|
||||
from isaaclab_rl.rsl_rl import RslRlBaseRunnerCfg, RslRlVecEnvWrapper
|
||||
from isaaclab_tasks.utils import get_checkpoint_path
|
||||
from isaaclab_tasks.utils.hydra import hydra_task_config
|
||||
|
||||
import robot_lab.tasks # noqa: F401
|
||||
|
||||
torch.backends.cuda.matmul.allow_tf32 = True
|
||||
torch.backends.cudnn.allow_tf32 = True
|
||||
torch.backends.cudnn.deterministic = False
|
||||
torch.backends.cudnn.benchmark = False
|
||||
|
||||
|
||||
@hydra_task_config(args_cli.task, args_cli.agent)
|
||||
def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agent_cfg: RslRlBaseRunnerCfg):
|
||||
"""Train with RSL-RL agent."""
|
||||
# override configurations with non-hydra CLI arguments
|
||||
agent_cfg = cli_args.update_rsl_rl_cfg(agent_cfg, args_cli)
|
||||
env_cfg.scene.num_envs = args_cli.num_envs if args_cli.num_envs is not None else env_cfg.scene.num_envs
|
||||
agent_cfg.max_iterations = (
|
||||
args_cli.max_iterations if args_cli.max_iterations is not None else agent_cfg.max_iterations
|
||||
)
|
||||
|
||||
# set the environment seed
|
||||
# note: certain randomizations occur in the environment initialization so we set the seed here
|
||||
env_cfg.seed = agent_cfg.seed
|
||||
env_cfg.sim.device = args_cli.device if args_cli.device is not None else env_cfg.sim.device
|
||||
# check for invalid combination of CPU device with distributed training
|
||||
if args_cli.distributed and args_cli.device is not None and "cpu" in args_cli.device:
|
||||
raise ValueError(
|
||||
"Distributed training is not supported when using CPU device. "
|
||||
"Please use GPU device (e.g., --device cuda) for distributed training."
|
||||
)
|
||||
|
||||
# multi-gpu training configuration
|
||||
if args_cli.distributed:
|
||||
env_cfg.sim.device = f"cuda:{app_launcher.local_rank}"
|
||||
agent_cfg.device = f"cuda:{app_launcher.local_rank}"
|
||||
|
||||
# set seed to have diversity in different threads
|
||||
seed = agent_cfg.seed + app_launcher.local_rank
|
||||
env_cfg.seed = seed
|
||||
agent_cfg.seed = seed
|
||||
|
||||
# specify directory for logging experiments
|
||||
log_root_path = os.path.join("logs", "rsl_rl", agent_cfg.experiment_name)
|
||||
log_root_path = os.path.abspath(log_root_path)
|
||||
print(f"[INFO] Logging experiment in directory: {log_root_path}")
|
||||
# specify directory for logging runs: {time-stamp}_{run_name}
|
||||
log_dir = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
|
||||
# The Ray Tune workflow extracts experiment name using the logging line below, hence, do not change it (see PR #2346, comment-2819298849)
|
||||
print(f"Exact experiment name requested from command line: {log_dir}")
|
||||
if agent_cfg.run_name:
|
||||
log_dir += f"_{agent_cfg.run_name}"
|
||||
log_dir = os.path.join(log_root_path, log_dir)
|
||||
|
||||
# set the IO descriptors export flag if requested
|
||||
if isinstance(env_cfg, ManagerBasedRLEnvCfg):
|
||||
env_cfg.export_io_descriptors = args_cli.export_io_descriptors
|
||||
else:
|
||||
omni.log.warn(
|
||||
"IO descriptors are only supported for manager based RL environments. No IO descriptors will be exported."
|
||||
)
|
||||
|
||||
# set the log directory for the environment (works for all environment types)
|
||||
env_cfg.log_dir = log_dir
|
||||
|
||||
# create isaac environment
|
||||
env = gym.make(args_cli.task, cfg=env_cfg, render_mode="rgb_array" if args_cli.video else None)
|
||||
|
||||
# convert to single-agent instance if required by the RL algorithm
|
||||
if isinstance(env.unwrapped, DirectMARLEnv):
|
||||
env = multi_agent_to_single_agent(env)
|
||||
|
||||
# save resume path before creating a new log_dir
|
||||
if agent_cfg.resume or agent_cfg.algorithm.class_name == "Distillation":
|
||||
resume_path = get_checkpoint_path(log_root_path, agent_cfg.load_run, agent_cfg.load_checkpoint)
|
||||
|
||||
# wrap for video recording
|
||||
if args_cli.video:
|
||||
video_kwargs = {
|
||||
"video_folder": os.path.join(log_dir, "videos", "train"),
|
||||
"step_trigger": lambda step: step % args_cli.video_interval == 0,
|
||||
"video_length": args_cli.video_length,
|
||||
"disable_logger": True,
|
||||
}
|
||||
print("[INFO] Recording videos during training.")
|
||||
print_dict(video_kwargs, nesting=4)
|
||||
env = gym.wrappers.RecordVideo(env, **video_kwargs)
|
||||
|
||||
# wrap around environment for rsl-rl
|
||||
env = RslRlVecEnvWrapper(env, clip_actions=agent_cfg.clip_actions)
|
||||
|
||||
# create runner from rsl-rl
|
||||
if agent_cfg.class_name == "OnPolicyRunner":
|
||||
runner = OnPolicyRunner(env, agent_cfg.to_dict(), log_dir=log_dir, device=agent_cfg.device)
|
||||
elif agent_cfg.class_name == "DistillationRunner":
|
||||
runner = DistillationRunner(env, agent_cfg.to_dict(), log_dir=log_dir, device=agent_cfg.device)
|
||||
elif agent_cfg.class_name == "OnPolicyRunnerCTS":
|
||||
runner = OnPolicyRunnerCTS(env, agent_cfg.to_dict(), log_dir=log_dir, device=agent_cfg.device)
|
||||
else:
|
||||
raise ValueError(f"Unsupported runner class: {agent_cfg.class_name}")
|
||||
# write git state to logs
|
||||
runner.add_git_repo_to_log(__file__)
|
||||
# load the checkpoint
|
||||
if agent_cfg.resume or agent_cfg.algorithm.class_name == "Distillation":
|
||||
print(f"[INFO]: Loading model checkpoint from: {resume_path}")
|
||||
# load previously trained model
|
||||
runner.load(resume_path)
|
||||
|
||||
# dump the configuration into log-directory
|
||||
dump_yaml(os.path.join(log_dir, "params", "env.yaml"), env_cfg)
|
||||
dump_yaml(os.path.join(log_dir, "params", "agent.yaml"), agent_cfg)
|
||||
sys.stdout = Logger(os.path.join(log_dir, "train.log"))
|
||||
# run training
|
||||
runner.learn(num_learning_iterations=agent_cfg.max_iterations, init_at_random_ep_len=True)
|
||||
|
||||
# close the simulator
|
||||
env.close()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# run the main function
|
||||
main()
|
||||
# close sim app
|
||||
simulation_app.close()
|
||||
177
scripts/reinforcement_learning/rsl_rl/utils.py
Normal file
177
scripts/reinforcement_learning/rsl_rl/utils.py
Normal file
@@ -0,0 +1,177 @@
|
||||
# base version: IsaacLab/source/isaaclab_rl/isaaclab_rl/rsl_rl/exporter.py
|
||||
|
||||
import copy
|
||||
import os
|
||||
import torch
|
||||
import re
|
||||
import os
|
||||
import sys
|
||||
from typing import NamedTuple
|
||||
|
||||
"Script to log terminal output to a file, stripping ANSI escape codes."
|
||||
class Logger:
|
||||
def __init__(self, filename):
|
||||
self.terminal = sys.stdout
|
||||
os.makedirs(os.path.dirname(filename), exist_ok=True)
|
||||
self.log = open(filename, 'w', encoding='utf-8')
|
||||
|
||||
self.ansi_escape = re.compile(r'\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])')
|
||||
|
||||
def write(self, message):
|
||||
clean_message = self.ansi_escape.sub('', message)
|
||||
|
||||
self.terminal.write(message)
|
||||
self.log.write(clean_message)
|
||||
self.log.flush()
|
||||
|
||||
def flush(self):
|
||||
self.terminal.flush()
|
||||
self.log.flush()
|
||||
|
||||
# Inputs of CTS Policy is a TensorDict with 'policy' and 'single_obs' keys, we simulate this with a NamedTuple.
|
||||
class CTSPolicyInputs(NamedTuple):
|
||||
policy: torch.Tensor
|
||||
single_obs: torch.Tensor
|
||||
|
||||
def export_cts_policy_as_jit(policy: object, actor_obs_normalizer: object | None, single_obs_normalizer: object | None, path: str, filename="policy.pt"):
|
||||
"""Export CTS policy into a Torch JIT file.
|
||||
|
||||
Args:
|
||||
policy: The CTS policy torch module.
|
||||
actor_obs_normalizer: The empirical normalizer module for actor observations. If None, Identity is used.
|
||||
single_obs_normalizer: The empirical normalizer module for single observations. If None, Identity is used.
|
||||
path: The path to the saving directory.
|
||||
filename: The name of exported JIT file. Defaults to "policy.pt".
|
||||
"""
|
||||
policy_exporter = _TorchPolicyExporter(policy, actor_obs_normalizer, single_obs_normalizer)
|
||||
policy_exporter.export(path, filename)
|
||||
|
||||
|
||||
def export_cts_policy_as_onnx(
|
||||
policy: object, path: str, actor_obs_normalizer: object | None = None, single_obs_normalizer: object | None = None, filename="policy.onnx", verbose=False
|
||||
):
|
||||
"""Export CTS policy into a Torch ONNX file.
|
||||
|
||||
Args:
|
||||
policy: The CTS policy torch module.
|
||||
actor_obs_normalizer: The empirical normalizer module for actor observations. If None, Identity is used.
|
||||
single_obs_normalizer: The empirical normalizer module for single observations. If None, Identity is used.
|
||||
path: The path to the saving directory.
|
||||
filename: The name of exported ONNX file. Defaults to "policy.onnx".
|
||||
verbose: Whether to print the model summary. Defaults to False.
|
||||
"""
|
||||
if not os.path.exists(path):
|
||||
os.makedirs(path, exist_ok=True)
|
||||
policy_exporter = _OnnxPolicyExporter(policy, actor_obs_normalizer, single_obs_normalizer, verbose)
|
||||
policy_exporter.export(path, filename)
|
||||
|
||||
|
||||
"""
|
||||
Helper Classes - Private.
|
||||
"""
|
||||
|
||||
|
||||
class _TorchPolicyExporter(torch.nn.Module):
|
||||
"""Exporter of actor-critic into JIT file."""
|
||||
|
||||
def __init__(self, policy, actor_obs_normalizer=None, single_obs_normalizer=None):
|
||||
assert not policy.is_recurrent, "CTS policy should not be recurrent"
|
||||
super().__init__()
|
||||
|
||||
# copy policy parameters
|
||||
if hasattr(policy, "actor"):
|
||||
self.actor = copy.deepcopy(policy.actor)
|
||||
elif hasattr(policy, "student"):
|
||||
self.actor = copy.deepcopy(policy.student)
|
||||
else:
|
||||
raise ValueError("Policy does not have an actor/student module.")
|
||||
self.student_moe_encoder = copy.deepcopy(policy.student_moe_encoder)
|
||||
self.state_dependent_std = policy.state_dependent_std
|
||||
|
||||
# copy normalizer if exists
|
||||
if actor_obs_normalizer:
|
||||
self.actor_obs_normalizer = copy.deepcopy(actor_obs_normalizer)
|
||||
else:
|
||||
self.actor_obs_normalizer = torch.nn.Identity()
|
||||
if single_obs_normalizer:
|
||||
self.single_obs_normalizer = copy.deepcopy(single_obs_normalizer)
|
||||
else:
|
||||
self.single_obs_normalizer = torch.nn.Identity()
|
||||
|
||||
def forward(self, x: CTSPolicyInputs):
|
||||
single_obs = self.single_obs_normalizer(x.single_obs)
|
||||
obs_a = self.actor_obs_normalizer(x.policy)
|
||||
latent, _ = self.student_moe_encoder(obs_a)
|
||||
latent_and_obs = torch.cat([latent, single_obs], dim=-1)
|
||||
if self.state_dependent_std:
|
||||
return self.actor(latent_and_obs)[..., 0, :]
|
||||
else:
|
||||
return self.actor(latent_and_obs)
|
||||
|
||||
@torch.jit.export
|
||||
def reset(self):
|
||||
pass
|
||||
|
||||
def export(self, path, filename):
|
||||
os.makedirs(path, exist_ok=True)
|
||||
path = os.path.join(path, filename)
|
||||
self.to("cpu")
|
||||
traced_script_module = torch.jit.script(self)
|
||||
traced_script_module.save(path)
|
||||
|
||||
|
||||
class _OnnxPolicyExporter(torch.nn.Module):
|
||||
"""Exporter of actor-critic into ONNX file."""
|
||||
|
||||
def __init__(self, policy, actor_obs_normalizer=None, single_obs_normalizer=None, verbose=False):
|
||||
assert not policy.is_recurrent, "CTS policy should not be recurrent"
|
||||
super().__init__()
|
||||
self.verbose = verbose
|
||||
|
||||
# copy policy parameters
|
||||
if hasattr(policy, "actor"):
|
||||
self.actor = copy.deepcopy(policy.actor)
|
||||
elif hasattr(policy, "student"):
|
||||
self.actor = copy.deepcopy(policy.student)
|
||||
else:
|
||||
raise ValueError("Policy does not have an actor/student module.")
|
||||
self.student_moe_encoder = copy.deepcopy(policy.student_moe_encoder)
|
||||
self.num_single_obs = policy.num_single_obs
|
||||
self.num_actor_obs = policy.num_actor_obs
|
||||
self.state_dependent_std = policy.state_dependent_std
|
||||
|
||||
# copy normalizer if exists
|
||||
if actor_obs_normalizer:
|
||||
self.actor_obs_normalizer = copy.deepcopy(actor_obs_normalizer)
|
||||
else:
|
||||
self.actor_obs_normalizer = torch.nn.Identity()
|
||||
if single_obs_normalizer:
|
||||
self.single_obs_normalizer = copy.deepcopy(single_obs_normalizer)
|
||||
else:
|
||||
self.single_obs_normalizer = torch.nn.Identity()
|
||||
|
||||
def forward(self, history, single_obs):
|
||||
single_obs = self.single_obs_normalizer(single_obs)
|
||||
obs_a = self.actor_obs_normalizer(history)
|
||||
latent, _ = self.student_moe_encoder(obs_a)
|
||||
latent_and_obs = torch.cat([latent, single_obs], dim=-1)
|
||||
if self.state_dependent_std:
|
||||
return self.actor(latent_and_obs)[..., 0, :]
|
||||
else:
|
||||
return self.actor(latent_and_obs)
|
||||
|
||||
def export(self, path, filename):
|
||||
self.to("cpu")
|
||||
self.eval()
|
||||
opset_version = 18 # was 11, but it caused problems with linux-aarch, and 18 worked well across all systems.
|
||||
torch.onnx.export(
|
||||
self,
|
||||
(torch.zeros(1, self.num_actor_obs), torch.zeros(1, self.num_single_obs)),
|
||||
os.path.join(path, filename),
|
||||
export_params=True,
|
||||
opset_version=opset_version,
|
||||
verbose=self.verbose,
|
||||
input_names=["obs"],
|
||||
output_names=["actions"],
|
||||
dynamic_axes={},
|
||||
)
|
||||
62
scripts/tools/clean_trash.py
Normal file
62
scripts/tools/clean_trash.py
Normal file
@@ -0,0 +1,62 @@
|
||||
# Copyright (c) 2024-2025 Ziqi Fan
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import os
|
||||
import re
|
||||
import shutil
|
||||
|
||||
|
||||
def clean_trash(folder_path):
|
||||
"""
|
||||
Delete folders that meet the following conditions:
|
||||
1. Contain `events.out.*` files.
|
||||
2. Also meet:
|
||||
- No `.pt` files, or
|
||||
- Less than 3 `.pt` files.
|
||||
|
||||
:param folder_path: Target folder path
|
||||
"""
|
||||
if not os.path.exists(folder_path):
|
||||
print(f"Folder {folder_path} does not exist.")
|
||||
return
|
||||
|
||||
# List of folders to delete
|
||||
folders_to_delete = []
|
||||
|
||||
# Traverse subfolders
|
||||
for root, dirs, files in os.walk(folder_path):
|
||||
# Regex match `events.out.*` and `.pt` files
|
||||
event_pattern = re.compile(r"events\.out.*")
|
||||
model_pattern = re.compile(r"model_\d+\.pt")
|
||||
|
||||
event_files = [f for f in files if event_pattern.match(f)]
|
||||
model_files = [f for f in files if model_pattern.match(f)]
|
||||
|
||||
# Check delete conditions: have `events.out.*` files, but no `.pt` files or less than 3 `.pt` files
|
||||
if event_files and (len(model_files) < 3):
|
||||
folders_to_delete.append(os.path.abspath(root))
|
||||
|
||||
# If there are folders that meet the conditions, prompt and delete
|
||||
if folders_to_delete:
|
||||
print("The following folders contain `events.out.*` and meet the deletion conditions, they will be deleted:")
|
||||
for folder in folders_to_delete:
|
||||
print(f" - {folder}")
|
||||
|
||||
# Confirm deletion
|
||||
confirm = input("Confirm deletion of these folders? (y/n): ").strip().lower()
|
||||
if confirm == "y":
|
||||
for folder in folders_to_delete:
|
||||
shutil.rmtree(folder) # Delete the entire folder
|
||||
print(f"Deleted: {folder}")
|
||||
print("All folders that meet the conditions have been deleted.")
|
||||
else:
|
||||
print("Deletion operation canceled.")
|
||||
else:
|
||||
print("No folders meet the conditions, no need to delete.")
|
||||
|
||||
|
||||
# Example call
|
||||
if __name__ == "__main__":
|
||||
# folder = input("Please enter the target folder path: ").strip()
|
||||
folder = "logs"
|
||||
clean_trash(folder)
|
||||
142
scripts/tools/convert_mjcf.py
Normal file
142
scripts/tools/convert_mjcf.py
Normal file
@@ -0,0 +1,142 @@
|
||||
# Copyright (c) 2024-2025 Ziqi Fan
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
|
||||
"""
|
||||
Utility to convert a MJCF into USD format.
|
||||
|
||||
MuJoCo XML Format (MJCF) is an XML file format used in MuJoCo to describe all elements of a robot.
|
||||
For more information, see: http://www.mujoco.org/book/XMLreference.html
|
||||
|
||||
This script uses the MJCF importer extension from Isaac Sim (``isaacsim.asset.importer.mjcf``) to convert
|
||||
a MJCF asset into USD format. It is designed as a convenience script for command-line use. For more information
|
||||
on the MJCF importer, see the documentation for the extension:
|
||||
https://docs.isaacsim.omniverse.nvidia.com/latest/robot_setup/ext_isaacsim_asset_importer_mjcf.html
|
||||
|
||||
|
||||
positional arguments:
|
||||
input The path to the input URDF file.
|
||||
output The path to store the USD file.
|
||||
|
||||
optional arguments:
|
||||
-h, --help Show this help message and exit
|
||||
--fix-base Fix the base to where it is imported. (default: False)
|
||||
--import-sites Import sites by parse <site> tag. (default: True)
|
||||
--make-instanceable Make the asset instanceable for efficient cloning. (default: False)
|
||||
|
||||
"""
|
||||
|
||||
"""Launch Isaac Sim Simulator first."""
|
||||
|
||||
import argparse
|
||||
|
||||
from isaaclab.app import AppLauncher
|
||||
|
||||
# add argparse arguments
|
||||
parser = argparse.ArgumentParser(description="Utility to convert a MJCF into USD format.")
|
||||
parser.add_argument("input", type=str, help="The path to the input MJCF file.")
|
||||
parser.add_argument("output", type=str, help="The path to store the USD file.")
|
||||
parser.add_argument("--fix-base", action="store_true", default=False, help="Fix the base to where it is imported.")
|
||||
parser.add_argument(
|
||||
"--import-sites", action="store_true", default=False, help="Import sites by parsing the <site> tag."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--make-instanceable",
|
||||
action="store_true",
|
||||
default=False,
|
||||
help="Make the asset instanceable for efficient cloning.",
|
||||
)
|
||||
|
||||
# append AppLauncher cli args
|
||||
AppLauncher.add_app_launcher_args(parser)
|
||||
# parse the arguments
|
||||
args_cli = parser.parse_args()
|
||||
|
||||
# launch omniverse app
|
||||
app_launcher = AppLauncher(args_cli)
|
||||
simulation_app = app_launcher.app
|
||||
|
||||
"""Rest everything follows."""
|
||||
|
||||
import contextlib
|
||||
import os
|
||||
|
||||
import carb
|
||||
import isaacsim.core.utils.stage as stage_utils
|
||||
import omni.kit.app
|
||||
|
||||
from isaaclab.sim.converters import MjcfConverter, MjcfConverterCfg
|
||||
from isaaclab.utils.assets import check_file_path
|
||||
from isaaclab.utils.dict import print_dict
|
||||
|
||||
|
||||
def main():
|
||||
# check valid file path
|
||||
mjcf_path = args_cli.input
|
||||
if not os.path.isabs(mjcf_path):
|
||||
mjcf_path = os.path.abspath(mjcf_path)
|
||||
if not check_file_path(mjcf_path):
|
||||
raise ValueError(f"Invalid file path: {mjcf_path}")
|
||||
# create destination path
|
||||
dest_path = args_cli.output
|
||||
if not os.path.isabs(dest_path):
|
||||
dest_path = os.path.abspath(dest_path)
|
||||
|
||||
# create the converter configuration
|
||||
mjcf_converter_cfg = MjcfConverterCfg(
|
||||
asset_path=mjcf_path,
|
||||
usd_dir=os.path.dirname(dest_path),
|
||||
usd_file_name=os.path.basename(dest_path),
|
||||
fix_base=args_cli.fix_base,
|
||||
import_sites=args_cli.import_sites,
|
||||
force_usd_conversion=True,
|
||||
make_instanceable=args_cli.make_instanceable,
|
||||
)
|
||||
|
||||
# Print info
|
||||
print("-" * 80)
|
||||
print("-" * 80)
|
||||
print(f"Input MJCF file: {mjcf_path}")
|
||||
print("MJCF importer config:")
|
||||
print_dict(mjcf_converter_cfg.to_dict(), nesting=0)
|
||||
print("-" * 80)
|
||||
print("-" * 80)
|
||||
|
||||
# Create mjcf converter and import the file
|
||||
mjcf_converter = MjcfConverter(mjcf_converter_cfg)
|
||||
# print output
|
||||
print("MJCF importer output:")
|
||||
print(f"Generated USD file: {mjcf_converter.usd_path}")
|
||||
print("-" * 80)
|
||||
print("-" * 80)
|
||||
|
||||
# Determine if there is a GUI to update:
|
||||
# acquire settings interface
|
||||
carb_settings_iface = carb.settings.get_settings()
|
||||
# read flag for whether a local GUI is enabled
|
||||
local_gui = carb_settings_iface.get("/app/window/enabled")
|
||||
# read flag for whether livestreaming GUI is enabled
|
||||
livestream_gui = carb_settings_iface.get("/app/livestream/enabled")
|
||||
|
||||
# Simulate scene (if not headless)
|
||||
if local_gui or livestream_gui:
|
||||
# Open the stage with USD
|
||||
stage_utils.open_stage(mjcf_converter.usd_path)
|
||||
# Reinitialize the simulation
|
||||
app = omni.kit.app.get_app_interface()
|
||||
# Run simulation
|
||||
with contextlib.suppress(KeyboardInterrupt):
|
||||
while app.is_running():
|
||||
# perform step
|
||||
app.update()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# run the main function
|
||||
main()
|
||||
# close sim app
|
||||
simulation_app.close()
|
||||
166
scripts/tools/convert_urdf.py
Normal file
166
scripts/tools/convert_urdf.py
Normal file
@@ -0,0 +1,166 @@
|
||||
# Copyright (c) 2024-2025 Ziqi Fan
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
|
||||
"""
|
||||
Utility to convert a URDF into USD format.
|
||||
|
||||
Unified Robot Description Format (URDF) is an XML file format used in ROS to describe all elements of
|
||||
a robot. For more information, see: http://wiki.ros.org/urdf
|
||||
|
||||
This script uses the URDF importer extension from Isaac Sim (``isaacsim.asset.importer.urdf``) to convert a
|
||||
URDF asset into USD format. It is designed as a convenience script for command-line use. For more
|
||||
information on the URDF importer, see the documentation for the extension:
|
||||
https://docs.isaacsim.omniverse.nvidia.com/latest/robot_setup/ext_isaacsim_asset_importer_urdf.html
|
||||
|
||||
|
||||
positional arguments:
|
||||
input The path to the input URDF file.
|
||||
output The path to store the USD file.
|
||||
|
||||
optional arguments:
|
||||
-h, --help Show this help message and exit
|
||||
--merge-joints Consolidate links that are connected by fixed joints. (default: False)
|
||||
--fix-base Fix the base to where it is imported. (default: False)
|
||||
--joint-stiffness The stiffness of the joint drive. (default: 100.0)
|
||||
--joint-damping The damping of the joint drive. (default: 1.0)
|
||||
--joint-target-type The type of control to use for the joint drive. (default: "position")
|
||||
|
||||
"""
|
||||
|
||||
"""Launch Isaac Sim Simulator first."""
|
||||
|
||||
import argparse
|
||||
|
||||
from isaaclab.app import AppLauncher
|
||||
|
||||
# add argparse arguments
|
||||
parser = argparse.ArgumentParser(description="Utility to convert a URDF into USD format.")
|
||||
parser.add_argument("input", type=str, help="The path to the input URDF file.")
|
||||
parser.add_argument("output", type=str, help="The path to store the USD file.")
|
||||
parser.add_argument(
|
||||
"--merge-joints",
|
||||
action="store_true",
|
||||
default=False,
|
||||
help="Consolidate links that are connected by fixed joints.",
|
||||
)
|
||||
parser.add_argument("--fix-base", action="store_true", default=False, help="Fix the base to where it is imported.")
|
||||
parser.add_argument(
|
||||
"--joint-stiffness",
|
||||
type=float,
|
||||
default=100.0,
|
||||
help="The stiffness of the joint drive.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--joint-damping",
|
||||
type=float,
|
||||
default=1.0,
|
||||
help="The damping of the joint drive.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--joint-target-type",
|
||||
type=str,
|
||||
default="position",
|
||||
choices=["position", "velocity", "none"],
|
||||
help="The type of control to use for the joint drive.",
|
||||
)
|
||||
|
||||
# append AppLauncher cli args
|
||||
AppLauncher.add_app_launcher_args(parser)
|
||||
# parse the arguments
|
||||
args_cli = parser.parse_args()
|
||||
|
||||
# launch omniverse app
|
||||
app_launcher = AppLauncher(args_cli)
|
||||
simulation_app = app_launcher.app
|
||||
|
||||
"""Rest everything follows."""
|
||||
|
||||
import contextlib
|
||||
import os
|
||||
|
||||
import carb
|
||||
import isaacsim.core.utils.stage as stage_utils
|
||||
import omni.kit.app
|
||||
|
||||
from isaaclab.sim.converters import UrdfConverter, UrdfConverterCfg
|
||||
from isaaclab.utils.assets import check_file_path
|
||||
from isaaclab.utils.dict import print_dict
|
||||
|
||||
|
||||
def main():
|
||||
# check valid file path
|
||||
urdf_path = args_cli.input
|
||||
if not os.path.isabs(urdf_path):
|
||||
urdf_path = os.path.abspath(urdf_path)
|
||||
if not check_file_path(urdf_path):
|
||||
raise ValueError(f"Invalid file path: {urdf_path}")
|
||||
# create destination path
|
||||
dest_path = args_cli.output
|
||||
if not os.path.isabs(dest_path):
|
||||
dest_path = os.path.abspath(dest_path)
|
||||
|
||||
# Create Urdf converter config
|
||||
urdf_converter_cfg = UrdfConverterCfg(
|
||||
asset_path=urdf_path,
|
||||
usd_dir=os.path.dirname(dest_path),
|
||||
usd_file_name=os.path.basename(dest_path),
|
||||
fix_base=args_cli.fix_base,
|
||||
merge_fixed_joints=args_cli.merge_joints,
|
||||
force_usd_conversion=True,
|
||||
joint_drive=UrdfConverterCfg.JointDriveCfg(
|
||||
gains=UrdfConverterCfg.JointDriveCfg.PDGainsCfg(
|
||||
stiffness=args_cli.joint_stiffness,
|
||||
damping=args_cli.joint_damping,
|
||||
),
|
||||
target_type=args_cli.joint_target_type,
|
||||
),
|
||||
)
|
||||
|
||||
# Print info
|
||||
print("-" * 80)
|
||||
print("-" * 80)
|
||||
print(f"Input URDF file: {urdf_path}")
|
||||
print("URDF importer config:")
|
||||
print_dict(urdf_converter_cfg.to_dict(), nesting=0)
|
||||
print("-" * 80)
|
||||
print("-" * 80)
|
||||
|
||||
# Create Urdf converter and import the file
|
||||
urdf_converter = UrdfConverter(urdf_converter_cfg)
|
||||
# print output
|
||||
print("URDF importer output:")
|
||||
print(f"Generated USD file: {urdf_converter.usd_path}")
|
||||
print("-" * 80)
|
||||
print("-" * 80)
|
||||
|
||||
# Determine if there is a GUI to update:
|
||||
# acquire settings interface
|
||||
carb_settings_iface = carb.settings.get_settings()
|
||||
# read flag for whether a local GUI is enabled
|
||||
local_gui = carb_settings_iface.get("/app/window/enabled")
|
||||
# read flag for whether livestreaming GUI is enabled
|
||||
livestream_gui = carb_settings_iface.get("/app/livestream/enabled")
|
||||
|
||||
# Simulate scene (if not headless)
|
||||
if local_gui or livestream_gui:
|
||||
# Open the stage with USD
|
||||
stage_utils.open_stage(urdf_converter.usd_path)
|
||||
# Reinitialize the simulation
|
||||
app = omni.kit.app.get_app_interface()
|
||||
# Run simulation
|
||||
with contextlib.suppress(KeyboardInterrupt):
|
||||
while app.is_running():
|
||||
# perform step
|
||||
app.update()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# run the main function
|
||||
main()
|
||||
# close sim app
|
||||
simulation_app.close()
|
||||
77
scripts/tools/list_envs.py
Normal file
77
scripts/tools/list_envs.py
Normal file
@@ -0,0 +1,77 @@
|
||||
# Copyright (c) 2024-2025 Ziqi Fan
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# Copyright (c) 2024-2025, The Isaac Lab Project Developers.
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
"""
|
||||
Script to print all the available environments in Isaac Lab.
|
||||
|
||||
The script iterates over all registered environments and stores the details in a table.
|
||||
It prints the name of the environment, the entry point and the config file.
|
||||
|
||||
All the environments are registered in the `isaaclab_tasks` extension. They start
|
||||
with `Isaac` in their name.
|
||||
"""
|
||||
|
||||
"""Launch Isaac Sim Simulator first."""
|
||||
|
||||
from isaaclab.app import AppLauncher
|
||||
|
||||
# launch omniverse app
|
||||
app_launcher = AppLauncher(headless=True)
|
||||
simulation_app = app_launcher.app
|
||||
|
||||
|
||||
"""Rest everything follows."""
|
||||
|
||||
import gymnasium as gym
|
||||
import textwrap
|
||||
from prettytable import PrettyTable
|
||||
|
||||
import robot_lab.tasks # noqa: F401
|
||||
|
||||
|
||||
def main():
|
||||
"""Print all environments registered in `isaaclab_tasks` extension."""
|
||||
# print all the available environments
|
||||
table = PrettyTable(["S. No.", "Task Name", "Entry Point", "Config"])
|
||||
table.title = "Available Environments in Isaac Lab"
|
||||
# set alignment of table columns
|
||||
table.align["Task Name"] = "l"
|
||||
table.align["Entry Point"] = "l"
|
||||
table.align["Config"] = "l"
|
||||
table.hrules = 1
|
||||
|
||||
# set max width for text wrapping
|
||||
max_width = 50
|
||||
|
||||
# count of environments
|
||||
index = 0
|
||||
# acquire all Isaac environments names
|
||||
for task_spec in gym.registry.values():
|
||||
if "RobotLab" in task_spec.id:
|
||||
# wrap long text in each column before adding it to the table
|
||||
task_name = textwrap.fill(task_spec.id, max_width)
|
||||
entry_point = textwrap.fill(task_spec.entry_point, max_width)
|
||||
config = textwrap.fill(task_spec.kwargs["env_cfg_entry_point"], max_width)
|
||||
|
||||
# add details to table
|
||||
table.add_row([index + 1, task_name, entry_point, config])
|
||||
# increment count
|
||||
index += 1
|
||||
|
||||
print(table)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
try:
|
||||
# run the main function
|
||||
main()
|
||||
except Exception as e:
|
||||
raise e
|
||||
finally:
|
||||
# close the app
|
||||
simulation_app.close()
|
||||
35
source/robot_lab/config/extension.toml
Normal file
35
source/robot_lab/config/extension.toml
Normal file
@@ -0,0 +1,35 @@
|
||||
[package]
|
||||
|
||||
# Semantic Versioning is used: https://semver.org/
|
||||
version = "2.3.0"
|
||||
|
||||
# Description
|
||||
category = "isaaclab"
|
||||
readme = "README.md"
|
||||
|
||||
title = "RL Extension Library for Robots, Based on IsaacLab."
|
||||
author = "Ziqi Fan"
|
||||
maintainer = "Ziqi Fan"
|
||||
description="RL Extension Library for Robots, Based on IsaacLab."
|
||||
repository = "https://github.com/fan-ziqi/robot_lab.git"
|
||||
keywords = ["extension", "robot_lab", "isaaclab"]
|
||||
|
||||
[dependencies]
|
||||
"isaaclab" = {}
|
||||
"isaaclab_assets" = {}
|
||||
"isaaclab_mimic" = {}
|
||||
"isaaclab_rl" = {}
|
||||
"isaaclab_tasks" = {}
|
||||
# NOTE: Add additional dependencies here
|
||||
|
||||
[[python.module]]
|
||||
name = "robot_lab"
|
||||
|
||||
[isaaclab_settings]
|
||||
# TODO: Uncomment and list any apt dependencies here.
|
||||
# If none, leave it commented out.
|
||||
# apt_deps = ["example_package"]
|
||||
# TODO: Uncomment and provide path to a ros_ws
|
||||
# with rosdeps to be installed. If none,
|
||||
# leave it commented out.
|
||||
# ros_ws = "path/from/extension_root/to/ros_ws"
|
||||
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
@@ -0,0 +1,761 @@
|
||||
<?xml version='1.0' encoding='utf-8'?>
|
||||
<robot name="go2_description">
|
||||
<link name="base">
|
||||
<inertial>
|
||||
<origin xyz="0.021112 0 -0.005366" rpy="0 0 0" />
|
||||
<mass value="6.921" />
|
||||
<inertia ixx="0.02448" ixy="0.00012166" ixz="0.0014849" iyy="0.098077" iyz="-3.12E-05" izz="0.107" />
|
||||
</inertial>
|
||||
<visual>
|
||||
<origin xyz="0 0 0" rpy="0 0 0" />
|
||||
<geometry>
|
||||
<mesh filename="package://go2_description/meshes/base.dae" />
|
||||
</geometry>
|
||||
<material name="深色橡胶_001-effect"><color rgba="0.6717055 0.6924257 0.7742702 1" /></material><material name="白色logo_001-effect"><color rgba="1 1 1 1" /></material><material name="黑色贴纸_007-effect"><color rgba="0.009018 0.009018 0.009018 1" /></material><material name="黑色金属_007-effect"><color rgba="0.01341 0.01341 0.01341 1" /></material><material name="黑色塑料_003-effect"><color rgba="0.023455 0.023455 0.023455 1" /></material></visual>
|
||||
<collision>
|
||||
<origin rpy="0 0 0" xyz="0 0 0" />
|
||||
<geometry>
|
||||
<box size="0.3762 0.0935 0.114" />
|
||||
</geometry>
|
||||
</collision>
|
||||
</link>
|
||||
<link name="Head_upper">
|
||||
<inertial>
|
||||
<origin xyz="0 0 0" rpy="0 0 0" />
|
||||
<mass value="0.001" />
|
||||
<inertia ixx="9.6e-06" ixy="0" ixz="0" iyy="9.6e-06" iyz="0" izz="9.6e-06" />
|
||||
</inertial>
|
||||
<collision>
|
||||
<origin xyz="0 0 0" rpy="0 0 0" />
|
||||
<geometry>
|
||||
<cylinder radius="0.05" length="0.09" />
|
||||
</geometry>
|
||||
</collision>
|
||||
</link>
|
||||
<joint name="Head_upper_joint" type="fixed" dont_collapse="true">
|
||||
<origin xyz="0.285 0 0.01" rpy="0 0 0" />
|
||||
<parent link="base" />
|
||||
<child link="Head_upper" />
|
||||
<axis xyz="0 0 0" />
|
||||
</joint>
|
||||
<link name="Head_lower">
|
||||
<inertial>
|
||||
<origin xyz="0 0 0" rpy="0 0 0" />
|
||||
<mass value="0.001" />
|
||||
<inertia ixx="9.6e-06" ixy="0" ixz="0" iyy="9.6e-06" iyz="0" izz="9.6e-06" />
|
||||
</inertial>
|
||||
<collision>
|
||||
<origin xyz="0 0 0" rpy="0 0 0" />
|
||||
<geometry>
|
||||
<sphere radius="0.047" />
|
||||
</geometry>
|
||||
</collision>
|
||||
</link>
|
||||
<joint name="Head_lower_joint" type="fixed" dont_collapse="true">
|
||||
<origin xyz="0.008 0 -0.07" rpy="0 0 0" />
|
||||
<parent link="Head_upper" />
|
||||
<child link="Head_lower" />
|
||||
<axis xyz="0 0 0" />
|
||||
</joint>
|
||||
<link name="FL_hip">
|
||||
<inertial>
|
||||
<origin xyz="-0.0054 0.00194 -0.000105" rpy="0 0 0" />
|
||||
<mass value="0.678" />
|
||||
<inertia ixx="0.00048" ixy="-3.01E-06" ixz="1.11E-06" iyy="0.000884" iyz="-1.42E-06" izz="0.000596" />
|
||||
</inertial>
|
||||
<visual>
|
||||
<origin xyz="0 0 0" rpy="0 0 0" />
|
||||
<geometry>
|
||||
<mesh filename="package://go2_description/meshes/hip.dae" />
|
||||
</geometry>
|
||||
<material name="深色橡胶_001-effect"><color rgba="0.6717055 0.6924257 0.7742702 1" /></material><material name="Material-effect"><color rgba="0.9 0.95 0.95 1" /></material></visual>
|
||||
<collision>
|
||||
<origin rpy="1.5707963267948966 0 0" xyz="0 0.08 0" />
|
||||
<geometry>
|
||||
<cylinder length="0.04" radius="0.046" />
|
||||
</geometry>
|
||||
</collision>
|
||||
</link>
|
||||
<joint name="FL_hip_joint" type="revolute">
|
||||
<origin xyz="0.1934 0.0465 0" rpy="0 0 0" />
|
||||
<parent link="base" />
|
||||
<child link="FL_hip" />
|
||||
<axis xyz="1 0 0" />
|
||||
<limit lower="-1.0472" upper="1.0472" effort="23.7" velocity="30.1" />
|
||||
</joint>
|
||||
<link name="FL_hip_rotor">
|
||||
<inertial>
|
||||
<origin rpy="0 0 0" xyz="0.0 0.0 0.0" />
|
||||
<mass value="0.089" />
|
||||
<inertia ixx="0.000111842" ixy="0.0" ixz="0.0" iyy="0.000059647" iyz="0.0" izz="0.000059647" />
|
||||
</inertial>
|
||||
</link>
|
||||
<joint name="FL_hip_rotor_joint" type="fixed">
|
||||
<origin rpy="0 0 0" xyz="0.11215 0.04675 0" />
|
||||
<parent link="base" />
|
||||
<child link="FL_hip_rotor" />
|
||||
</joint>
|
||||
<link name="FL_thigh">
|
||||
<inertial>
|
||||
<origin xyz="-0.00374 -0.0223 -0.0327" rpy="0 0 0" />
|
||||
<mass value="1.152" />
|
||||
<inertia ixx="0.00584" ixy="8.72E-05" ixz="-0.000289" iyy="0.0058" iyz="0.000808" izz="0.00103" />
|
||||
</inertial>
|
||||
<visual>
|
||||
<origin xyz="0 0 0" rpy="0 0 0" />
|
||||
<geometry>
|
||||
<mesh filename="package://go2_description/meshes/thigh.dae" />
|
||||
</geometry>
|
||||
<material name="深色橡胶_003-effect"><color rgba="0.6717055 0.6924257 0.7742702 1" /></material><material name="Material_009-effect"><color rgba="0.9 0.95 0.95 1" /></material></visual>
|
||||
<collision>
|
||||
<origin rpy="0 1.5707963267948966 0" xyz="0 0 -0.1065" />
|
||||
<geometry>
|
||||
<box size="0.213 0.0245 0.034" />
|
||||
</geometry>
|
||||
</collision>
|
||||
</link>
|
||||
<joint name="FL_thigh_joint" type="revolute">
|
||||
<origin xyz="0 0.0955 0" rpy="0 0 0" />
|
||||
<parent link="FL_hip" />
|
||||
<child link="FL_thigh" />
|
||||
<axis xyz="0 1 0" />
|
||||
<limit lower="-1.5708" upper="3.4907" effort="23.7" velocity="30.1" />
|
||||
</joint>
|
||||
<link name="FL_thigh_rotor">
|
||||
<inertial>
|
||||
<origin rpy="0 0 0" xyz="0.0 0.0 0.0" />
|
||||
<mass value="0.089" />
|
||||
<inertia ixx="0.000059647" ixy="0.0" ixz="0.0" iyy="0.000111842" iyz="0.0" izz="0.000059647" />
|
||||
</inertial>
|
||||
</link>
|
||||
<joint name="FL_thigh_rotor_joint" type="fixed">
|
||||
<origin rpy="0 0 0" xyz="0 -0.00015 0" />
|
||||
<parent link="FL_hip" />
|
||||
<child link="FL_thigh_rotor" />
|
||||
</joint>
|
||||
<link name="FL_calf">
|
||||
<inertial>
|
||||
<origin xyz="0.00548 -0.000975 -0.115" rpy="0 0 0" />
|
||||
<mass value="0.154" />
|
||||
<inertia ixx="0.00108" ixy="3.4E-07" ixz="1.72E-05" iyy="0.0011" iyz="8.28E-06" izz="3.29E-05" />
|
||||
</inertial>
|
||||
<visual>
|
||||
<origin xyz="0 0 0" rpy="0 0 0" />
|
||||
<geometry>
|
||||
<mesh filename="package://go2_description/meshes/calf.dae" />
|
||||
</geometry>
|
||||
<material name="黑色足端_005-effect"><color rgba="0 0 0 1" /></material><material name="深色橡胶_005-effect"><color rgba="0.6717055 0.6924257 0.7742702 1" /></material></visual>
|
||||
<collision>
|
||||
<origin rpy="0 -0.21 0" xyz="0.008 0 -0.06" />
|
||||
<geometry>
|
||||
<cylinder length="0.12" radius="0.012" />
|
||||
</geometry>
|
||||
</collision>
|
||||
</link>
|
||||
<joint name="FL_calf_joint" type="revolute">
|
||||
<origin xyz="0 0 -0.213" rpy="0 0 0" />
|
||||
<parent link="FL_thigh" />
|
||||
<child link="FL_calf" />
|
||||
<axis xyz="0 1 0" />
|
||||
<limit lower="-2.7227" upper="-0.83776" effort="45.43" velocity="15.70" />
|
||||
</joint>
|
||||
<link name="FL_calflower">
|
||||
<collision>
|
||||
<origin rpy="0 0 0" xyz="0 0 0" />
|
||||
<geometry>
|
||||
<cylinder length="0.065" radius="0.011" />
|
||||
</geometry>
|
||||
</collision>
|
||||
</link>
|
||||
<joint name="FL_calflower_joint" type="fixed">
|
||||
<origin xyz="0.020 0 -0.148" rpy="0 0.05 0" />
|
||||
<parent link="FL_calf" />
|
||||
<child link="FL_calflower" />
|
||||
<axis xyz="0 0 0" />
|
||||
</joint>
|
||||
<link name="FL_calflower1">
|
||||
<collision>
|
||||
<origin rpy="0 0 0" xyz="0 0 0" />
|
||||
<geometry>
|
||||
<cylinder length="0.03" radius="0.0155" />
|
||||
</geometry>
|
||||
</collision>
|
||||
</link>
|
||||
<joint name="FL_calflower1_joint" type="fixed">
|
||||
<origin xyz="-0.01 0 -0.04" rpy="0 0.48 0" />
|
||||
<parent link="FL_calflower" />
|
||||
<child link="FL_calflower1" />
|
||||
<axis xyz="0 0 0" />
|
||||
</joint>
|
||||
<link name="FL_calf_rotor">
|
||||
<inertial>
|
||||
<origin rpy="0 0 0" xyz="0.0 0.0 0.0" />
|
||||
<mass value="0.089" />
|
||||
<inertia ixx="0.000059647" ixy="0.0" ixz="0.0" iyy="0.000111842" iyz="0.0" izz="0.000059647" />
|
||||
</inertial>
|
||||
</link>
|
||||
<joint name="FL_calf_rotor_joint" type="fixed">
|
||||
<origin rpy="0 0 0" xyz="0 -0.03235 0" />
|
||||
<parent link="FL_thigh" />
|
||||
<child link="FL_calf_rotor" />
|
||||
</joint>
|
||||
<link name="FL_foot">
|
||||
<inertial>
|
||||
<origin xyz="0 0 0" rpy="0 0 0" />
|
||||
<mass value="0.04" />
|
||||
<inertia ixx="9.6e-06" ixy="0" ixz="0" iyy="9.6e-06" iyz="0" izz="9.6e-06" />
|
||||
</inertial>
|
||||
<visual>
|
||||
<origin xyz="0 0 0" rpy="0 0 0" />
|
||||
<geometry>
|
||||
<mesh filename="package://go2_description/meshes/foot.dae" />
|
||||
</geometry>
|
||||
<material name="黑色足端_007-effect"><color rgba="0 0 0 1" /></material></visual>
|
||||
<collision>
|
||||
<origin rpy="0 0 0" xyz="-0.002 0 0" />
|
||||
<geometry>
|
||||
<sphere radius="0.022" />
|
||||
</geometry>
|
||||
</collision>
|
||||
</link>
|
||||
<joint name="FL_foot_joint" type="fixed" dont_collapse="true">
|
||||
<origin xyz="0 0 -0.213" rpy="0 0 0" />
|
||||
<parent link="FL_calf" />
|
||||
<child link="FL_foot" />
|
||||
<axis xyz="0 0 0" />
|
||||
</joint>
|
||||
<link name="FR_hip">
|
||||
<inertial>
|
||||
<origin xyz="-0.0054 -0.00194 -0.000105" rpy="0 0 0" />
|
||||
<mass value="0.678" />
|
||||
<inertia ixx="0.00048" ixy="3.01E-06" ixz="1.11E-06" iyy="0.000884" iyz="1.42E-06" izz="0.000596" />
|
||||
</inertial>
|
||||
<visual>
|
||||
<origin xyz="0 0 0" rpy="3.1415 0 0" />
|
||||
<geometry>
|
||||
<mesh filename="package://go2_description/meshes/hip.dae" />
|
||||
</geometry>
|
||||
<material name="深色橡胶_001-effect"><color rgba="0.6717055 0.6924257 0.7742702 1" /></material><material name="Material-effect"><color rgba="0.9 0.95 0.95 1" /></material></visual>
|
||||
<collision>
|
||||
<origin rpy="1.5707963267948966 0 0" xyz="0 -0.08 0" />
|
||||
<geometry>
|
||||
<cylinder length="0.04" radius="0.046" />
|
||||
</geometry>
|
||||
</collision>
|
||||
</link>
|
||||
<joint name="FR_hip_joint" type="revolute">
|
||||
<origin xyz="0.1934 -0.0465 0" rpy="0 0 0" />
|
||||
<parent link="base" />
|
||||
<child link="FR_hip" />
|
||||
<axis xyz="1 0 0" />
|
||||
<limit lower="-1.0472" upper="1.0472" effort="23.7" velocity="30.1" />
|
||||
</joint>
|
||||
<link name="FR_hip_rotor">
|
||||
<inertial>
|
||||
<origin rpy="0 0 0" xyz="0.0 0.0 0.0" />
|
||||
<mass value="0.089" />
|
||||
<inertia ixx="0.000111842" ixy="0.0" ixz="0.0" iyy="0.000059647" iyz="0.0" izz="0.000059647" />
|
||||
</inertial>
|
||||
</link>
|
||||
<joint name="FR_hip_rotor_joint" type="fixed">
|
||||
<origin rpy="0 0 0" xyz="0.11215 -0.04675 0" />
|
||||
<parent link="base" />
|
||||
<child link="FR_hip_rotor" />
|
||||
</joint>
|
||||
<link name="FR_thigh">
|
||||
<inertial>
|
||||
<origin xyz="-0.00374 0.0223 -0.0327" rpy="0 0 0" />
|
||||
<mass value="1.152" />
|
||||
<inertia ixx="0.00584" ixy="-8.72E-05" ixz="-0.000289" iyy="0.0058" iyz="-0.000808" izz="0.00103" />
|
||||
</inertial>
|
||||
<visual>
|
||||
<origin xyz="0 0 0" rpy="0 0 0" />
|
||||
<geometry>
|
||||
<mesh filename="package://go2_description/meshes/thigh_mirror.dae" />
|
||||
</geometry>
|
||||
<material name="深色橡胶_012-effect"><color rgba="0.6717055 0.6924257 0.7742702 1" /></material><material name="Material_011-effect"><color rgba="0.9 0.95 0.95 1" /></material></visual>
|
||||
<collision>
|
||||
<origin rpy="0 1.5707963267948966 0" xyz="0 0 -0.1065" />
|
||||
<geometry>
|
||||
<box size="0.213 0.0245 0.034" />
|
||||
</geometry>
|
||||
</collision>
|
||||
</link>
|
||||
<joint name="FR_thigh_joint" type="revolute">
|
||||
<origin xyz="0 -0.0955 0" rpy="0 0 0" />
|
||||
<parent link="FR_hip" />
|
||||
<child link="FR_thigh" />
|
||||
<axis xyz="0 1 0" />
|
||||
<limit lower="-1.5708" upper="3.4907" effort="23.7" velocity="30.1" />
|
||||
</joint>
|
||||
<link name="FR_thigh_rotor">
|
||||
<inertial>
|
||||
<origin rpy="0 0 0" xyz="0.0 0.0 0.0" />
|
||||
<mass value="0.089" />
|
||||
<inertia ixx="0.000059647" ixy="0.0" ixz="0.0" iyy="0.000111842" iyz="0.0" izz="0.000059647" />
|
||||
</inertial>
|
||||
</link>
|
||||
<joint name="FR_thigh_rotor_joint" type="fixed">
|
||||
<origin rpy="0 0 0" xyz="0 0.00015 0" />
|
||||
<parent link="FR_hip" />
|
||||
<child link="FR_thigh_rotor" />
|
||||
</joint>
|
||||
<link name="FR_calf">
|
||||
<inertial>
|
||||
<origin xyz="0.00548 0.000975 -0.115" rpy="0 0 0" />
|
||||
<mass value="0.154" />
|
||||
<inertia ixx="0.00108" ixy="-3.4E-07" ixz="1.72E-05" iyy="0.0011" iyz="-8.28E-06" izz="3.29E-05" />
|
||||
</inertial>
|
||||
<visual>
|
||||
<origin xyz="0 0 0" rpy="0 0 0" />
|
||||
<geometry>
|
||||
<mesh filename="package://go2_description/meshes/calf_mirror.dae" />
|
||||
</geometry>
|
||||
<material name="黑色足端_006-effect"><color rgba="0 0 0 1" /></material><material name="深色橡胶_007-effect"><color rgba="0.6717055 0.6924257 0.7742702 1" /></material></visual>
|
||||
<collision>
|
||||
<origin rpy="0 -0.2 0" xyz="0.01 0 -0.06" />
|
||||
<geometry>
|
||||
<cylinder length="0.12" radius="0.013" />
|
||||
</geometry>
|
||||
</collision>
|
||||
</link>
|
||||
<joint name="FR_calf_joint" type="revolute">
|
||||
<origin xyz="0 0 -0.213" rpy="0 0 0" />
|
||||
<parent link="FR_thigh" />
|
||||
<child link="FR_calf" />
|
||||
<axis xyz="0 1 0" />
|
||||
<limit lower="-2.7227" upper="-0.83776" effort="45.43" velocity="15.70" />
|
||||
</joint>
|
||||
<link name="FR_calflower">
|
||||
<collision>
|
||||
<origin rpy="0 0 0" xyz="0 0 0" />
|
||||
<geometry>
|
||||
<cylinder length="0.065" radius="0.011" />
|
||||
</geometry>
|
||||
</collision>
|
||||
</link>
|
||||
<joint name="FR_calflower_joint" type="fixed">
|
||||
<origin xyz="0.020 0 -0.148" rpy="0 0.05 0" />
|
||||
<parent link="FR_calf" />
|
||||
<child link="FR_calflower" />
|
||||
<axis xyz="0 0 0" />
|
||||
</joint>
|
||||
<link name="FR_calflower1">
|
||||
<collision>
|
||||
<origin rpy="0 0 0" xyz="0 0 0" />
|
||||
<geometry>
|
||||
<cylinder length="0.03" radius="0.0155" />
|
||||
</geometry>
|
||||
</collision>
|
||||
</link>
|
||||
<joint name="FR_calflower1_joint" type="fixed">
|
||||
<origin xyz="-0.01 0 -0.04" rpy="0 0.48 0" />
|
||||
<parent link="FR_calflower" />
|
||||
<child link="FR_calflower1" />
|
||||
<axis xyz="0 0 0" />
|
||||
</joint>
|
||||
<link name="FR_calf_rotor">
|
||||
<inertial>
|
||||
<origin rpy="0 0 0" xyz="0.0 0.0 0.0" />
|
||||
<mass value="0.089" />
|
||||
<inertia ixx="0.000059647" ixy="0.0" ixz="0.0" iyy="0.000111842" iyz="0.0" izz="0.000059647" />
|
||||
</inertial>
|
||||
</link>
|
||||
<joint name="FR_calf_rotor_joint" type="fixed">
|
||||
<origin rpy="0 0 0" xyz="0 0.03235 0" />
|
||||
<parent link="FR_thigh" />
|
||||
<child link="FR_calf_rotor" />
|
||||
</joint>
|
||||
<link name="FR_foot">
|
||||
<inertial>
|
||||
<origin xyz="0 0 0" rpy="0 0 0" />
|
||||
<mass value="0.04" />
|
||||
<inertia ixx="9.6e-06" ixy="0" ixz="0" iyy="9.6e-06" iyz="0" izz="9.6e-06" />
|
||||
</inertial>
|
||||
<visual>
|
||||
<origin xyz="0 0 0" rpy="0 0 0" />
|
||||
<geometry>
|
||||
<mesh filename="package://go2_description/meshes/foot.dae" />
|
||||
</geometry>
|
||||
<material name="黑色足端_007-effect"><color rgba="0 0 0 1" /></material></visual>
|
||||
<collision>
|
||||
<origin rpy="0 0 0" xyz="-0.002 0 0" />
|
||||
<geometry>
|
||||
<sphere radius="0.022" />
|
||||
</geometry>
|
||||
</collision>
|
||||
</link>
|
||||
<joint name="FR_foot_joint" type="fixed" dont_collapse="true">
|
||||
<origin xyz="0 0 -0.213" rpy="0 0 0" />
|
||||
<parent link="FR_calf" />
|
||||
<child link="FR_foot" />
|
||||
<axis xyz="0 0 0" />
|
||||
</joint>
|
||||
<link name="RL_hip">
|
||||
<inertial>
|
||||
<origin xyz="0.0054 0.00194 -0.000105" rpy="0 0 0" />
|
||||
<mass value="0.678" />
|
||||
<inertia ixx="0.00048" ixy="3.01E-06" ixz="-1.11E-06" iyy="0.000884" iyz="-1.42E-06" izz="0.000596" />
|
||||
</inertial>
|
||||
<visual>
|
||||
<origin xyz="0 0 0" rpy="0 3.1415 0" />
|
||||
<geometry>
|
||||
<mesh filename="package://go2_description/meshes/hip.dae" />
|
||||
</geometry>
|
||||
<material name="深色橡胶_001-effect"><color rgba="0.6717055 0.6924257 0.7742702 1" /></material><material name="Material-effect"><color rgba="0.9 0.95 0.95 1" /></material></visual>
|
||||
<collision>
|
||||
<origin rpy="1.5707963267948966 0 0" xyz="0 0.08 0" />
|
||||
<geometry>
|
||||
<cylinder length="0.04" radius="0.046" />
|
||||
</geometry>
|
||||
</collision>
|
||||
</link>
|
||||
<joint name="RL_hip_joint" type="revolute">
|
||||
<origin xyz="-0.1934 0.0465 0" rpy="0 0 0" />
|
||||
<parent link="base" />
|
||||
<child link="RL_hip" />
|
||||
<axis xyz="1 0 0" />
|
||||
<limit lower="-1.0472" upper="1.0472" effort="23.7" velocity="30.1" />
|
||||
</joint>
|
||||
<link name="RL_hip_rotor">
|
||||
<inertial>
|
||||
<origin rpy="0 0 0" xyz="0.0 0.0 0.0" />
|
||||
<mass value="0.089" />
|
||||
<inertia ixx="0.000111842" ixy="0.0" ixz="0.0" iyy="0.000059647" iyz="0.0" izz="0.000059647" />
|
||||
</inertial>
|
||||
</link>
|
||||
<joint name="RL_hip_rotor_joint" type="fixed">
|
||||
<origin rpy="0 0 0" xyz="-0.11215 0.04675 0" />
|
||||
<parent link="base" />
|
||||
<child link="RL_hip_rotor" />
|
||||
</joint>
|
||||
<link name="RL_thigh">
|
||||
<inertial>
|
||||
<origin xyz="-0.00374 -0.0223 -0.0327" rpy="0 0 0" />
|
||||
<mass value="1.152" />
|
||||
<inertia ixx="0.00584" ixy="8.72E-05" ixz="-0.000289" iyy="0.0058" iyz="0.000808" izz="0.00103" />
|
||||
</inertial>
|
||||
<visual>
|
||||
<origin xyz="0 0 0" rpy="0 0 0" />
|
||||
<geometry>
|
||||
<mesh filename="package://go2_description/meshes/thigh.dae" />
|
||||
</geometry>
|
||||
<material name="深色橡胶_003-effect"><color rgba="0.6717055 0.6924257 0.7742702 1" /></material><material name="Material_009-effect"><color rgba="0.9 0.95 0.95 1" /></material></visual>
|
||||
<collision>
|
||||
<origin rpy="0 1.5707963267948966 0" xyz="0 0 -0.1065" />
|
||||
<geometry>
|
||||
<box size="0.213 0.0245 0.034" />
|
||||
</geometry>
|
||||
</collision>
|
||||
</link>
|
||||
<joint name="RL_thigh_joint" type="revolute">
|
||||
<origin xyz="0 0.0955 0" rpy="0 0 0" />
|
||||
<parent link="RL_hip" />
|
||||
<child link="RL_thigh" />
|
||||
<axis xyz="0 1 0" />
|
||||
<limit lower="-0.5236" upper="4.5379" effort="23.7" velocity="30.1" />
|
||||
</joint>
|
||||
<link name="RL_thigh_rotor">
|
||||
<inertial>
|
||||
<origin rpy="0 0 0" xyz="0.0 0.0 0.0" />
|
||||
<mass value="0.089" />
|
||||
<inertia ixx="0.000059647" ixy="0.0" ixz="0.0" iyy="0.000111842" iyz="0.0" izz="0.000059647" />
|
||||
</inertial>
|
||||
</link>
|
||||
<joint name="RL_thigh_rotor_joint" type="fixed">
|
||||
<origin rpy="0 0 0" xyz="0 -0.00015 0" />
|
||||
<parent link="RL_hip" />
|
||||
<child link="RL_thigh_rotor" />
|
||||
</joint>
|
||||
<link name="RL_calf">
|
||||
<inertial>
|
||||
<origin xyz="0.00548 -0.000975 -0.115" rpy="0 0 0" />
|
||||
<mass value="0.154" />
|
||||
<inertia ixx="0.00108" ixy="3.4E-07" ixz="1.72E-05" iyy="0.0011" iyz="8.28E-06" izz="3.29E-05" />
|
||||
</inertial>
|
||||
<visual>
|
||||
<origin xyz="0 0 0" rpy="0 0 0" />
|
||||
<geometry>
|
||||
<mesh filename="package://go2_description/meshes/calf.dae" />
|
||||
</geometry>
|
||||
<material name="黑色足端_005-effect"><color rgba="0 0 0 1" /></material><material name="深色橡胶_005-effect"><color rgba="0.6717055 0.6924257 0.7742702 1" /></material></visual>
|
||||
<collision>
|
||||
<origin rpy="0 -0.2 0" xyz="0.01 0 -0.06" />
|
||||
<geometry>
|
||||
<cylinder length="0.12" radius="0.013" />
|
||||
</geometry>
|
||||
</collision>
|
||||
</link>
|
||||
<joint name="RL_calf_joint" type="revolute">
|
||||
<origin xyz="0 0 -0.213" rpy="0 0 0" />
|
||||
<parent link="RL_thigh" />
|
||||
<child link="RL_calf" />
|
||||
<axis xyz="0 1 0" />
|
||||
<limit lower="-2.7227" upper="-0.83776" effort="45.43" velocity="15.70" />
|
||||
</joint>
|
||||
<link name="RL_calflower">
|
||||
<collision>
|
||||
<origin rpy="0 0 0" xyz="0 0 0" />
|
||||
<geometry>
|
||||
<cylinder length="0.065" radius="0.011" />
|
||||
</geometry>
|
||||
</collision>
|
||||
</link>
|
||||
<joint name="RL_calflower_joint" type="fixed">
|
||||
<origin xyz="0.020 0 -0.148" rpy="0 0.05 0" />
|
||||
<parent link="RL_calf" />
|
||||
<child link="RL_calflower" />
|
||||
<axis xyz="0 0 0" />
|
||||
</joint>
|
||||
<link name="RL_calflower1">
|
||||
<collision>
|
||||
<origin rpy="0 0 0" xyz="0 0 0" />
|
||||
<geometry>
|
||||
<cylinder length="0.03" radius="0.0155" />
|
||||
</geometry>
|
||||
</collision>
|
||||
</link>
|
||||
<joint name="RL_calflower1_joint" type="fixed">
|
||||
<origin xyz="-0.01 0 -0.04" rpy="0 0.48 0" />
|
||||
<parent link="RL_calflower" />
|
||||
<child link="RL_calflower1" />
|
||||
<axis xyz="0 0 0" />
|
||||
</joint>
|
||||
<link name="RL_calf_rotor">
|
||||
<inertial>
|
||||
<origin rpy="0 0 0" xyz="0.0 0.0 0.0" />
|
||||
<mass value="0.089" />
|
||||
<inertia ixx="0.000059647" ixy="0.0" ixz="0.0" iyy="0.000111842" iyz="0.0" izz="0.000059647" />
|
||||
</inertial>
|
||||
</link>
|
||||
<joint name="RL_calf_rotor_joint" type="fixed">
|
||||
<origin rpy="0 0 0" xyz="0 -0.03235 0" />
|
||||
<parent link="RL_thigh" />
|
||||
<child link="RL_calf_rotor" />
|
||||
</joint>
|
||||
<link name="RL_foot">
|
||||
<inertial>
|
||||
<origin xyz="0 0 0" rpy="0 0 0" />
|
||||
<mass value="0.04" />
|
||||
<inertia ixx="9.6e-06" ixy="0" ixz="0" iyy="9.6e-06" iyz="0" izz="9.6e-06" />
|
||||
</inertial>
|
||||
<visual>
|
||||
<origin xyz="0 0 0" rpy="0 0 0" />
|
||||
<geometry>
|
||||
<mesh filename="package://go2_description/meshes/foot.dae" />
|
||||
</geometry>
|
||||
<material name="黑色足端_007-effect"><color rgba="0 0 0 1" /></material></visual>
|
||||
<collision>
|
||||
<origin rpy="0 0 0" xyz="-0.002 0 0" />
|
||||
<geometry>
|
||||
<sphere radius="0.022" />
|
||||
</geometry>
|
||||
</collision>
|
||||
</link>
|
||||
<joint name="RL_foot_joint" type="fixed" dont_collapse="true">
|
||||
<origin xyz="0 0 -0.213" rpy="0 0 0" />
|
||||
<parent link="RL_calf" />
|
||||
<child link="RL_foot" />
|
||||
<axis xyz="0 0 0" />
|
||||
</joint>
|
||||
<link name="RR_hip">
|
||||
<inertial>
|
||||
<origin xyz="0.0054 -0.00194 -0.000105" rpy="0 0 0" />
|
||||
<mass value="0.678" />
|
||||
<inertia ixx="0.00048" ixy="-3.01E-06" ixz="-1.11E-06" iyy="0.000884" iyz="1.42E-06" izz="0.000596" />
|
||||
</inertial>
|
||||
<visual>
|
||||
<origin xyz="0 0 0" rpy="3.1415 3.1415 0" />
|
||||
<geometry>
|
||||
<mesh filename="package://go2_description/meshes/hip.dae" />
|
||||
</geometry>
|
||||
<material name="深色橡胶_001-effect"><color rgba="0.6717055 0.6924257 0.7742702 1" /></material><material name="Material-effect"><color rgba="0.9 0.95 0.95 1" /></material></visual>
|
||||
<collision>
|
||||
<origin rpy="1.5707963267948966 0 0" xyz="0 -0.08 0" />
|
||||
<geometry>
|
||||
<cylinder length="0.04" radius="0.046" />
|
||||
</geometry>
|
||||
</collision>
|
||||
</link>
|
||||
<joint name="RR_hip_joint" type="revolute">
|
||||
<origin xyz="-0.1934 -0.0465 0" rpy="0 0 0" />
|
||||
<parent link="base" />
|
||||
<child link="RR_hip" />
|
||||
<axis xyz="1 0 0" />
|
||||
<limit lower="-1.0472" upper="1.0472" effort="23.7" velocity="30.1" />
|
||||
</joint>
|
||||
<link name="RR_hip_rotor">
|
||||
<inertial>
|
||||
<origin rpy="0 0 0" xyz="0.0 0.0 0.0" />
|
||||
<mass value="0.089" />
|
||||
<inertia ixx="0.000111842" ixy="0.0" ixz="0.0" iyy="0.000059647" iyz="0.0" izz="0.000059647" />
|
||||
</inertial>
|
||||
</link>
|
||||
<joint name="RR_hip_rotor_joint" type="fixed">
|
||||
<origin rpy="0 0 0" xyz="-0.11215 -0.04675 0" />
|
||||
<parent link="base" />
|
||||
<child link="RR_hip_rotor" />
|
||||
</joint>
|
||||
<link name="RR_thigh">
|
||||
<inertial>
|
||||
<origin xyz="-0.00374 0.0223 -0.0327" rpy="0 0 0" />
|
||||
<mass value="1.152" />
|
||||
<inertia ixx="0.00584" ixy="-8.72E-05" ixz="-0.000289" iyy="0.0058" iyz="-0.000808" izz="0.00103" />
|
||||
</inertial>
|
||||
<visual>
|
||||
<origin xyz="0 0 0" rpy="0 0 0" />
|
||||
<geometry>
|
||||
<mesh filename="package://go2_description/meshes/thigh_mirror.dae" />
|
||||
</geometry>
|
||||
<material name="深色橡胶_012-effect"><color rgba="0.6717055 0.6924257 0.7742702 1" /></material><material name="Material_011-effect"><color rgba="0.9 0.95 0.95 1" /></material></visual>
|
||||
<collision>
|
||||
<origin rpy="0 1.5707963267948966 0" xyz="0 0 -0.1065" />
|
||||
<geometry>
|
||||
<box size="0.213 0.0245 0.034" />
|
||||
</geometry>
|
||||
</collision>
|
||||
</link>
|
||||
<joint name="RR_thigh_joint" type="revolute">
|
||||
<origin xyz="0 -0.0955 0" rpy="0 0 0" />
|
||||
<parent link="RR_hip" />
|
||||
<child link="RR_thigh" />
|
||||
<axis xyz="0 1 0" />
|
||||
<limit lower="-0.5236" upper="4.5379" effort="23.7" velocity="30.1" />
|
||||
</joint>
|
||||
<link name="RR_thigh_rotor">
|
||||
<inertial>
|
||||
<origin rpy="0 0 0" xyz="0.0 0.0 0.0" />
|
||||
<mass value="0.089" />
|
||||
<inertia ixx="0.000059647" ixy="0.0" ixz="0.0" iyy="0.000111842" iyz="0.0" izz="0.000059647" />
|
||||
</inertial>
|
||||
</link>
|
||||
<joint name="RR_thigh_rotor_joint" type="fixed">
|
||||
<origin rpy="0 0 0" xyz="0 0.00015 0" />
|
||||
<parent link="RR_hip" />
|
||||
<child link="RR_thigh_rotor" />
|
||||
</joint>
|
||||
<link name="RR_calf">
|
||||
<inertial>
|
||||
<origin xyz="0.00548 0.000975 -0.115" rpy="0 0 0" />
|
||||
<mass value="0.154" />
|
||||
<inertia ixx="0.00108" ixy="-3.4E-07" ixz="1.72E-05" iyy="0.0011" iyz="-8.28E-06" izz="3.29E-05" />
|
||||
</inertial>
|
||||
<visual>
|
||||
<origin xyz="0 0 0" rpy="0 0 0" />
|
||||
<geometry>
|
||||
<mesh filename="package://go2_description/meshes/calf_mirror.dae" />
|
||||
</geometry>
|
||||
<material name="黑色足端_006-effect"><color rgba="0 0 0 1" /></material><material name="深色橡胶_007-effect"><color rgba="0.6717055 0.6924257 0.7742702 1" /></material></visual>
|
||||
<collision>
|
||||
<origin rpy="0 -0.2 0" xyz="0.01 0 -0.06" />
|
||||
<geometry>
|
||||
<cylinder length="0.12" radius="0.013" />
|
||||
</geometry>
|
||||
</collision>
|
||||
</link>
|
||||
<joint name="RR_calf_joint" type="revolute">
|
||||
<origin xyz="0 0 -0.213" rpy="0 0 0" />
|
||||
<parent link="RR_thigh" />
|
||||
<child link="RR_calf" />
|
||||
<axis xyz="0 1 0" />
|
||||
<limit lower="-2.7227" upper="-0.83776" effort="45.43" velocity="15.70" />
|
||||
</joint>
|
||||
<link name="RR_calflower">
|
||||
<collision>
|
||||
<origin rpy="0 0 0" xyz="0 0 0" />
|
||||
<geometry>
|
||||
<cylinder length="0.065" radius="0.011" />
|
||||
</geometry>
|
||||
</collision>
|
||||
</link>
|
||||
<joint name="RR_calflower_joint" type="fixed">
|
||||
<origin xyz="0.020 0 -0.148" rpy="0 0.05 0" />
|
||||
<parent link="RR_calf" />
|
||||
<child link="RR_calflower" />
|
||||
<axis xyz="0 0 0" />
|
||||
</joint>
|
||||
<link name="RR_calflower1">
|
||||
<collision>
|
||||
<origin rpy="0 0 0" xyz="0 0 0" />
|
||||
<geometry>
|
||||
<cylinder length="0.03" radius="0.0155" />
|
||||
</geometry>
|
||||
</collision>
|
||||
</link>
|
||||
<joint name="RR_calflower1_joint" type="fixed">
|
||||
<origin xyz="-0.01 0 -0.04" rpy="0 0.48 0" />
|
||||
<parent link="RR_calflower" />
|
||||
<child link="RR_calflower1" />
|
||||
<axis xyz="0 0 0" />
|
||||
</joint>
|
||||
<link name="RR_calf_rotor">
|
||||
<inertial>
|
||||
<origin rpy="0 0 0" xyz="0.0 0.0 0.0" />
|
||||
<mass value="0.089" />
|
||||
<inertia ixx="0.000059647" ixy="0.0" ixz="0.0" iyy="0.000111842" iyz="0.0" izz="0.000059647" />
|
||||
</inertial>
|
||||
</link>
|
||||
<joint name="RR_calf_rotor_joint" type="fixed">
|
||||
<origin rpy="0 0 0" xyz="0 0.03235 0" />
|
||||
<parent link="RR_thigh" />
|
||||
<child link="RR_calf_rotor" />
|
||||
</joint>
|
||||
<link name="RR_foot">
|
||||
<inertial>
|
||||
<origin xyz="0 0 0" rpy="0 0 0" />
|
||||
<mass value="0.04" />
|
||||
<inertia ixx="9.6e-06" ixy="0" ixz="0" iyy="9.6e-06" iyz="0" izz="9.6e-06" />
|
||||
</inertial>
|
||||
<visual>
|
||||
<origin xyz="0 0 0" rpy="0 0 0" />
|
||||
<geometry>
|
||||
<mesh filename="package://go2_description/meshes/foot.dae" />
|
||||
</geometry>
|
||||
<material name="黑色足端_007-effect"><color rgba="0 0 0 1" /></material></visual>
|
||||
<collision>
|
||||
<origin rpy="0 0 0" xyz="-0.002 0 0" />
|
||||
<geometry>
|
||||
<sphere radius="0.022" />
|
||||
</geometry>
|
||||
</collision>
|
||||
</link>
|
||||
<joint name="RR_foot_joint" type="fixed" dont_collapse="true">
|
||||
<origin xyz="0 0 -0.213" rpy="0 0 0" />
|
||||
<parent link="RR_calf" />
|
||||
<child link="RR_foot" />
|
||||
<axis xyz="0 0 0" />
|
||||
</joint>
|
||||
<link name="imu">
|
||||
<inertial>
|
||||
<origin xyz="0 0 0" rpy="0 0 0" />
|
||||
<mass value="0" />
|
||||
<inertia ixx="0" ixy="0" ixz="0" iyy="0" iyz="0" izz="0" />
|
||||
</inertial>
|
||||
</link>
|
||||
<joint name="imu_joint" type="fixed">
|
||||
<origin xyz="-0.02557 0 0.04232" rpy="0 0 0" />
|
||||
<parent link="base" />
|
||||
<child link="imu" />
|
||||
<axis xyz="0 0 0" />
|
||||
</joint>
|
||||
<link name="radar">
|
||||
<inertial>
|
||||
<origin xyz="0 0 0" rpy="0 0 0" />
|
||||
<mass value="0" />
|
||||
<inertia ixx="0" ixy="0" ixz="0" iyy="0" iyz="0" izz="0" />
|
||||
</inertial>
|
||||
</link>
|
||||
<joint name="radar_joint" type="fixed">
|
||||
<origin xyz="0.28945 0 -0.046825" rpy="0 2.8782 0" />
|
||||
<parent link="base" />
|
||||
<child link="radar" />
|
||||
<axis xyz="0 0 0" />
|
||||
</joint>
|
||||
<joint name="front_camera_joint" type="fixed">
|
||||
<origin rpy="0 0 0" xyz="0.32715 -0.00003 0.04297" />
|
||||
<parent link="base" />
|
||||
<child link="front_camera" />
|
||||
</joint>
|
||||
<link name="front_camera">
|
||||
</link>
|
||||
</robot>
|
||||
3
source/robot_lab/pyproject.toml
Normal file
3
source/robot_lab/pyproject.toml
Normal file
@@ -0,0 +1,3 @@
|
||||
[build-system]
|
||||
requires = ["setuptools", "wheel", "toml"]
|
||||
build-backend = "setuptools.build_meta"
|
||||
12
source/robot_lab/robot_lab/__init__.py
Normal file
12
source/robot_lab/robot_lab/__init__.py
Normal file
@@ -0,0 +1,12 @@
|
||||
# Copyright (c) 2024-2025 Ziqi Fan
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
"""
|
||||
Python module serving as a project/extension template.
|
||||
"""
|
||||
|
||||
# Register Gym environments.
|
||||
# from .tasks import *
|
||||
|
||||
# Register UI extensions.
|
||||
from .ui_extension_example import *
|
||||
29
source/robot_lab/robot_lab/assets/__init__.py
Normal file
29
source/robot_lab/robot_lab/assets/__init__.py
Normal file
@@ -0,0 +1,29 @@
|
||||
# Copyright (c) 2024-2025 Ziqi Fan
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# Copyright (c) 2022-2025, The Isaac Lab Project Developers.
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
|
||||
"""Package containing asset and sensor configurations."""
|
||||
|
||||
import os
|
||||
import toml
|
||||
|
||||
##
|
||||
# Configuration for different assets.
|
||||
##
|
||||
|
||||
# Conveniences to other module directories via relative paths
|
||||
ISAACLAB_ASSETS_EXT_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../"))
|
||||
"""Path to the extension source directory."""
|
||||
|
||||
ISAACLAB_ASSETS_DATA_DIR = os.path.join(ISAACLAB_ASSETS_EXT_DIR, "data")
|
||||
"""Path to the extension data directory."""
|
||||
|
||||
ISAACLAB_ASSETS_METADATA = toml.load(os.path.join(ISAACLAB_ASSETS_EXT_DIR, "config", "extension.toml"))
|
||||
"""Extension metadata dictionary parsed from the extension.toml file."""
|
||||
|
||||
# Configure the module-level variables
|
||||
__version__ = ISAACLAB_ASSETS_METADATA["package"]["version"]
|
||||
70
source/robot_lab/robot_lab/assets/unitree.py
Normal file
70
source/robot_lab/robot_lab/assets/unitree.py
Normal file
@@ -0,0 +1,70 @@
|
||||
# Copyright (c) 2024-2025 Ziqi Fan
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
"""Configuration for Unitree robots.
|
||||
Reference: https://github.com/unitreerobotics/unitree_ros
|
||||
"""
|
||||
|
||||
import isaaclab.sim as sim_utils
|
||||
from isaaclab.actuators import DCMotorCfg
|
||||
from isaaclab.assets.articulation import ArticulationCfg
|
||||
|
||||
from robot_lab.assets import ISAACLAB_ASSETS_DATA_DIR
|
||||
|
||||
##
|
||||
# Configuration
|
||||
##
|
||||
|
||||
UNITREE_GO2_CFG = ArticulationCfg(
|
||||
spawn=sim_utils.UrdfFileCfg(
|
||||
fix_base=False,
|
||||
merge_fixed_joints=True,
|
||||
replace_cylinders_with_capsules=False,
|
||||
asset_path=f"{ISAACLAB_ASSETS_DATA_DIR}/Robots/unitree/go2_description/urdf/go2_description.urdf",
|
||||
activate_contact_sensors=True,
|
||||
rigid_props=sim_utils.RigidBodyPropertiesCfg(
|
||||
disable_gravity=False,
|
||||
retain_accelerations=False,
|
||||
linear_damping=0.0,
|
||||
angular_damping=0.0,
|
||||
max_linear_velocity=1000.0,
|
||||
max_angular_velocity=1000.0,
|
||||
max_depenetration_velocity=1.0,
|
||||
),
|
||||
articulation_props=sim_utils.ArticulationRootPropertiesCfg(
|
||||
enabled_self_collisions=False,
|
||||
solver_position_iteration_count=4,
|
||||
solver_velocity_iteration_count=0,
|
||||
),
|
||||
joint_drive=sim_utils.UrdfConverterCfg.JointDriveCfg(
|
||||
gains=sim_utils.UrdfConverterCfg.JointDriveCfg.PDGainsCfg(
|
||||
stiffness=0, damping=0
|
||||
)
|
||||
),
|
||||
),
|
||||
init_state=ArticulationCfg.InitialStateCfg(
|
||||
pos=(0.0, 0.0, 0.38),
|
||||
joint_pos={
|
||||
".*L_hip_joint": 0.0,
|
||||
".*R_hip_joint": -0.0,
|
||||
"F.*_thigh_joint": 0.8,
|
||||
"R.*_thigh_joint": 0.8,
|
||||
".*_calf_joint": -1.5,
|
||||
},
|
||||
joint_vel={".*": 0.0},
|
||||
),
|
||||
soft_joint_pos_limit_factor=0.9,
|
||||
actuators={
|
||||
"legs": DCMotorCfg(
|
||||
joint_names_expr=[".*"],
|
||||
effort_limit=23.5,
|
||||
saturation_effort=23.5,
|
||||
velocity_limit=30.0,
|
||||
stiffness=25.0,
|
||||
damping=0.5,
|
||||
friction=0.0,
|
||||
),
|
||||
},
|
||||
)
|
||||
"""Configuration of Unitree Go2 using DC motor.
|
||||
"""
|
||||
24
source/robot_lab/robot_lab/tasks/__init__.py
Normal file
24
source/robot_lab/robot_lab/tasks/__init__.py
Normal file
@@ -0,0 +1,24 @@
|
||||
# Copyright (c) 2024-2025 Ziqi Fan
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# Copyright (c) 2022-2025, The Isaac Lab Project Developers.
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
|
||||
"""Package containing task implementations for various robotic environments."""
|
||||
|
||||
import os
|
||||
import toml
|
||||
|
||||
from isaaclab_tasks.utils import import_packages
|
||||
|
||||
##
|
||||
# Register Gym environments.
|
||||
##
|
||||
|
||||
|
||||
# The blacklist is used to prevent importing configs from sub-packages
|
||||
_BLACKLIST_PKGS = ["utils"]
|
||||
# Import all configs in this package
|
||||
import_packages(__name__, _BLACKLIST_PKGS)
|
||||
34
source/robot_lab/robot_lab/tasks/go2/__init__.py
Normal file
34
source/robot_lab/robot_lab/tasks/go2/__init__.py
Normal file
@@ -0,0 +1,34 @@
|
||||
# Copyright (c) 2024-2025 Ziqi Fan
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# Copyright (c) 2022-2025, The Isaac Lab Project Developers.
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
|
||||
"""Package containing task implementations for various robotic environments."""
|
||||
|
||||
import os
|
||||
import toml
|
||||
import gymnasium as gym
|
||||
|
||||
from isaaclab_tasks.utils import import_packages
|
||||
|
||||
##
|
||||
# Register Gym environments.
|
||||
##
|
||||
gym.register(
|
||||
id="Robotlab-Go2-v0",
|
||||
entry_point="robot_lab.tasks.go2.env.go2_env:ActionDelayGo2Env",
|
||||
disable_env_checker=True,
|
||||
kwargs={
|
||||
"env_cfg_entry_point": f"{__name__}.env_cfg:Go2EnvCfg",
|
||||
# "env_cfg_entry_point": f"{__name__}.origin_env_cfg:UnitreeGo2RoughEnvCfg",
|
||||
"rsl_rl_cfg_entry_point": f"{__name__}.rsl_rl_cfg:MoECTSRunnerCfg",
|
||||
},
|
||||
)
|
||||
|
||||
# The blacklist is used to prevent importing configs from sub-packages
|
||||
_BLACKLIST_PKGS = ["utils"]
|
||||
# Import all configs in this package
|
||||
import_packages(__name__, _BLACKLIST_PKGS)
|
||||
128
source/robot_lab/robot_lab/tasks/go2/env/go2_env.py
vendored
Normal file
128
source/robot_lab/robot_lab/tasks/go2/env/go2_env.py
vendored
Normal file
@@ -0,0 +1,128 @@
|
||||
from isaaclab.envs import ManagerBasedRLEnv, ManagerBasedRLEnvCfg, VecEnvStepReturn
|
||||
from robot_lab.tasks.go2.manager.action_manager import ActionManagerWithDelay
|
||||
|
||||
import torch
|
||||
|
||||
from isaaclab.ui.widgets import ManagerLiveVisualizer
|
||||
|
||||
|
||||
class ActionDelayGo2Env(ManagerBasedRLEnv):
|
||||
cfg: ManagerBasedRLEnvCfg
|
||||
|
||||
def __init__(self, cfg: ManagerBasedRLEnvCfg, render_mode: str | None = None, **kwargs):
|
||||
"""
|
||||
Initialize the ActionDelayGo2Env with the given configuration.
|
||||
|
||||
Args:
|
||||
cfg: The configuration for the environment.
|
||||
render_mode: Rendering mode for the environment, e.g., "human" or "rgb_array". Default is None.
|
||||
**kwargs: Additional keyword arguments for customization.
|
||||
"""
|
||||
# Call the parent class initializer
|
||||
super().__init__(cfg=cfg, render_mode=render_mode, **kwargs)
|
||||
print(
|
||||
"[WARNING] You are using ActionDelayGo2Env; "
|
||||
"make sure all ActionTerms support multiple calls to process_actions() "
|
||||
"within a single step()."
|
||||
)
|
||||
|
||||
def load_managers(self):
|
||||
super().load_managers()
|
||||
# override action manager
|
||||
self.action_manager = ActionManagerWithDelay(self.cfg.actions, self)
|
||||
print("[INFO] Overriding action manager with ActionManagerWithDelay: ", self.action_manager)
|
||||
|
||||
def step(self, action: torch.Tensor) -> VecEnvStepReturn:
|
||||
"""Execute one time-step of the environment's dynamics and reset terminated environments.
|
||||
|
||||
Args:
|
||||
action: The actions to apply on the environment. Shape is (num_envs, action_dim).
|
||||
|
||||
Returns:
|
||||
A tuple containing the observations, rewards, resets (terminated and truncated) and extras.
|
||||
|
||||
IMPORTANT NOTE:
|
||||
We intentionally call action_manager.process_action_with_delay() multiple times inside step
|
||||
(contrary to its original intent) to implement action-delay functionality.
|
||||
So, we assume that all ActionTerm.process_actions() are allowed to be
|
||||
called multiple times within a single step call.
|
||||
"""
|
||||
# call update_action once per step to set action and prev action
|
||||
self.action_manager.update_action(action.to(self.device))
|
||||
|
||||
# randomly determine when to start applying actions within the decimation steps for each environment
|
||||
actions_start_decimation = torch.randint(0, self.cfg.decimation+1, (self.num_envs, 1), device=self.device)
|
||||
|
||||
self.recorder_manager.record_pre_step()
|
||||
|
||||
# check if we need to do rendering within the physics loop
|
||||
# note: checked here once to avoid multiple checks within the loop
|
||||
is_rendering = self.sim.has_gui() or self.sim.has_rtx_sensors()
|
||||
|
||||
# perform physics stepping
|
||||
for i in range(self.cfg.decimation):
|
||||
self._sim_step_counter += 1
|
||||
|
||||
# determine which envs should apply delayed action at this decimation step
|
||||
action_delay_masks = (i <= actions_start_decimation)
|
||||
self.action_manager.process_action_with_delay(action_delay_masks)
|
||||
|
||||
# set actions into buffers
|
||||
self.action_manager.apply_action()
|
||||
# set actions into simulator
|
||||
self.scene.write_data_to_sim()
|
||||
# simulate
|
||||
self.sim.step(render=False)
|
||||
# render between steps only if the GUI or an RTX sensor needs it
|
||||
# note: we assume the render interval to be the shortest accepted rendering interval.
|
||||
# If a camera needs rendering at a faster frequency, this will lead to unexpected behavior.
|
||||
if self._sim_step_counter % self.cfg.sim.render_interval == 0 and is_rendering:
|
||||
self.sim.render()
|
||||
# update buffers at sim dt
|
||||
self.scene.update(dt=self.physics_dt)
|
||||
|
||||
# post-step:
|
||||
# -- update env counters (used for curriculum generation)
|
||||
self.episode_length_buf += 1 # step in current episode (per env)
|
||||
self.common_step_counter += 1 # total step (common for all envs)
|
||||
# -- check terminations
|
||||
self.reset_buf = self.termination_manager.compute()
|
||||
self.reset_terminated = self.termination_manager.terminated
|
||||
self.reset_time_outs = self.termination_manager.time_outs
|
||||
# -- reward computation
|
||||
self.reward_buf = self.reward_manager.compute(dt=self.step_dt)
|
||||
|
||||
if len(self.recorder_manager.active_terms) > 0:
|
||||
# update observations for recording if needed
|
||||
self.obs_buf = self.observation_manager.compute()
|
||||
self.recorder_manager.record_post_step()
|
||||
|
||||
# -- reset envs that terminated/timed-out and log the episode information
|
||||
reset_env_ids = self.reset_buf.nonzero(as_tuple=False).squeeze(-1)
|
||||
if len(reset_env_ids) > 0:
|
||||
# trigger recorder terms for pre-reset calls
|
||||
self.recorder_manager.record_pre_reset(reset_env_ids)
|
||||
|
||||
self._reset_idx(reset_env_ids)
|
||||
# update articulation kinematics
|
||||
self.scene.write_data_to_sim()
|
||||
self.sim.forward()
|
||||
|
||||
# if sensors are added to the scene, make sure we render to reflect changes in reset
|
||||
if self.sim.has_rtx_sensors() and self.cfg.rerender_on_reset:
|
||||
self.sim.render()
|
||||
|
||||
# trigger recorder terms for post-reset calls
|
||||
self.recorder_manager.record_post_reset(reset_env_ids)
|
||||
|
||||
# -- update command
|
||||
self.command_manager.compute(dt=self.step_dt)
|
||||
# -- step interval events
|
||||
if "interval" in self.event_manager.available_modes:
|
||||
self.event_manager.apply(mode="interval", dt=self.step_dt)
|
||||
# -- compute observations
|
||||
# note: done after reset to get the correct observations for reset envs
|
||||
self.obs_buf = self.observation_manager.compute(update_history=True)
|
||||
|
||||
# return observations, rewards, resets and extras
|
||||
return self.obs_buf, self.reward_buf, self.reset_terminated, self.reset_time_outs, self.extras
|
||||
615
source/robot_lab/robot_lab/tasks/go2/env_cfg.py
Normal file
615
source/robot_lab/robot_lab/tasks/go2/env_cfg.py
Normal file
@@ -0,0 +1,615 @@
|
||||
import math
|
||||
import isaaclab.sim as sim_utils
|
||||
from isaaclab.assets import ArticulationCfg, AssetBaseCfg
|
||||
from isaaclab.envs import ManagerBasedRLEnvCfg
|
||||
from isaaclab.managers import CurriculumTermCfg as CurrTerm
|
||||
from isaaclab.managers import EventTermCfg as EventTerm
|
||||
from isaaclab.managers import ObservationGroupCfg as ObsGroup
|
||||
from isaaclab.managers import ObservationTermCfg as ObsTerm
|
||||
from isaaclab.managers import RewardTermCfg as RewTerm
|
||||
from isaaclab.managers import SceneEntityCfg
|
||||
from isaaclab.managers import TerminationTermCfg as DoneTerm
|
||||
from isaaclab.scene import InteractiveSceneCfg
|
||||
from isaaclab.sensors import ContactSensorCfg, RayCasterCfg, patterns
|
||||
from isaaclab.terrains import TerrainImporterCfg
|
||||
from isaaclab.utils import configclass
|
||||
from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR, ISAACLAB_NUCLEUS_DIR
|
||||
from isaaclab.utils.noise import AdditiveUniformNoiseCfg as Unoise
|
||||
import isaaclab.terrains as terrain_gen
|
||||
|
||||
import robot_lab.tasks.go2.mdp as mdp
|
||||
from robot_lab.assets.unitree import UNITREE_GO2_CFG
|
||||
|
||||
JOINT_NAMES = [
|
||||
"FL_hip_joint", "FL_thigh_joint", "FL_calf_joint",
|
||||
"FR_hip_joint", "FR_thigh_joint", "FR_calf_joint",
|
||||
"RL_hip_joint", "RL_thigh_joint", "RL_calf_joint",
|
||||
"RR_hip_joint", "RR_thigh_joint", "RR_calf_joint",
|
||||
]
|
||||
|
||||
BASE_LINK_NAME = "base"
|
||||
FOOT_LINK_NAME = ".*_foot"
|
||||
|
||||
|
||||
##
|
||||
# Terrain definition
|
||||
##
|
||||
|
||||
TERRAIN_CFG = terrain_gen.TerrainGeneratorCfg(
|
||||
size=(8.0, 8.0),
|
||||
border_width=20.0,
|
||||
num_rows=10,
|
||||
num_cols=20,
|
||||
horizontal_scale=0.1,
|
||||
vertical_scale=0.005,
|
||||
slope_threshold=0.75,
|
||||
use_cache=False,
|
||||
sub_terrains={
|
||||
"pyramid_stairs": terrain_gen.MeshPyramidStairsTerrainCfg(
|
||||
proportion=0.15,
|
||||
step_height_range=(0.05, 0.23),
|
||||
step_width=0.3,
|
||||
platform_width=3.0,
|
||||
border_width=1.0,
|
||||
holes=False,
|
||||
),
|
||||
"pyramid_stairs_inv": terrain_gen.MeshInvertedPyramidStairsTerrainCfg(
|
||||
proportion=0.20,
|
||||
step_height_range=(0.05, 0.23),
|
||||
step_width=0.3,
|
||||
platform_width=3.0,
|
||||
border_width=1.0,
|
||||
holes=False,
|
||||
),
|
||||
"boxes": terrain_gen.MeshRandomGridTerrainCfg(
|
||||
proportion=0.15, grid_width=0.45, grid_height_range=(0.05, 0.2), platform_width=2.0
|
||||
),
|
||||
"random_rough": terrain_gen.HfRandomUniformTerrainCfg(
|
||||
proportion=0.1, noise_range=(0.02, 0.10), noise_step=0.02, border_width=0.25
|
||||
),
|
||||
"flat": terrain_gen.MeshPlaneTerrainCfg(proportion=0.2),
|
||||
"hf_pyramid_slope": terrain_gen.HfPyramidSlopedTerrainCfg(
|
||||
proportion=0.1, slope_range=(0.0, 0.4), platform_width=2.0, border_width=0.25
|
||||
),
|
||||
"hf_pyramid_slope_inv": terrain_gen.HfInvertedPyramidSlopedTerrainCfg(
|
||||
proportion=0.1, slope_range=(0.0, 0.4), platform_width=2.0, border_width=0.25
|
||||
),
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
##
|
||||
# Scene definition
|
||||
##
|
||||
|
||||
@configclass
|
||||
class Go2SceneCfg(InteractiveSceneCfg):
|
||||
"""Configuration for the terrain scene with the Go2 robot."""
|
||||
|
||||
terrain = TerrainImporterCfg(
|
||||
prim_path="/World/ground",
|
||||
terrain_type="generator",
|
||||
terrain_generator=TERRAIN_CFG,
|
||||
max_init_terrain_level=5,
|
||||
collision_group=-1,
|
||||
physics_material=sim_utils.RigidBodyMaterialCfg(
|
||||
friction_combine_mode="multiply",
|
||||
static_friction=1.0,
|
||||
dynamic_friction=1.0,
|
||||
),
|
||||
visual_material=sim_utils.MdlFileCfg(
|
||||
mdl_path=f"{ISAACLAB_NUCLEUS_DIR}/Materials/TilesMarbleSpiderWhiteBrickBondHoned/TilesMarbleSpiderWhiteBrickBondHoned.mdl",
|
||||
project_uvw=True,
|
||||
texture_scale=(0.25, 0.25),
|
||||
),
|
||||
debug_vis=False
|
||||
)
|
||||
|
||||
robot: ArticulationCfg = UNITREE_GO2_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot")
|
||||
|
||||
height_scanner = RayCasterCfg(
|
||||
prim_path="{ENV_REGEX_NS}/Robot/base",
|
||||
offset=RayCasterCfg.OffsetCfg(pos=(0.0, 0.0, 20.0)),
|
||||
ray_alignment="yaw",
|
||||
pattern_cfg=patterns.GridPatternCfg(resolution=0.1, size=[1.6, 1.0]),
|
||||
debug_vis=False,
|
||||
mesh_prim_paths=["/World/ground"],
|
||||
)
|
||||
|
||||
height_scanner_base = RayCasterCfg(
|
||||
prim_path="{ENV_REGEX_NS}/Robot/base",
|
||||
offset=RayCasterCfg.OffsetCfg(pos=(0.0, 0.0, 20.0)),
|
||||
ray_alignment="yaw",
|
||||
pattern_cfg=patterns.GridPatternCfg(resolution=0.05, size=(0.1, 0.1)),
|
||||
debug_vis=False,
|
||||
mesh_prim_paths=["/World/ground"],
|
||||
)
|
||||
|
||||
contact_forces = ContactSensorCfg(
|
||||
prim_path="{ENV_REGEX_NS}/Robot/.*",
|
||||
history_length=3,
|
||||
track_air_time=True,
|
||||
)
|
||||
|
||||
# 灯光
|
||||
sky_light = AssetBaseCfg(
|
||||
prim_path="/World/skyLight",
|
||||
spawn=sim_utils.DomeLightCfg(
|
||||
intensity=750.0,
|
||||
texture_file=f"{ISAAC_NUCLEUS_DIR}/Materials/Textures/Skies/PolyHaven/kloofendal_43d_clear_puresky_4k.hdr",
|
||||
),
|
||||
)
|
||||
|
||||
##
|
||||
# MDP settings
|
||||
##
|
||||
|
||||
@configclass
|
||||
class CommandsCfg:
|
||||
"""Command specifications for the MDP."""
|
||||
base_velocity = mdp.UniformVelTerrainCmdCfg(
|
||||
asset_name="robot",
|
||||
resampling_time_range=(5.0, 5.0),
|
||||
rel_standing_envs=0.1,
|
||||
rel_heading_envs=1.0,
|
||||
heading_command=False, # note here, no heading command!
|
||||
heading_control_stiffness=0.5,
|
||||
debug_vis=True,
|
||||
cycle_time=0.5,
|
||||
use_phase_cmd=False,
|
||||
ranges = mdp.UniformVelocityCommandCfg.Ranges(
|
||||
lin_vel_x=(-0.5, 0.5), lin_vel_y=(-0.5, 0.5), ang_vel_z=(-1.0, 1.0), heading=(-1.57, 1.57)
|
||||
),
|
||||
terrain_max_ranges={
|
||||
"pyramid_stairs": mdp.UniformVelocityCommandCfg.Ranges(
|
||||
lin_vel_x=(-1.0, 1.0), lin_vel_y=(-1.0, 1.0), ang_vel_z=(-1.0, 1.0), heading=(-1.57, 1.57)
|
||||
),
|
||||
"pyramid_stairs_inv": mdp.UniformVelocityCommandCfg.Ranges(
|
||||
lin_vel_x=(-1.0, 1.0), lin_vel_y=(-1.0, 1.0), ang_vel_z=(-1.0, 1.0), heading=(-1.57, 1.57)
|
||||
),
|
||||
"box": mdp.UniformVelocityCommandCfg.Ranges(
|
||||
lin_vel_x=(-1.0, 1.0), lin_vel_y=(-1.0, 1.0), ang_vel_z=(-1.5, 1.5), heading=(-1.57, 1.57)
|
||||
),
|
||||
"random_rough": mdp.UniformVelocityCommandCfg.Ranges(
|
||||
lin_vel_x=(-1.5, 1.5), lin_vel_y=(-1.0, 1.0), ang_vel_z=(-1.5, 1.5), heading=(-1.57, 1.57)
|
||||
),
|
||||
"flat": mdp.UniformVelocityCommandCfg.Ranges(
|
||||
lin_vel_x=(-2.0, 2.0), lin_vel_y=(-1.0, 1.0), ang_vel_z=(-2.0, 2.0), heading=(-1.57, 1.57)
|
||||
),
|
||||
"hf_pyramid_slope": mdp.UniformVelocityCommandCfg.Ranges(
|
||||
lin_vel_x=(-1.5, 1.5), lin_vel_y=(-1.0, 1.0), ang_vel_z=(-1.5, 1.5), heading=(-1.57, 1.57)
|
||||
),
|
||||
"hf_pyramid_slope_inv": mdp.UniformVelocityCommandCfg.Ranges(
|
||||
lin_vel_x=(-1.5, 1.5), lin_vel_y=(-1.0, 1.0), ang_vel_z=(-1.5, 1.5), heading=(-1.57, 1.57)
|
||||
),
|
||||
},
|
||||
curriculum_schedule=[
|
||||
{'iter': 50000, 'lin_vel_x': [-2.0, 2.0], 'lin_vel_y': [-1.0, 1.0], 'ang_vel_yaw': [-2.0, 2.0], 'heading': [-1.57, 1.57]},
|
||||
{'iter': 20000, 'lin_vel_x': [-1.0, 1.0], 'lin_vel_y': [-1.0, 1.0], 'ang_vel_yaw': [-1.5, 1.5], 'heading': [-1.57, 1.57]},
|
||||
], # 根据iter倒序排布
|
||||
)
|
||||
|
||||
@configclass
|
||||
class ActionsCfg:
|
||||
"""Action specifications for the MDP."""
|
||||
|
||||
# 腿部关节：位置控制
|
||||
joint_pos = mdp.JointPositionActionCfg(
|
||||
asset_name="robot",
|
||||
joint_names=JOINT_NAMES,
|
||||
scale={".*_hip_joint": 0.25, "^(?!.*_hip_joint).*": 0.25},
|
||||
use_default_offset=True,
|
||||
clip=None,
|
||||
preserve_order=True
|
||||
)
|
||||
|
||||
@configclass
|
||||
class ObservationsCfg:
|
||||
@configclass
|
||||
class PolicyCfg(ObsGroup):
|
||||
"""Observations for policy group."""
|
||||
base_ang_vel = ObsTerm(
|
||||
func=mdp.base_ang_vel,
|
||||
noise=Unoise(n_min=-0.2, n_max=0.2),
|
||||
clip=(-100.0, 100.0),
|
||||
scale=0.25,
|
||||
)
|
||||
projected_gravity = ObsTerm(
|
||||
func=mdp.projected_gravity,
|
||||
noise=Unoise(n_min=-0.05, n_max=0.05),
|
||||
clip=(-100.0, 100.0),
|
||||
scale=1.0,
|
||||
)
|
||||
velocity_commands = ObsTerm(
|
||||
func=mdp.generated_commands,
|
||||
params={"command_name": "base_velocity"},
|
||||
clip=(-100.0, 100.0),
|
||||
scale=1.0,
|
||||
)
|
||||
joint_pos = ObsTerm(
|
||||
func=mdp.joint_pos_rel,
|
||||
params={"asset_cfg": SceneEntityCfg("robot", joint_names=JOINT_NAMES, preserve_order=True)},
|
||||
noise=Unoise(n_min=-0.03, n_max=0.03),
|
||||
clip=(-100.0, 100.0),
|
||||
scale=1.0,
|
||||
)
|
||||
joint_vel = ObsTerm(
|
||||
func=mdp.joint_vel_rel,
|
||||
params={"asset_cfg": SceneEntityCfg("robot", joint_names=JOINT_NAMES, preserve_order=True)},
|
||||
noise=Unoise(n_min=-2.0, n_max=2.0),
|
||||
clip=(-100.0, 100.0),
|
||||
scale=0.05,
|
||||
)
|
||||
actions = ObsTerm(
|
||||
func=mdp.last_action,
|
||||
clip=(-100.0, 100.0),
|
||||
scale=1.0,
|
||||
)
|
||||
def __post_init__(self):
|
||||
self.history_length = 5
|
||||
self.enable_corruption = True
|
||||
self.concatenate_terms = True
|
||||
self.flatten_history_dim = True
|
||||
|
||||
@configclass
|
||||
class CriticCfg(ObsGroup):
|
||||
base_lin_vel = ObsTerm(
|
||||
func=mdp.base_lin_vel,
|
||||
clip=(-100.0, 100.0),
|
||||
scale=2.0,
|
||||
)
|
||||
base_ang_vel = ObsTerm(
|
||||
func=mdp.base_ang_vel,
|
||||
clip=(-100.0, 100.0),
|
||||
scale=0.25,
|
||||
)
|
||||
projected_gravity = ObsTerm(
|
||||
func=mdp.projected_gravity,
|
||||
clip=(-100.0, 100.0),
|
||||
scale=1.0,
|
||||
)
|
||||
velocity_commands = ObsTerm(
|
||||
func=mdp.generated_commands,
|
||||
params={"command_name": "base_velocity"},
|
||||
clip=(-100.0, 100.0),
|
||||
scale=1.0,
|
||||
)
|
||||
joint_pos = ObsTerm(
|
||||
func=mdp.joint_pos_rel,
|
||||
params={"asset_cfg": SceneEntityCfg("robot", joint_names=JOINT_NAMES, preserve_order=True)},
|
||||
clip=(-100.0, 100.0),
|
||||
scale=1.0,
|
||||
)
|
||||
joint_vel = ObsTerm(
|
||||
func=mdp.joint_vel_rel,
|
||||
params={"asset_cfg": SceneEntityCfg("robot", joint_names=JOINT_NAMES, preserve_order=True)},
|
||||
clip=(-100.0, 100.0),
|
||||
scale=0.05,
|
||||
)
|
||||
actions = ObsTerm(
|
||||
func=mdp.last_action,
|
||||
clip=(-100.0, 100.0),
|
||||
scale=1.0,
|
||||
)
|
||||
joint_acc = ObsTerm(
|
||||
func=mdp.joint_acc,
|
||||
params={"asset_cfg": SceneEntityCfg("robot", joint_names=JOINT_NAMES, preserve_order=True)},
|
||||
clip=(-100.0, 100.0),
|
||||
scale=1e-4,
|
||||
)
|
||||
joint_torque = ObsTerm(
|
||||
func=mdp.joint_effort,
|
||||
params={"asset_cfg": SceneEntityCfg("robot", joint_names=JOINT_NAMES, preserve_order=True)},
|
||||
clip=(-100.0, 100.0),
|
||||
scale=0.01,
|
||||
)
|
||||
contact_force = ObsTerm(
|
||||
func=mdp.foot_contact_force_norm,
|
||||
params={"sensor_cfg": SceneEntityCfg("contact_forces", body_names=FOOT_LINK_NAME)},
|
||||
clip=(-100.0, 100.0),
|
||||
scale=1e-3,
|
||||
)
|
||||
height_scan = ObsTerm(
|
||||
func=mdp.height_scan,
|
||||
params={"sensor_cfg": SceneEntityCfg("height_scanner")},
|
||||
clip=(-1.0, 1.0),
|
||||
scale=2.5,
|
||||
)
|
||||
def __post_init__(self):
|
||||
self.enable_corruption = False
|
||||
self.concatenate_terms = True
|
||||
|
||||
@configclass
|
||||
class SingleObsCfg(PolicyCfg):
|
||||
def __post_init__(self):
|
||||
super().__post_init__()
|
||||
self.history_length = 1
|
||||
|
||||
policy: PolicyCfg = PolicyCfg()
|
||||
critic: CriticCfg = CriticCfg()
|
||||
single_obs: SingleObsCfg = SingleObsCfg() # Used to obtain the current-timestep observation for the MoE CTS model
|
||||
|
||||
@configclass
|
||||
class EventCfg:
|
||||
"""Configuration for events."""
|
||||
|
||||
randomize_rigid_body_mass_base = EventTerm(
|
||||
func=mdp.randomize_rigid_body_mass,
|
||||
mode="startup",
|
||||
params={
|
||||
"asset_cfg": SceneEntityCfg("robot", body_names=BASE_LINK_NAME), # Override
|
||||
"mass_distribution_params": (-1.0, 1.0),
|
||||
"operation": "add",
|
||||
"recompute_inertia": True,
|
||||
},
|
||||
)
|
||||
|
||||
randomize_rigid_body_mass_others = EventTerm(
|
||||
func=mdp.randomize_rigid_body_mass,
|
||||
mode="startup",
|
||||
params={
|
||||
"asset_cfg": SceneEntityCfg("robot", body_names="^(?!.*base).*"), # Override
|
||||
"mass_distribution_params": (0.8, 1.2),
|
||||
"operation": "scale",
|
||||
"recompute_inertia": True,
|
||||
},
|
||||
)
|
||||
|
||||
randomize_com_positions = EventTerm(
|
||||
func=mdp.randomize_rigid_body_com,
|
||||
mode="startup",
|
||||
params={
|
||||
"asset_cfg": SceneEntityCfg("robot", body_names=BASE_LINK_NAME), # Override
|
||||
"com_range": {"x": (-0.05, 0.05), "y": (-0.05, 0.05), "z": (-0.05, 0.05)},
|
||||
},
|
||||
)
|
||||
|
||||
randomize_com_positions_other = EventTerm(
|
||||
func=mdp.randomize_rigid_body_com,
|
||||
mode="startup",
|
||||
params={
|
||||
"asset_cfg": SceneEntityCfg("robot", body_names="^(?!.*base).*"), # Override
|
||||
"com_range": {"x": (-0.05, 0.05), "y": (-0.05, 0.05), "z": (-0.05, 0.05)},
|
||||
},
|
||||
)
|
||||
|
||||
reset_robot_joints = EventTerm(
|
||||
func=mdp.reset_joints_by_scale,
|
||||
mode="reset",
|
||||
params={
|
||||
"position_range": (0.8, 1.2),
|
||||
"velocity_range": (0.0, 0.0),
|
||||
},
|
||||
)
|
||||
|
||||
# when use DelayedPDActuator, the range should be narrowed, range that is too wide like (0.5, 2.0) will cause training to fail.
|
||||
randomize_actuator_gains = EventTerm(
|
||||
func=mdp.randomize_actuator_gains,
|
||||
mode="reset",
|
||||
params={
|
||||
"asset_cfg": SceneEntityCfg("robot", joint_names=".*"),
|
||||
"stiffness_distribution_params": (0.9, 1.1),
|
||||
"damping_distribution_params": (0.9, 1.1),
|
||||
"operation": "scale",
|
||||
"distribution": "uniform",
|
||||
},
|
||||
)
|
||||
randomize_push_robot = EventTerm(
|
||||
func=mdp.push_by_setting_velocity,
|
||||
mode="interval",
|
||||
interval_range_s=(4.0, 4.0),
|
||||
params={
|
||||
"velocity_range": {
|
||||
"x": (-0.5, 0.5),
|
||||
"y": (-0.5, 0.5),
|
||||
"roll": (-0.7, 0.7),
|
||||
"pitch": (-0.7, 0.7),
|
||||
"yaw": (-0.7, 0.7)
|
||||
}
|
||||
}
|
||||
)
|
||||
randomize_rigid_body_mass_base = EventTerm(
|
||||
func=mdp.randomize_rigid_body_mass,
|
||||
mode="startup",
|
||||
params={
|
||||
"asset_cfg": SceneEntityCfg("robot", body_names=BASE_LINK_NAME),
|
||||
"mass_distribution_params": (-5.0, 10.0),
|
||||
"operation": "add",
|
||||
"recompute_inertia": True,
|
||||
},
|
||||
)
|
||||
randomize_rigid_body_material = EventTerm(
|
||||
func=mdp.randomize_rigid_body_material,
|
||||
mode="startup",
|
||||
params={
|
||||
"asset_cfg": SceneEntityCfg("robot", body_names=".*"),
|
||||
"static_friction_range": (0.5, 1.5),
|
||||
"dynamic_friction_range": (0.5, 1.5),
|
||||
"restitution_range": (0.0, 0.5),
|
||||
"num_buckets": 512,
|
||||
"make_consistent": True
|
||||
},
|
||||
)
|
||||
randomize_apply_external_force_torque = EventTerm(
|
||||
func=mdp.apply_external_force_torque,
|
||||
mode="reset",
|
||||
params={
|
||||
"asset_cfg": SceneEntityCfg("robot", body_names=BASE_LINK_NAME),
|
||||
"force_range": (-10.0, 10.0),
|
||||
"torque_range": (-10.0, 10.0),
|
||||
},
|
||||
)
|
||||
reset_base = EventTerm(
|
||||
func=mdp.reset_root_state_uniform,
|
||||
mode="reset",
|
||||
params={
|
||||
"pose_range": {"x": (-0.5, 0.5), "y": (-0.5, 0.5), "z": (0.0, 0.2), "yaw": (-3.14, 3.14)},
|
||||
"velocity_range": {
|
||||
"x": (-0.5, 0.5),
|
||||
"y": (-0.5, 0.5),
|
||||
"z": (-0.5, 0.5),
|
||||
"roll": (-0.5, 0.5),
|
||||
"pitch": (-0.5, 0.5),
|
||||
"yaw": (-0.5, 0.5),
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@configclass
|
||||
class RewardsCfg:
|
||||
"""Reward terms for the MDP."""
|
||||
|
||||
track_lin_vel_xy_exp = RewTerm(
|
||||
func=mdp.track_lin_vel_xy_exp,
|
||||
weight=1.0,
|
||||
params={"command_name": "base_velocity", "std": 0.5}
|
||||
)
|
||||
|
||||
track_ang_vel_z_exp = RewTerm(
|
||||
func=mdp.track_ang_vel_z_exp,
|
||||
weight=0.5,
|
||||
params={"command_name": "base_velocity", "std": 0.5}
|
||||
)
|
||||
|
||||
lin_vel_z_l2 = RewTerm(func=mdp.lin_vel_z_l2, weight=-2.0)
|
||||
|
||||
ang_vel_xy_l2 = RewTerm(func=mdp.ang_vel_xy_l2, weight=-0.05)
|
||||
|
||||
flat_orientation_l2 = RewTerm(
|
||||
func=mdp.flat_orientation_l2,
|
||||
weight=-3.0
|
||||
)
|
||||
|
||||
base_height_l2 = RewTerm(
|
||||
func=mdp.base_height_l2,
|
||||
weight=-50.0,
|
||||
params={
|
||||
"asset_cfg": SceneEntityCfg("robot", body_names=BASE_LINK_NAME),
|
||||
"target_height": 0.38
|
||||
}
|
||||
)
|
||||
|
||||
dof_acc_l2 = RewTerm(
|
||||
func=mdp.joint_acc_l2,
|
||||
weight=-2.0e-7,
|
||||
params={"asset_cfg": SceneEntityCfg("robot", joint_names=JOINT_NAMES)}
|
||||
)
|
||||
|
||||
action_rate_l2 = RewTerm(func=mdp.action_rate_l2, weight=-0.01)
|
||||
|
||||
action_smoothness_l2 = RewTerm(func=mdp.action_smoothness_l2, weight=-0.01)
|
||||
|
||||
undesired_contacts = RewTerm(
|
||||
func=mdp.undesired_contacts,
|
||||
weight=-1.0,
|
||||
params={"sensor_cfg": SceneEntityCfg("contact_forces", body_names=".*_thigh|.*_calf"), "threshold": 5.0},
|
||||
)
|
||||
|
||||
hip_pos_penalty = RewTerm(
|
||||
func=mdp.hip_pos_penalty,
|
||||
weight=-0.85,
|
||||
params={
|
||||
"command_name": "base_velocity",
|
||||
"asset_cfg": SceneEntityCfg("robot", joint_names=".*_hip_joint"),
|
||||
"stand_still_scale": 1.0,
|
||||
"command_threshold": 0.1,
|
||||
},
|
||||
)
|
||||
|
||||
joint_torques_l2 = RewTerm(
|
||||
func=mdp.joint_torques_l2,
|
||||
weight=-1e-4,
|
||||
params={"asset_cfg": SceneEntityCfg("robot", joint_names=JOINT_NAMES)}
|
||||
)
|
||||
|
||||
joint_power = RewTerm(
|
||||
func=mdp.joint_power,
|
||||
weight=-2e-5,
|
||||
params={"asset_cfg": SceneEntityCfg("robot", joint_names=JOINT_NAMES)}
|
||||
)
|
||||
|
||||
@configclass
|
||||
class TerminationsCfg:
|
||||
"""Termination terms for the MDP."""
|
||||
time_out = DoneTerm(func=mdp.time_out, time_out=True)
|
||||
bad_orientation = DoneTerm(
|
||||
func=mdp.bad_orientation,
|
||||
params={
|
||||
"asset_cfg": SceneEntityCfg("robot", body_names=BASE_LINK_NAME),
|
||||
"limit_angle": 2 * math.pi / 3
|
||||
},
|
||||
)
|
||||
|
||||
@configclass
|
||||
class CurriculumCfg:
|
||||
"""Curriculum terms for the MDP."""
|
||||
terrain_levels = CurrTerm(func=mdp.terrain_levels_vel_gym)
|
||||
base_linear_velocity = CurrTerm(mdp.gradual_reward_weight_modification, params={
|
||||
"term_name": "lin_vel_z_l2", "initial_weight": -2.0, "final_weight": -0.0, "start_it": 0, "end_it": 2000
|
||||
})
|
||||
ref_stand_envs = CurrTerm(mdp.gradual_ref_stand_modification, params={
|
||||
"term_name": "base_velocity", "initial": 0.0, "final": 0.1, "start_it": 0, "end_it": 1500
|
||||
})
|
||||
command_curr = CurrTerm(mdp.command_curriculum, params={
|
||||
"command_term_name": "base_velocity", "num_steps_per_iter": 24,
|
||||
})
|
||||
|
||||
|
||||
##
|
||||
# Environment configuration
|
||||
##
|
||||
|
||||
@configclass
|
||||
class Go2EnvCfg(ManagerBasedRLEnvCfg):
|
||||
"""Merged configuration for the Go2 robot on rough terrain."""
|
||||
|
||||
# Scene settings
|
||||
scene: Go2SceneCfg = Go2SceneCfg(num_envs=8192, env_spacing=2.5)
|
||||
# Basic settings
|
||||
observations: ObservationsCfg = ObservationsCfg()
|
||||
actions: ActionsCfg = ActionsCfg()
|
||||
commands: CommandsCfg = CommandsCfg()
|
||||
# MDP settings
|
||||
rewards: RewardsCfg = RewardsCfg()
|
||||
terminations: TerminationsCfg = TerminationsCfg()
|
||||
events: EventCfg = EventCfg()
|
||||
curriculum: CurriculumCfg = CurriculumCfg()
|
||||
|
||||
def __post_init__(self):
|
||||
"""Post initialization."""
|
||||
# General settings
|
||||
self.decimation = 4
|
||||
self.episode_length_s = 25.0
|
||||
# Simulation settings
|
||||
self.sim.dt = 0.005
|
||||
self.sim.render_interval = self.decimation
|
||||
|
||||
# Physics material settings from subclass
|
||||
self.sim.physics_material = self.scene.terrain.physics_material
|
||||
self.sim.physx.gpu_max_rigid_patch_count = 10 * 2**15
|
||||
self.sim.physx.gpu_collision_stack_size = int(128 * 1024 * 1024) # 128 MB
|
||||
self.sim.physx.enable_external_forces_every_iteration = True
|
||||
|
||||
# Update sensor periods
|
||||
if self.scene.height_scanner is not None:
|
||||
self.scene.height_scanner.update_period = self.decimation * self.sim.dt
|
||||
if self.scene.contact_forces is not None:
|
||||
self.scene.contact_forces.update_period = self.sim.dt
|
||||
|
||||
# Scale down terrain grid for small robot
|
||||
if self.scene.terrain.terrain_generator is not None:
|
||||
if "boxes" in self.scene.terrain.terrain_generator.sub_terrains:
|
||||
self.scene.terrain.terrain_generator.sub_terrains["boxes"].grid_height_range = (0.025, 0.1)
|
||||
if "random_rough" in self.scene.terrain.terrain_generator.sub_terrains:
|
||||
self.scene.terrain.terrain_generator.sub_terrains["random_rough"].noise_range = (0.01, 0.06)
|
||||
self.scene.terrain.terrain_generator.sub_terrains["random_rough"].noise_step = 0.01
|
||||
|
||||
# Handle curriculum for terrain generator
|
||||
if getattr(self.curriculum, "terrain_levels", None) is not None:
|
||||
if self.scene.terrain.terrain_generator is not None:
|
||||
self.scene.terrain.terrain_generator.curriculum = True
|
||||
else:
|
||||
if self.scene.terrain.terrain_generator is not None:
|
||||
self.scene.terrain.terrain_generator.curriculum = False
|
||||
|
||||
@@ -0,0 +1,47 @@
|
||||
from isaaclab.managers import ActionManager
|
||||
import torch
|
||||
from collections.abc import Sequence
|
||||
|
||||
class ActionManagerWithDelay(ActionManager):
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
self._prev_prev_action = torch.zeros_like(self._action)
|
||||
|
||||
def reset(self, env_ids: Sequence[int] | None = None) -> dict[str, torch.Tensor]:
|
||||
super().reset(env_ids)
|
||||
self._prev_prev_action.zero_()
|
||||
return {}
|
||||
|
||||
def process_action(self, action: torch.Tensor):
|
||||
raise NotImplementedError("Use update_action() and process_action_with_delay() instead for ActionManagerWithDelay.")
|
||||
|
||||
def process_action_with_delay(self, action_delay_masks: torch.Tensor):
|
||||
"""Processes the actions sent to the environment.
|
||||
|
||||
Important Note:
|
||||
This function can be called multiple times within a single step() call to implement action delay.
|
||||
|
||||
Args:
|
||||
action_delay_masks: A tensor of shape (num_envs, 1) indicating which actions to apply at this time step.
|
||||
"""
|
||||
# action_delay_masks == True means delay (use prev action)
|
||||
action = torch.where(action_delay_masks, self._prev_action, self._action)
|
||||
|
||||
# split the actions and apply to each tensor
|
||||
# NOTE: we also assume that all term.process_actions can be called multiple times within a single step() call
|
||||
idx = 0
|
||||
for term in self._terms.values():
|
||||
term_actions = action[:, idx : idx + term.action_dim]
|
||||
term.process_actions(term_actions)
|
||||
idx += term.action_dim
|
||||
|
||||
def update_action(self, action: torch.Tensor):
|
||||
if self.total_action_dim != action.shape[1]:
|
||||
raise ValueError(f"Invalid action shape, expected: {self.total_action_dim}, received: {action.shape[1]}.")
|
||||
self._prev_prev_action[:] = self._prev_action
|
||||
self._prev_action[:] = self._action
|
||||
self._action[:] = action.to(self.device)
|
||||
|
||||
@property
|
||||
def prev_prev_action(self):
|
||||
return self._prev_prev_action
|
||||
19
source/robot_lab/robot_lab/tasks/go2/mdp/__init__.py
Normal file
19
source/robot_lab/robot_lab/tasks/go2/mdp/__init__.py
Normal file
@@ -0,0 +1,19 @@
|
||||
# Copyright (c) 2024-2025 Ziqi Fan
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# Copyright (c) 2022-2025, The Isaac Lab Project Developers.
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
|
||||
"""This sub-module contains the functions that are specific to the locomotion environments."""
|
||||
|
||||
from isaaclab.envs.mdp import * # noqa: F401, F403
|
||||
from isaaclab_tasks.manager_based.locomotion.velocity.mdp import * # noqa: F401, F403
|
||||
|
||||
from .commands import * # noqa: F401, F403
|
||||
from .curriculums import * # noqa: F401, F403
|
||||
from .events import * # noqa: F401, F403
|
||||
from .observations import * # noqa: F401, F403
|
||||
from .rewards import * # noqa: F401, F403
|
||||
from .utils import * # noqa: F401, F403
|
||||
397
source/robot_lab/robot_lab/tasks/go2/mdp/commands.py
Normal file
397
source/robot_lab/robot_lab/tasks/go2/mdp/commands.py
Normal file
@@ -0,0 +1,397 @@
|
||||
# Copyright (c) 2024-2025 Ziqi Fan
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
import torch
|
||||
import copy
|
||||
from collections.abc import Sequence
|
||||
from typing import TYPE_CHECKING
|
||||
from dataclasses import MISSING
|
||||
from itertools import product
|
||||
|
||||
from isaaclab.managers import CommandTerm
|
||||
from isaaclab.utils import configclass
|
||||
from isaaclab.assets import Articulation
|
||||
from isaaclab.envs.mdp import UniformVelocityCommandCfg
|
||||
import isaaclab.utils.math as math_utils
|
||||
from isaaclab.markers import VisualizationMarkers
|
||||
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
# Copyright (c) 2022-2025, The Isaac Lab Project Developers
|
||||
|
||||
from .utils import is_robot_on_terrain
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from isaaclab.envs import ManagerBasedEnv
|
||||
|
||||
|
||||
class UniformVelTerrainCmd(CommandTerm):
|
||||
cfg: UniformVelTerrainCmdCfg
|
||||
|
||||
def __init__(self, cfg: UniformVelTerrainCmdCfg, env: ManagerBasedRLEnv):
|
||||
super().__init__(cfg, env)
|
||||
# obtain the robot asset
|
||||
# -- robot
|
||||
self.robot: Articulation = env.scene[cfg.asset_name]
|
||||
|
||||
# crete buffers to store the command
|
||||
# -- command: x vel, y vel, yaw vel, heading
|
||||
self.vel_command_b = torch.zeros(self.num_envs, 3, device=self.device)
|
||||
self.heading_target = torch.zeros(self.num_envs, device=self.device)
|
||||
self.is_heading_env = torch.zeros(self.num_envs, dtype=torch.bool, device=self.device)
|
||||
self.is_standing_env = torch.zeros_like(self.is_heading_env)
|
||||
|
||||
# -- metrics
|
||||
self.metrics["error_vel_xy"] = torch.zeros(self.num_envs, device=self.device)
|
||||
self.metrics["error_vel_yaw"] = torch.zeros(self.num_envs, device=self.device)
|
||||
if self.cfg.use_phase_cmd:
|
||||
self.phase = torch.zeros(self.num_envs, 1, device=self.device)
|
||||
self.cycle_time = cfg.cycle_time
|
||||
self.max_angular_envs = cfg.max_angular_envs
|
||||
self.limit_vel_envs = cfg.limit_vel_envs
|
||||
self.stop_heading = torch.zeros(self.num_envs, dtype=torch.bool, device=self.device)
|
||||
|
||||
# 获取当前环境的地形类型索引
|
||||
self.terrain_types = list(env.scene.terrain.cfg.terrain_generator.sub_terrains.keys())
|
||||
if not isinstance(self.cfg.ranges, dict):
|
||||
single_range = cfg.ranges
|
||||
expanded_ranges = {}
|
||||
for t_name in self.terrain_types:
|
||||
expanded_ranges[t_name] = copy.deepcopy(single_range)
|
||||
self.cfg.ranges = expanded_ranges
|
||||
assert set(self.terrain_types) == set(list(self.cfg.ranges.keys())), \
|
||||
"Terrain types in cfg.ranges do not match those in terrain generator config."
|
||||
self.terrain_type_to_id = {name: i for i, name in enumerate(self.terrain_types)}
|
||||
self.global_env_terrain_idx = torch.full((self.num_envs,), -1, dtype=torch.long, device=self.device)
|
||||
for t_type in self.terrain_types:
|
||||
ids = is_robot_on_terrain(self._env, t_type).nonzero(as_tuple=True)[0]
|
||||
if len(ids) > 0:
|
||||
self.global_env_terrain_idx[ids] = self.terrain_type_to_id[t_type]
|
||||
# 初始化累计指令和最大位移
|
||||
self.commands_xy_accumulation = torch.zeros(self.num_envs, 2, device=self.device)
|
||||
self.dt = self._env.step_dt
|
||||
self.max_episode_length = self._env.max_episode_length
|
||||
self.target_dist = env.scene.terrain.cfg.terrain_generator.size[0] * 0.625
|
||||
self.max_move_distance = torch.zeros(self.num_envs, device=self.device)
|
||||
self.env_origins = env.scene.env_origins
|
||||
|
||||
# 极限指令组合
|
||||
limit_options_x = [-1, 1]
|
||||
limit_options_y = [-1, 1]
|
||||
limit_options_z = [-1, 0, 1] # 允许 Yaw 为 0，即直线冲刺
|
||||
self.limit_vel_combinations = torch.tensor(
|
||||
list(product(limit_options_x, limit_options_y, limit_options_z)),
|
||||
dtype=torch.long,
|
||||
device=self.device
|
||||
)
|
||||
|
||||
self.last_is_limit_vel = torch.zeros(self.num_envs, dtype=torch.bool, device=self.device)
|
||||
|
||||
def __str__(self) -> str:
|
||||
"""Return a string representation of the command generator."""
|
||||
msg = "UniformVelocityCommand:\n"
|
||||
msg += f"\tCommand dimension: {tuple(self.command.shape[1:])}\n"
|
||||
msg += f"\tResampling time range: {self.cfg.resampling_time_range}\n"
|
||||
msg += f"\tHeading command: {self.cfg.heading_command}\n"
|
||||
if self.cfg.heading_command:
|
||||
msg += f"\tHeading probability: {self.cfg.rel_heading_envs}\n"
|
||||
msg += f"\tStanding probability: {self.cfg.rel_standing_envs}\n"
|
||||
msg += f"\tMax angular vel envs: {self.max_angular_envs}\n"
|
||||
msg += f"\tLimit linear vel envs: {self.limit_vel_envs}\n"
|
||||
for t_type in self.terrain_types:
|
||||
msg += f"\tTerrain type '{t_type}' command ranges: {self.cfg.ranges[t_type]}"
|
||||
return msg
|
||||
|
||||
@property
|
||||
def command(self) -> torch.Tensor:
|
||||
if self.cfg.use_phase_cmd:
|
||||
phase = self.phase * 2 * torch.pi
|
||||
return torch.cat([self.vel_command_b, torch.sin(phase), torch.cos(phase)], dim=-1)
|
||||
return self.vel_command_b
|
||||
|
||||
def _update_metrics(self):
|
||||
# time for which the command was executed
|
||||
max_command_time = self.cfg.resampling_time_range[1]
|
||||
max_command_step = max_command_time / self._env.step_dt
|
||||
# logs data
|
||||
self.metrics["error_vel_xy"] += (
|
||||
torch.norm(self.vel_command_b[:, :2] - self.robot.data.root_lin_vel_b[:, :2], dim=-1) / max_command_step
|
||||
)
|
||||
self.metrics["error_vel_yaw"] += (
|
||||
torch.abs(self.vel_command_b[:, 2] - self.robot.data.root_ang_vel_b[:, 2]) / max_command_step
|
||||
)
|
||||
|
||||
def reset(self, env_ids: Sequence[int] | None = None):
|
||||
self.commands_xy_accumulation[env_ids] = 0.0
|
||||
self.max_move_distance[env_ids] = 0.0
|
||||
extra = super().reset(env_ids)
|
||||
return extra
|
||||
|
||||
def _terrain_cmd_sample(self, dim: str, env_ids: torch.Tensor, min_abs_vel: torch.Tensor | None = None):
|
||||
"""根据不同地形sample不同指令"""
|
||||
cmd = torch.zeros(len(env_ids), device=self.device)
|
||||
batch_terrain_idxs = self.global_env_terrain_idx[env_ids]
|
||||
|
||||
for t_type in self.terrain_types:
|
||||
t_id = self.terrain_type_to_id[t_type]
|
||||
mask = (batch_terrain_idxs == t_id)
|
||||
|
||||
if mask.any():
|
||||
count = mask.sum().item()
|
||||
if dim == "x":
|
||||
r = self.cfg.ranges[t_type].lin_vel_x
|
||||
elif dim == "y":
|
||||
r = self.cfg.ranges[t_type].lin_vel_y
|
||||
elif dim == "z":
|
||||
r = self.cfg.ranges[t_type].ang_vel_z
|
||||
elif dim == "heading":
|
||||
r = self.cfg.ranges[t_type].heading
|
||||
|
||||
sampled_vals = torch.empty(count, device=self.device).uniform_(*r) # type: ignore
|
||||
|
||||
if min_abs_vel is not None and (dim == "x" or dim == "y"):
|
||||
lower_bound = min_abs_vel[mask]
|
||||
|
||||
max_abs = max(abs(r[0]), abs(r[1]))
|
||||
lower_bound = torch.clamp(lower_bound, max=max_abs)
|
||||
|
||||
# 重新采样 Magnitude: [lower_bound, max_abs]
|
||||
mag = torch.empty(count, device=self.device).uniform_(0, 1) * (max_abs - lower_bound) + lower_bound
|
||||
|
||||
if r[0] >= 0:
|
||||
sampled_vals = mag
|
||||
elif r[1] <= 0:
|
||||
sampled_vals = -mag
|
||||
else:
|
||||
sign = torch.sign(torch.empty(count, device=self.device).uniform_(-1, 1))
|
||||
sign = torch.where(sign == 0, torch.ones_like(sign), sign)
|
||||
sampled_vals = mag * sign
|
||||
|
||||
cmd[mask] = sampled_vals
|
||||
|
||||
return cmd
|
||||
|
||||
def _get_cmd_board(self, dim: str, env_ids: torch.Tensor):
|
||||
"""根据预计算的地形类型获取速度指令边界"""
|
||||
cmd_min = torch.zeros(len(env_ids), device=self.device)
|
||||
cmd_max = torch.zeros(len(env_ids), device=self.device)
|
||||
|
||||
batch_terrain_idxs = self.global_env_terrain_idx[env_ids]
|
||||
|
||||
for t_type in self.terrain_types:
|
||||
t_id = self.terrain_type_to_id[t_type]
|
||||
mask = (batch_terrain_idxs == t_id)
|
||||
|
||||
if mask.any():
|
||||
if dim == "x":
|
||||
r = self.cfg.ranges[t_type].lin_vel_x
|
||||
elif dim == "y":
|
||||
r = self.cfg.ranges[t_type].lin_vel_y
|
||||
elif dim == "z":
|
||||
r = self.cfg.ranges[t_type].ang_vel_z
|
||||
else:
|
||||
continue
|
||||
|
||||
cmd_min[mask] = r[0]
|
||||
cmd_max[mask] = r[1]
|
||||
|
||||
return cmd_min, cmd_max
|
||||
|
||||
def _resample_command(self, env_ids: Sequence[int]):
|
||||
_env_ids = torch.tensor(env_ids, device=self.device)
|
||||
|
||||
# 计算剩余距离: 目标距离 - 已经指令累积走过的距离 * 上一次重采样时间
|
||||
dist_covered = torch.norm(self.commands_xy_accumulation[env_ids], dim=1) * self.cfg.resampling_time_range[0]
|
||||
remaining_dist = torch.clamp(self.target_dist - dist_covered, min=0.0)
|
||||
|
||||
# 计算剩余时间: (最大步数 - 当前步数) * dt
|
||||
time_left = (self.max_episode_length - self._env.episode_length_buf[env_ids]) * self.dt
|
||||
# 计算下限速度: 距离 / 时间
|
||||
vel_low_bound = torch.zeros(len(env_ids), device=self.device)
|
||||
# 避免除以零或负数时间
|
||||
valid_time_mask = time_left > 1e-4
|
||||
if valid_time_mask.any():
|
||||
vel_low_bound[valid_time_mask] = remaining_dist[valid_time_mask] / time_left[valid_time_mask]
|
||||
|
||||
# 根据不同地形类型采样速度指令
|
||||
self.vel_command_b[env_ids, 0] = self._terrain_cmd_sample("x", _env_ids, min_abs_vel=vel_low_bound)
|
||||
self.vel_command_b[env_ids, 0] = torch.where(
|
||||
self.vel_command_b[env_ids, 0].abs() < 0.1,
|
||||
torch.zeros_like(self.vel_command_b[env_ids, 0]),
|
||||
self.vel_command_b[env_ids, 0]
|
||||
)
|
||||
self.vel_command_b[env_ids, 1] = self._terrain_cmd_sample("y", _env_ids, min_abs_vel=vel_low_bound)
|
||||
self.vel_command_b[env_ids, 1] = torch.where(
|
||||
self.vel_command_b[env_ids, 1].abs() < 0.1,
|
||||
torch.zeros_like(self.vel_command_b[env_ids, 1]),
|
||||
self.vel_command_b[env_ids, 1]
|
||||
)
|
||||
self.vel_command_b[env_ids, 2] = self._terrain_cmd_sample("z", _env_ids)
|
||||
self.vel_command_b[env_ids, 2] = torch.where(
|
||||
self.vel_command_b[env_ids, 2].abs() < 0.1,
|
||||
torch.zeros_like(self.vel_command_b[env_ids, 2]),
|
||||
self.vel_command_b[env_ids, 2]
|
||||
)
|
||||
|
||||
# 计算heading和stand
|
||||
r = torch.empty(len(env_ids), device=self.device)
|
||||
if self.cfg.heading_command:
|
||||
self.heading_target[env_ids] = self._terrain_cmd_sample("heading", _env_ids)
|
||||
# update heading envs
|
||||
self.is_heading_env[env_ids] = r.uniform_(0.0, 1.0) <= self.cfg.rel_heading_envs
|
||||
# update standing envs
|
||||
self.is_standing_env[env_ids] = r.uniform_(0.0, 1.0) <= self.cfg.rel_standing_envs
|
||||
|
||||
# 静止环境sample最大转向
|
||||
standing_env_ids = self.is_standing_env.nonzero(as_tuple=False).flatten()
|
||||
if len(standing_env_ids) > 0:
|
||||
self.vel_command_b[standing_env_ids, :] = 0.0
|
||||
ang_vel_rand = torch.rand(len(standing_env_ids), device=self.device)
|
||||
add_ang_mask = ang_vel_rand < self.max_angular_envs
|
||||
add_ang_env_ids = standing_env_ids[add_ang_mask]
|
||||
if self.max_angular_envs > 0 and len(add_ang_env_ids) > 0:
|
||||
direction_rand = torch.rand(len(add_ang_env_ids), device=self.device)
|
||||
min_z, max_z = self._get_cmd_board("z", add_ang_env_ids)
|
||||
self.vel_command_b[add_ang_env_ids, 2] = torch.where(direction_rand < 0.5, min_z, max_z)
|
||||
self.stop_heading[add_ang_env_ids] = True
|
||||
|
||||
limit_vel_env_ids = (self.is_standing_env == 0).nonzero(as_tuple=False).flatten()
|
||||
|
||||
# 非静止环境sample极限速度
|
||||
current_ids = env_ids
|
||||
mask = torch.zeros(self.num_envs, dtype=torch.bool, device=self.device)
|
||||
mask[current_ids] = True
|
||||
target_limit_ids = limit_vel_env_ids[mask[limit_vel_env_ids]]
|
||||
|
||||
if self.limit_vel_envs > 0 and len(target_limit_ids) > 0:
|
||||
limit_prob = torch.rand(len(target_limit_ids), device=self.device)
|
||||
execute_limit_mask = limit_prob < self.limit_vel_envs
|
||||
execute_ids = target_limit_ids[execute_limit_mask]
|
||||
|
||||
if len(execute_ids) > 0:
|
||||
|
||||
num_combs = self.limit_vel_combinations.shape[0]
|
||||
comb_indices = torch.randint(0, num_combs, (len(execute_ids),), device=self.device)
|
||||
selected_combs = self.limit_vel_combinations[comb_indices] # (N, 3)
|
||||
|
||||
min_x, max_x = self._get_cmd_board("x", execute_ids)
|
||||
min_y, max_y = self._get_cmd_board("y", execute_ids)
|
||||
min_z, max_z = self._get_cmd_board("z", execute_ids)
|
||||
|
||||
vals_x = torch.zeros_like(min_x)
|
||||
vals_x = torch.where(selected_combs[:, 0] == -1, min_x, vals_x)
|
||||
vals_x = torch.where(selected_combs[:, 0] == 1, max_x, vals_x)
|
||||
|
||||
vals_y = torch.zeros_like(min_y)
|
||||
vals_y = torch.where(selected_combs[:, 1] == -1, min_y, vals_y)
|
||||
vals_y = torch.where(selected_combs[:, 1] == 1, max_y, vals_y)
|
||||
|
||||
vals_z = torch.zeros_like(min_z)
|
||||
vals_z = torch.where(selected_combs[:, 2] == -1, min_z, vals_z)
|
||||
vals_z = torch.where(selected_combs[:, 2] == 1, max_z, vals_z)
|
||||
|
||||
self.vel_command_b[execute_ids, 0] = vals_x
|
||||
self.vel_command_b[execute_ids, 1] = vals_y
|
||||
self.vel_command_b[execute_ids, 2] = vals_z
|
||||
|
||||
self.stop_heading[execute_ids] = True
|
||||
|
||||
if self.cfg.use_phase_cmd:
|
||||
self.phase[env_ids] = torch.rand((len(env_ids), 1), device=self.device)
|
||||
self.commands_xy_accumulation[env_ids] += self.vel_command_b[env_ids, :2]
|
||||
|
||||
def _update_command(self):
|
||||
# 停止heading更新
|
||||
stop_heading_env_ids = self.stop_heading.nonzero(as_tuple=False).flatten()
|
||||
if len(stop_heading_env_ids) > 0:
|
||||
self.is_heading_env[stop_heading_env_ids] = False
|
||||
|
||||
# Compute angular velocity from heading direction
|
||||
if self.cfg.heading_command:
|
||||
# resolve indices of heading envs
|
||||
env_ids = self.is_heading_env.nonzero(as_tuple=False).flatten()
|
||||
# compute angular velocity
|
||||
if len(env_ids) > 0:
|
||||
heading_error = math_utils.wrap_to_pi(self.heading_target[env_ids] - self.robot.data.heading_w[env_ids])
|
||||
min_z, max_z = self._get_cmd_board("z", env_ids)
|
||||
|
||||
self.vel_command_b[env_ids, 2] = torch.clip(
|
||||
self.cfg.heading_control_stiffness * heading_error,
|
||||
min=min_z,
|
||||
max=max_z,
|
||||
)
|
||||
if self.cfg.use_phase_cmd:
|
||||
self.phase = self._env.episode_length_buf[:, None] * self._env.step_dt / self.cycle_time
|
||||
current_dist = torch.norm(self.robot.data.root_pos_w[:, :2] - self.env_origins[:, :2], dim=1)
|
||||
self.max_move_distance = torch.max(self.max_move_distance, current_dist)
|
||||
|
||||
def _set_debug_vis_impl(self, debug_vis: bool):
|
||||
# set visibility of markers
|
||||
# note: parent only deals with callbacks. not their visibility
|
||||
if debug_vis:
|
||||
# create markers if necessary for the first time
|
||||
if not hasattr(self, "goal_vel_visualizer"):
|
||||
# -- goal
|
||||
self.goal_vel_visualizer = VisualizationMarkers(self.cfg.goal_vel_visualizer_cfg)
|
||||
# -- current
|
||||
self.current_vel_visualizer = VisualizationMarkers(self.cfg.current_vel_visualizer_cfg)
|
||||
# set their visibility to true
|
||||
self.goal_vel_visualizer.set_visibility(True)
|
||||
self.current_vel_visualizer.set_visibility(True)
|
||||
else:
|
||||
if hasattr(self, "goal_vel_visualizer"):
|
||||
self.goal_vel_visualizer.set_visibility(False)
|
||||
self.current_vel_visualizer.set_visibility(False)
|
||||
|
||||
def _debug_vis_callback(self, event):
|
||||
# check if robot is initialized
|
||||
# note: this is needed in-case the robot is de-initialized. we can't access the data
|
||||
if not self.robot.is_initialized:
|
||||
return
|
||||
# get marker location
|
||||
# -- base state
|
||||
base_pos_w = self.robot.data.root_pos_w.clone()
|
||||
base_pos_w[:, 2] += 0.5
|
||||
# -- resolve the scales and quaternions
|
||||
vel_des_arrow_scale, vel_des_arrow_quat = self._resolve_xy_velocity_to_arrow(self.command[:, :2])
|
||||
vel_arrow_scale, vel_arrow_quat = self._resolve_xy_velocity_to_arrow(self.robot.data.root_lin_vel_b[:, :2])
|
||||
# display markers
|
||||
self.goal_vel_visualizer.visualize(base_pos_w, vel_des_arrow_quat, vel_des_arrow_scale)
|
||||
self.current_vel_visualizer.visualize(base_pos_w, vel_arrow_quat, vel_arrow_scale)
|
||||
|
||||
"""
|
||||
Internal helpers.
|
||||
"""
|
||||
|
||||
def _resolve_xy_velocity_to_arrow(self, xy_velocity: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Converts the XY base velocity command to arrow direction rotation."""
|
||||
# obtain default scale of the marker
|
||||
default_scale = self.goal_vel_visualizer.cfg.markers["arrow"].scale
|
||||
# arrow-scale
|
||||
arrow_scale = torch.tensor(default_scale, device=self.device).repeat(xy_velocity.shape[0], 1)
|
||||
arrow_scale[:, 0] *= torch.linalg.norm(xy_velocity, dim=1) * 3.0
|
||||
# arrow-direction
|
||||
heading_angle = torch.atan2(xy_velocity[:, 1], xy_velocity[:, 0])
|
||||
zeros = torch.zeros_like(heading_angle)
|
||||
arrow_quat = math_utils.quat_from_euler_xyz(zeros, zeros, heading_angle)
|
||||
# convert everything back from base to world frame
|
||||
base_quat_w = self.robot.data.root_quat_w
|
||||
arrow_quat = math_utils.quat_mul(base_quat_w, arrow_quat)
|
||||
|
||||
return arrow_scale, arrow_quat
|
||||
|
||||
|
||||
@configclass
|
||||
class UniformVelTerrainCmdCfg(UniformVelocityCommandCfg):
|
||||
|
||||
class_type: type = UniformVelTerrainCmd
|
||||
cycle_time:float = 0.5
|
||||
max_angular_envs: float = 0.2
|
||||
limit_vel_envs: float = 0.2
|
||||
ranges: dict[str, UniformVelocityCommandCfg.Ranges] | UniformVelocityCommandCfg.Ranges = MISSING # 当前指令范围
|
||||
terrain_max_ranges: dict[str, UniformVelocityCommandCfg.Ranges] = MISSING # type:ignore 地形最大指令范围
|
||||
curriculum_schedule: list[dict] | None = None
|
||||
use_phase_cmd: bool = True
|
||||
|
||||
216
source/robot_lab/robot_lab/tasks/go2/mdp/curriculums.py
Normal file
216
source/robot_lab/robot_lab/tasks/go2/mdp/curriculums.py
Normal file
@@ -0,0 +1,216 @@
|
||||
# Copyright (c) 2024-2025 Ziqi Fan
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
"""Common functions that can be used to create curriculum for the learning environment.
|
||||
|
||||
The functions can be passed to the :class:`isaaclab.managers.CurriculumTermCfg` object to enable
|
||||
the curriculum introduced by the function.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
from collections.abc import Sequence
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from isaaclab.envs import ManagerBasedRLEnv
|
||||
|
||||
|
||||
def command_levels_lin_vel(
|
||||
env: ManagerBasedRLEnv,
|
||||
env_ids: Sequence[int],
|
||||
reward_term_name: str,
|
||||
range_multiplier: Sequence[float] = (0.1, 1.0),
|
||||
) -> None:
|
||||
"""command_levels_lin_vel"""
|
||||
base_velocity_ranges = env.command_manager.get_term("base_velocity").cfg.ranges
|
||||
# Get original velocity ranges (ONLY ON FIRST EPISODE)
|
||||
if env.common_step_counter == 0:
|
||||
env._original_vel_x = torch.tensor(base_velocity_ranges.lin_vel_x, device=env.device)
|
||||
env._original_vel_y = torch.tensor(base_velocity_ranges.lin_vel_y, device=env.device)
|
||||
env._initial_vel_x = env._original_vel_x * range_multiplier[0]
|
||||
env._final_vel_x = env._original_vel_x * range_multiplier[1]
|
||||
env._initial_vel_y = env._original_vel_y * range_multiplier[0]
|
||||
env._final_vel_y = env._original_vel_y * range_multiplier[1]
|
||||
|
||||
# Initialize command ranges to initial values
|
||||
base_velocity_ranges.lin_vel_x = env._initial_vel_x.tolist()
|
||||
base_velocity_ranges.lin_vel_y = env._initial_vel_y.tolist()
|
||||
|
||||
# avoid updating command curriculum at each step since the maximum command is common to all envs
|
||||
if env.common_step_counter % env.max_episode_length == 0:
|
||||
episode_sums = env.reward_manager._episode_sums[reward_term_name]
|
||||
reward_term_cfg = env.reward_manager.get_term_cfg(reward_term_name)
|
||||
delta_command = torch.tensor([-0.1, 0.1], device=env.device)
|
||||
|
||||
# If the tracking reward is above 80% of the maximum, increase the range of commands
|
||||
if torch.mean(episode_sums[env_ids]) / env.max_episode_length_s > 0.8 * reward_term_cfg.weight:
|
||||
new_vel_x = torch.tensor(base_velocity_ranges.lin_vel_x, device=env.device) + delta_command
|
||||
new_vel_y = torch.tensor(base_velocity_ranges.lin_vel_y, device=env.device) + delta_command
|
||||
|
||||
# Clamp to ensure we don't exceed final ranges
|
||||
new_vel_x = torch.clamp(new_vel_x, min=env._final_vel_x[0], max=env._final_vel_x[1])
|
||||
new_vel_y = torch.clamp(new_vel_y, min=env._final_vel_y[0], max=env._final_vel_y[1])
|
||||
|
||||
# Update ranges
|
||||
base_velocity_ranges.lin_vel_x = new_vel_x.tolist()
|
||||
base_velocity_ranges.lin_vel_y = new_vel_y.tolist()
|
||||
|
||||
return torch.tensor(base_velocity_ranges.lin_vel_x[1], device=env.device)
|
||||
|
||||
|
||||
def command_levels_ang_vel(
|
||||
env: ManagerBasedRLEnv,
|
||||
env_ids: Sequence[int],
|
||||
reward_term_name: str,
|
||||
range_multiplier: Sequence[float] = (0.1, 1.0),
|
||||
) -> None:
|
||||
"""command_levels_ang_vel"""
|
||||
base_velocity_ranges = env.command_manager.get_term("base_velocity").cfg.ranges
|
||||
# Get original angular velocity ranges (ONLY ON FIRST EPISODE)
|
||||
if env.common_step_counter == 0:
|
||||
env._original_ang_vel_z = torch.tensor(base_velocity_ranges.ang_vel_z, device=env.device)
|
||||
env._initial_ang_vel_z = env._original_ang_vel_z * range_multiplier[0]
|
||||
env._final_ang_vel_z = env._original_ang_vel_z * range_multiplier[1]
|
||||
|
||||
# Initialize command ranges to initial values
|
||||
base_velocity_ranges.ang_vel_z = env._initial_ang_vel_z.tolist()
|
||||
|
||||
# avoid updating command curriculum at each step since the maximum command is common to all envs
|
||||
if env.common_step_counter % env.max_episode_length == 0:
|
||||
episode_sums = env.reward_manager._episode_sums[reward_term_name]
|
||||
reward_term_cfg = env.reward_manager.get_term_cfg(reward_term_name)
|
||||
delta_command = torch.tensor([-0.1, 0.1], device=env.device)
|
||||
|
||||
# If the tracking reward is above 80% of the maximum, increase the range of commands
|
||||
if torch.mean(episode_sums[env_ids]) / env.max_episode_length_s > 0.8 * reward_term_cfg.weight:
|
||||
new_ang_vel_z = torch.tensor(base_velocity_ranges.ang_vel_z, device=env.device) + delta_command
|
||||
|
||||
# Clamp to ensure we don't exceed final ranges
|
||||
new_ang_vel_z = torch.clamp(new_ang_vel_z, min=env._final_ang_vel_z[0], max=env._final_ang_vel_z[1])
|
||||
|
||||
# Update ranges
|
||||
base_velocity_ranges.ang_vel_z = new_ang_vel_z.tolist()
|
||||
|
||||
return torch.tensor(base_velocity_ranges.ang_vel_z[1], device=env.device)
|
||||
|
||||
def command_curriculum(
|
||||
env: ManagerBasedRLEnv,
|
||||
env_ids: Sequence[int],
|
||||
command_term_name: str,
|
||||
num_steps_per_iter: int = 24,
|
||||
) -> float:
|
||||
"""
|
||||
阶跃式指令课程 (数据存储在 CommandCfg 中)。
|
||||
"""
|
||||
try:
|
||||
cmd_term = env.command_manager.get_term(command_term_name)
|
||||
cmd_cfg = cmd_term.cfg
|
||||
except LookupError:
|
||||
return 0.0
|
||||
|
||||
current_iter = env.common_step_counter // num_steps_per_iter
|
||||
schedule = cmd_cfg.curriculum_schedule
|
||||
|
||||
for i in range(len(schedule) - 1, -1, -1):
|
||||
stage = schedule[i]
|
||||
|
||||
if current_iter >= stage['iter']:
|
||||
for t_name, active_range in cmd_cfg.ranges.items():
|
||||
|
||||
hard_limit = cmd_cfg.terrain_max_ranges.get(t_name)
|
||||
if hard_limit is None:
|
||||
continue
|
||||
|
||||
def get_intersection(target_val, limit_val):
|
||||
new_min = max(target_val[0], limit_val[0])
|
||||
new_max = min(target_val[1], limit_val[1])
|
||||
return (new_min, new_max)
|
||||
|
||||
if 'lin_vel_x' in stage:
|
||||
new_range = get_intersection(stage['lin_vel_x'], hard_limit.lin_vel_x)
|
||||
active_range.lin_vel_x = new_range
|
||||
|
||||
if 'lin_vel_y' in stage:
|
||||
new_range = get_intersection(stage['lin_vel_y'], hard_limit.lin_vel_y)
|
||||
active_range.lin_vel_y = new_range
|
||||
|
||||
if 'ang_vel_yaw' in stage:
|
||||
new_range = get_intersection(stage['ang_vel_yaw'], hard_limit.ang_vel_z)
|
||||
active_range.ang_vel_z = new_range
|
||||
|
||||
if 'heading' in stage and hasattr(active_range, 'heading'):
|
||||
new_range = get_intersection(stage['heading'], hard_limit.heading)
|
||||
active_range.heading = new_range
|
||||
|
||||
schedule.pop(i)
|
||||
break
|
||||
|
||||
last_key = list(cmd_cfg.ranges.keys())[-1]
|
||||
return cmd_cfg.ranges[last_key].lin_vel_x[1]
|
||||
|
||||
def gradual_ref_stand_modification(
|
||||
env: ManagerBasedRLEnv,
|
||||
env_ids: Sequence[int],
|
||||
term_name: str,
|
||||
initial: float,
|
||||
final: float,
|
||||
start_it: int,
|
||||
end_it: int,
|
||||
):
|
||||
current_it = env.common_step_counter // 24
|
||||
if current_it < start_it:
|
||||
return
|
||||
|
||||
if current_it >= end_it:
|
||||
new = final
|
||||
else:
|
||||
new = (current_it - start_it) / (end_it - start_it) * (final - initial) + initial
|
||||
|
||||
term = env.command_manager.get_term(term_name)
|
||||
term.cfg.rel_standing_envs = new
|
||||
|
||||
|
||||
def gradual_reward_weight_modification(
|
||||
env: ManagerBasedRLEnv,
|
||||
env_ids: Sequence[int],
|
||||
term_name: str,
|
||||
initial_weight: float,
|
||||
final_weight: float,
|
||||
start_it: int,
|
||||
end_it: int,
|
||||
):
|
||||
"""Curriculum that gradually modifies a reward weight between an initial and final value over a range of steps."""
|
||||
current_it = env.common_step_counter // 24
|
||||
if current_it < start_it:
|
||||
return
|
||||
|
||||
if current_it >= end_it:
|
||||
new_weight = final_weight
|
||||
else:
|
||||
new_weight = (current_it - start_it) / (end_it - start_it) * (final_weight - initial_weight) + initial_weight
|
||||
|
||||
term_cfg = env.reward_manager.get_term_cfg(term_name)
|
||||
term_cfg.weight = new_weight
|
||||
env.reward_manager.set_term_cfg(term_name, term_cfg)
|
||||
|
||||
def terrain_levels_vel_gym(env: ManagerBasedRLEnv, env_ids: Sequence[int]) -> float:
|
||||
"""
|
||||
使用 max_move_distance 而非 reset 时的瞬间位移, 比较标准基于 commands_xy_accumulation
|
||||
"""
|
||||
terrain = env.scene.terrain
|
||||
command = env.command_manager.get_term("base_velocity")
|
||||
|
||||
max_move_dist = command.max_move_distance[env_ids]
|
||||
cmd_accum = command.commands_xy_accumulation[env_ids]
|
||||
|
||||
resampling_time = command.cfg.resampling_time_range[0]
|
||||
zero_prob = command.cfg.rel_standing_envs
|
||||
|
||||
move_up = max_move_dist > terrain.cfg.terrain_generator.size[0] / 2
|
||||
target_dist = torch.norm(cmd_accum, dim=1) * (resampling_time * (1 - zero_prob))
|
||||
move_down = (max_move_dist < target_dist * 0.5) * ~move_up
|
||||
terrain.update_env_origins(env_ids, move_up, move_down)
|
||||
|
||||
return torch.mean(terrain.terrain_levels.float()) # type: ignore
|
||||
269
source/robot_lab/robot_lab/tasks/go2/mdp/events.py
Normal file
269
source/robot_lab/robot_lab/tasks/go2/mdp/events.py
Normal file
@@ -0,0 +1,269 @@
|
||||
# Copyright (c) 2024-2025 Ziqi Fan
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
from typing import TYPE_CHECKING, Literal
|
||||
|
||||
import isaaclab.utils.math as math_utils
|
||||
from isaaclab.assets import Articulation, RigidObject
|
||||
from isaaclab.managers import SceneEntityCfg
|
||||
|
||||
from .utils import is_env_assigned_to_terrain
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from isaaclab.envs import ManagerBasedEnv
|
||||
|
||||
|
||||
def randomize_rigid_body_inertia(
|
||||
env: ManagerBasedEnv,
|
||||
env_ids: torch.Tensor | None,
|
||||
asset_cfg: SceneEntityCfg,
|
||||
inertia_distribution_params: tuple[float, float],
|
||||
operation: Literal["add", "scale", "abs"],
|
||||
distribution: Literal["uniform", "log_uniform", "gaussian"] = "uniform",
|
||||
):
|
||||
"""Randomize the inertia tensors of the bodies by adding, scaling, or setting random values.
|
||||
|
||||
This function allows randomizing only the diagonal inertia tensor components (xx, yy, zz) of the bodies.
|
||||
The function samples random values from the given distribution parameters and adds, scales, or sets the values
|
||||
into the physics simulation based on the operation.
|
||||
|
||||
.. tip::
|
||||
This function uses CPU tensors to assign the body inertias. It is recommended to use this function
|
||||
only during the initialization of the environment.
|
||||
"""
|
||||
# extract the used quantities (to enable type-hinting)
|
||||
asset: RigidObject | Articulation = env.scene[asset_cfg.name]
|
||||
|
||||
# resolve environment ids
|
||||
if env_ids is None:
|
||||
env_ids = torch.arange(env.scene.num_envs, device="cpu")
|
||||
else:
|
||||
env_ids = env_ids.cpu()
|
||||
|
||||
# resolve body indices
|
||||
if asset_cfg.body_ids == slice(None):
|
||||
body_ids = torch.arange(asset.num_bodies, dtype=torch.int, device="cpu")
|
||||
else:
|
||||
body_ids = torch.tensor(asset_cfg.body_ids, dtype=torch.int, device="cpu")
|
||||
|
||||
# get the current inertia tensors of the bodies (num_assets, num_bodies, 9 for articulations or 9 for rigid objects)
|
||||
inertias = asset.root_physx_view.get_inertias()
|
||||
|
||||
# apply randomization on default values
|
||||
inertias[env_ids[:, None], body_ids, :] = asset.data.default_inertia[env_ids[:, None], body_ids, :].clone()
|
||||
|
||||
# randomize each diagonal element (xx, yy, zz -> indices 0, 4, 8)
|
||||
for idx in [0, 4, 8]:
|
||||
# Extract and randomize the specific diagonal element
|
||||
randomized_inertias = _randomize_prop_by_op(
|
||||
inertias[:, :, idx],
|
||||
inertia_distribution_params,
|
||||
env_ids,
|
||||
body_ids,
|
||||
operation,
|
||||
distribution,
|
||||
)
|
||||
# Assign the randomized values back to the inertia tensor
|
||||
inertias[env_ids[:, None], body_ids, idx] = randomized_inertias
|
||||
|
||||
# set the inertia tensors into the physics simulation
|
||||
asset.root_physx_view.set_inertias(inertias, env_ids)
|
||||
|
||||
|
||||
def randomize_com_positions(
|
||||
env: ManagerBasedEnv,
|
||||
env_ids: torch.Tensor | None,
|
||||
asset_cfg: SceneEntityCfg,
|
||||
com_distribution_params: tuple[float, float],
|
||||
operation: Literal["add", "scale", "abs"],
|
||||
distribution: Literal["uniform", "log_uniform", "gaussian"] = "uniform",
|
||||
):
|
||||
"""Randomize the center of mass (COM) positions for the rigid bodies.
|
||||
|
||||
This function allows randomizing the COM positions of the bodies in the physics simulation. The positions can be
|
||||
randomized by adding, scaling, or setting random values sampled from the specified distribution.
|
||||
|
||||
.. tip::
|
||||
This function is intended for initialization or offline adjustments, as it modifies physics properties directly.
|
||||
|
||||
Args:
|
||||
env (ManagerBasedEnv): The simulation environment.
|
||||
env_ids (torch.Tensor | None): Specific environment indices to apply randomization, or None for all environments.
|
||||
asset_cfg (SceneEntityCfg): The configuration for the target asset whose COM will be randomized.
|
||||
com_distribution_params (tuple[float, float]): Parameters of the distribution (e.g., min and max for uniform).
|
||||
operation (Literal["add", "scale", "abs"]): The operation to apply for randomization.
|
||||
distribution (Literal["uniform", "log_uniform", "gaussian"]): The distribution to sample random values from.
|
||||
"""
|
||||
# Extract the asset (Articulation or RigidObject)
|
||||
asset: RigidObject | Articulation = env.scene[asset_cfg.name]
|
||||
|
||||
# Resolve environment indices
|
||||
if env_ids is None:
|
||||
env_ids = torch.arange(env.scene.num_envs, device="cpu")
|
||||
else:
|
||||
env_ids = env_ids.cpu()
|
||||
|
||||
# Resolve body indices
|
||||
if asset_cfg.body_ids == slice(None):
|
||||
body_ids = torch.arange(asset.num_bodies, dtype=torch.int, device="cpu")
|
||||
else:
|
||||
body_ids = torch.tensor(asset_cfg.body_ids, dtype=torch.int, device="cpu")
|
||||
|
||||
# Get the current COM offsets (num_assets, num_bodies, 3)
|
||||
com_offsets = asset.root_physx_view.get_coms()
|
||||
|
||||
for dim_idx in range(3): # Randomize x, y, z independently
|
||||
randomized_offset = _randomize_prop_by_op(
|
||||
com_offsets[:, :, dim_idx],
|
||||
com_distribution_params,
|
||||
env_ids,
|
||||
body_ids,
|
||||
operation,
|
||||
distribution,
|
||||
)
|
||||
com_offsets[env_ids[:, None], body_ids, dim_idx] = randomized_offset[env_ids[:, None], body_ids]
|
||||
|
||||
# Set the randomized COM offsets into the simulation
|
||||
asset.root_physx_view.set_coms(com_offsets, env_ids)
|
||||
|
||||
|
||||
"""
|
||||
Internal helper functions.
|
||||
"""
|
||||
|
||||
|
||||
def _randomize_prop_by_op(
|
||||
data: torch.Tensor,
|
||||
distribution_parameters: tuple[float | torch.Tensor, float | torch.Tensor],
|
||||
dim_0_ids: torch.Tensor | None,
|
||||
dim_1_ids: torch.Tensor | slice,
|
||||
operation: Literal["add", "scale", "abs"],
|
||||
distribution: Literal["uniform", "log_uniform", "gaussian"],
|
||||
) -> torch.Tensor:
|
||||
"""Perform data randomization based on the given operation and distribution.
|
||||
|
||||
Args:
|
||||
data: The data tensor to be randomized. Shape is (dim_0, dim_1).
|
||||
distribution_parameters: The parameters for the distribution to sample values from.
|
||||
dim_0_ids: The indices of the first dimension to randomize.
|
||||
dim_1_ids: The indices of the second dimension to randomize.
|
||||
operation: The operation to perform on the data. Options: 'add', 'scale', 'abs'.
|
||||
distribution: The distribution to sample the random values from. Options: 'uniform', 'log_uniform'.
|
||||
|
||||
Returns:
|
||||
The data tensor after randomization. Shape is (dim_0, dim_1).
|
||||
|
||||
Raises:
|
||||
NotImplementedError: If the operation or distribution is not supported.
|
||||
"""
|
||||
# resolve shape
|
||||
# -- dim 0
|
||||
if dim_0_ids is None:
|
||||
n_dim_0 = data.shape[0]
|
||||
dim_0_ids = slice(None)
|
||||
else:
|
||||
n_dim_0 = len(dim_0_ids)
|
||||
if not isinstance(dim_1_ids, slice):
|
||||
dim_0_ids = dim_0_ids[:, None]
|
||||
# -- dim 1
|
||||
if isinstance(dim_1_ids, slice):
|
||||
n_dim_1 = data.shape[1]
|
||||
else:
|
||||
n_dim_1 = len(dim_1_ids)
|
||||
|
||||
# resolve the distribution
|
||||
if distribution == "uniform":
|
||||
dist_fn = math_utils.sample_uniform
|
||||
elif distribution == "log_uniform":
|
||||
dist_fn = math_utils.sample_log_uniform
|
||||
elif distribution == "gaussian":
|
||||
dist_fn = math_utils.sample_gaussian
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
f"Unknown distribution: '{distribution}' for joint properties randomization."
|
||||
" Please use 'uniform', 'log_uniform', 'gaussian'."
|
||||
)
|
||||
# perform the operation
|
||||
if operation == "add":
|
||||
data[dim_0_ids, dim_1_ids] += dist_fn(*distribution_parameters, (n_dim_0, n_dim_1), device=data.device)
|
||||
elif operation == "scale":
|
||||
data[dim_0_ids, dim_1_ids] *= dist_fn(*distribution_parameters, (n_dim_0, n_dim_1), device=data.device)
|
||||
elif operation == "abs":
|
||||
data[dim_0_ids, dim_1_ids] = dist_fn(*distribution_parameters, (n_dim_0, n_dim_1), device=data.device)
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
f"Unknown operation: '{operation}' for property randomization. Please use 'add', 'scale', or 'abs'."
|
||||
)
|
||||
return data
|
||||
|
||||
|
||||
def reset_root_state_uniform(
|
||||
env: ManagerBasedEnv,
|
||||
env_ids: torch.Tensor,
|
||||
pose_range: dict[str, tuple[float, float]],
|
||||
velocity_range: dict[str, tuple[float, float]],
|
||||
asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"),
|
||||
):
|
||||
"""Reset the asset root state to a random position and velocity uniformly within the given ranges.
|
||||
|
||||
This function randomizes the root position and velocity of the asset.
|
||||
|
||||
* It samples the root position from the given ranges and adds them to the default root position, before setting
|
||||
them into the physics simulation.
|
||||
* It samples the root orientation from the given ranges and sets them into the physics simulation.
|
||||
* It samples the root velocity from the given ranges and sets them into the physics simulation.
|
||||
|
||||
The function takes a dictionary of pose and velocity ranges for each axis and rotation. The keys of the
|
||||
dictionary are ``x``, ``y``, ``z``, ``roll``, ``pitch``, and ``yaw``. The values are tuples of the form
|
||||
``(min, max)``. If the dictionary does not contain a key, the position or velocity is set to zero for that axis.
|
||||
|
||||
Note: If "pits" terrain exists, environments on pit terrain will be reset to default state without random
|
||||
perturbations to avoid the robot falling into the pit.
|
||||
"""
|
||||
# extract the used quantities (to enable type-hinting)
|
||||
asset: RigidObject | Articulation = env.scene[asset_cfg.name]
|
||||
|
||||
# Separate pit and non-pit environments
|
||||
# Check which environments are assigned to pit terrain (not random reset)
|
||||
assigned_to_pits = is_env_assigned_to_terrain(env, "pits")
|
||||
pit_env_ids = env_ids[assigned_to_pits[env_ids]]
|
||||
non_pit_env_ids = env_ids[~assigned_to_pits[env_ids]]
|
||||
|
||||
# Reset pit environments to default state (no random perturbations)
|
||||
if len(pit_env_ids) > 0:
|
||||
root_states = asset.data.default_root_state[pit_env_ids].clone()
|
||||
positions = root_states[:, 0:3] + env.scene.env_origins[pit_env_ids]
|
||||
orientations = root_states[:, 3:7]
|
||||
velocities = torch.zeros_like(root_states[:, 7:13])
|
||||
asset.write_root_pose_to_sim(torch.cat([positions, orientations], dim=-1), env_ids=pit_env_ids)
|
||||
asset.write_root_velocity_to_sim(velocities, env_ids=pit_env_ids)
|
||||
|
||||
# Reset non-pit environments with random perturbations
|
||||
if len(non_pit_env_ids) > 0:
|
||||
root_states = asset.data.default_root_state[non_pit_env_ids].clone()
|
||||
|
||||
# poses
|
||||
range_list = [pose_range.get(key, (0.0, 0.0)) for key in ["x", "y", "z", "roll", "pitch", "yaw"]]
|
||||
ranges = torch.tensor(range_list, device=asset.device)
|
||||
rand_samples = math_utils.sample_uniform(
|
||||
ranges[:, 0], ranges[:, 1], (len(non_pit_env_ids), 6), device=asset.device
|
||||
)
|
||||
|
||||
positions = root_states[:, 0:3] + env.scene.env_origins[non_pit_env_ids] + rand_samples[:, 0:3]
|
||||
orientations_delta = math_utils.quat_from_euler_xyz(rand_samples[:, 3], rand_samples[:, 4], rand_samples[:, 5])
|
||||
orientations = math_utils.quat_mul(root_states[:, 3:7], orientations_delta)
|
||||
# velocities
|
||||
range_list = [velocity_range.get(key, (0.0, 0.0)) for key in ["x", "y", "z", "roll", "pitch", "yaw"]]
|
||||
ranges = torch.tensor(range_list, device=asset.device)
|
||||
rand_samples = math_utils.sample_uniform(
|
||||
ranges[:, 0], ranges[:, 1], (len(non_pit_env_ids), 6), device=asset.device
|
||||
)
|
||||
|
||||
velocities = root_states[:, 7:13] + rand_samples
|
||||
|
||||
# set into the physics simulation
|
||||
asset.write_root_pose_to_sim(torch.cat([positions, orientations], dim=-1), env_ids=non_pit_env_ids)
|
||||
asset.write_root_velocity_to_sim(velocities, env_ids=non_pit_env_ids)
|
||||
49
source/robot_lab/robot_lab/tasks/go2/mdp/observations.py
Normal file
49
source/robot_lab/robot_lab/tasks/go2/mdp/observations.py
Normal file
@@ -0,0 +1,49 @@
|
||||
# Copyright (c) 2024-2025 Ziqi Fan
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from isaaclab.assets import Articulation
|
||||
from isaaclab.managers import SceneEntityCfg
|
||||
from isaaclab.sensors import ContactSensor
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from isaaclab.envs import ManagerBasedEnv, ManagerBasedRLEnv
|
||||
|
||||
|
||||
def joint_pos_rel_without_wheel(
|
||||
env: ManagerBasedEnv,
|
||||
asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"),
|
||||
wheel_asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"),
|
||||
) -> torch.Tensor:
|
||||
"""The joint positions of the asset w.r.t. the default joint positions.(Without the wheel joints)"""
|
||||
# extract the used quantities (to enable type-hinting)
|
||||
asset: Articulation = env.scene[asset_cfg.name]
|
||||
joint_pos_rel = asset.data.joint_pos[:, asset_cfg.joint_ids] - asset.data.default_joint_pos[:, asset_cfg.joint_ids]
|
||||
joint_pos_rel[:, wheel_asset_cfg.joint_ids] = 0
|
||||
return joint_pos_rel
|
||||
|
||||
|
||||
def phase(env: ManagerBasedRLEnv, cycle_time: float) -> torch.Tensor:
|
||||
if not hasattr(env, "episode_length_buf") or env.episode_length_buf is None:
|
||||
env.episode_length_buf = torch.zeros(env.num_envs, device=env.device, dtype=torch.long)
|
||||
phase = env.episode_length_buf[:, None] * env.step_dt / cycle_time
|
||||
phase_tensor = torch.cat([torch.sin(2 * torch.pi * phase), torch.cos(2 * torch.pi * phase)], dim=-1)
|
||||
return phase_tensor
|
||||
|
||||
def joint_acc(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
|
||||
asset: Articulation = env.scene[asset_cfg.name]
|
||||
return asset.data.joint_acc[:, asset_cfg.joint_ids]
|
||||
|
||||
def foot_contact_force_norm(env: ManagerBasedRLEnv, sensor_cfg: SceneEntityCfg) -> torch.Tensor:
|
||||
contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name]
|
||||
net_contact_forces = contact_sensor.data.net_forces_w_history # [B, T_hist, num_bodies, 3]
|
||||
|
||||
contact_force_norm = torch.norm(net_contact_forces[:, :, sensor_cfg.body_ids], dim=-1) # [B, T_hist, num_legs]
|
||||
max_contact_force_norm, _ = torch.max(contact_force_norm, dim=1) # [B, num_legs]
|
||||
contact_force_norm = torch.concat([max_contact_force_norm.unsqueeze(1), contact_force_norm], dim=1) # [B, T_hist+1, num_legs]
|
||||
|
||||
return contact_force_norm.flatten(start_dim=-2) # [B, (T_hist+1)*num_legs]
|
||||
557
source/robot_lab/robot_lab/tasks/go2/mdp/rewards.py
Normal file
557
source/robot_lab/robot_lab/tasks/go2/mdp/rewards.py
Normal file
@@ -0,0 +1,557 @@
|
||||
# Copyright (c) 2024-2025 Ziqi Fan
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import isaaclab.utils.math as math_utils
|
||||
from isaaclab.assets import Articulation, RigidObject
|
||||
from isaaclab.envs import mdp
|
||||
from isaaclab.managers import ManagerTermBase
|
||||
from isaaclab.managers import RewardTermCfg as RewTerm
|
||||
from isaaclab.managers import SceneEntityCfg
|
||||
from isaaclab.sensors import ContactSensor, RayCaster
|
||||
from isaaclab.utils.math import quat_apply_inverse, yaw_quat
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from isaaclab.envs import ManagerBasedRLEnv
|
||||
|
||||
|
||||
def track_lin_vel_xy_exp(
|
||||
env: ManagerBasedRLEnv, std: float, command_name: str, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")
|
||||
) -> torch.Tensor:
|
||||
"""Reward tracking of linear velocity commands (xy axes) using exponential kernel."""
|
||||
# extract the used quantities (to enable type-hinting)
|
||||
asset: RigidObject = env.scene[asset_cfg.name]
|
||||
# compute the error
|
||||
lin_vel_error = torch.sum(
|
||||
torch.square(env.command_manager.get_command(command_name)[:, :2] - asset.data.root_lin_vel_b[:, :2]),
|
||||
dim=1,
|
||||
)
|
||||
reward = torch.exp(-lin_vel_error / std**2)
|
||||
reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
|
||||
return reward
|
||||
|
||||
|
||||
def track_ang_vel_z_exp(
|
||||
env: ManagerBasedRLEnv, std: float, command_name: str, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")
|
||||
) -> torch.Tensor:
|
||||
"""Reward tracking of angular velocity commands (yaw) using exponential kernel."""
|
||||
# extract the used quantities (to enable type-hinting)
|
||||
asset: RigidObject = env.scene[asset_cfg.name]
|
||||
# compute the error
|
||||
ang_vel_error = torch.square(env.command_manager.get_command(command_name)[:, 2] - asset.data.root_ang_vel_b[:, 2])
|
||||
reward = torch.exp(-ang_vel_error / std**2)
|
||||
reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
|
||||
return reward
|
||||
|
||||
|
||||
def joint_power(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
|
||||
"""Reward joint_power"""
|
||||
# extract the used quantities (to enable type-hinting)
|
||||
asset: Articulation = env.scene[asset_cfg.name]
|
||||
# compute the reward
|
||||
reward = torch.sum(
|
||||
torch.abs(asset.data.joint_vel[:, asset_cfg.joint_ids] * asset.data.applied_torque[:, asset_cfg.joint_ids]),
|
||||
dim=1,
|
||||
)
|
||||
return reward
|
||||
|
||||
|
||||
def stand_still(
|
||||
env: ManagerBasedRLEnv,
|
||||
command_name: str,
|
||||
command_threshold: float = 0.06,
|
||||
asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"),
|
||||
) -> torch.Tensor:
|
||||
"""Penalize offsets from the default joint positions when the command is very small."""
|
||||
# Penalize motion when command is nearly zero.
|
||||
reward = mdp.joint_deviation_l1(env, asset_cfg)
|
||||
reward *= torch.norm(env.command_manager.get_command(command_name), dim=1) < command_threshold
|
||||
reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
|
||||
return reward
|
||||
|
||||
|
||||
def joint_pos_penalty(
|
||||
env: ManagerBasedRLEnv,
|
||||
command_name: str,
|
||||
asset_cfg: SceneEntityCfg,
|
||||
stand_still_scale: float,
|
||||
velocity_threshold: float,
|
||||
command_threshold: float,
|
||||
) -> torch.Tensor:
|
||||
"""Penalize joint position error from default on the articulation."""
|
||||
# extract the used quantities (to enable type-hinting)
|
||||
asset: Articulation = env.scene[asset_cfg.name]
|
||||
cmd = torch.linalg.norm(env.command_manager.get_command(command_name), dim=1)
|
||||
body_vel = torch.linalg.norm(asset.data.root_lin_vel_b[:, :2], dim=1)
|
||||
running_reward = torch.linalg.norm(
|
||||
(asset.data.joint_pos[:, asset_cfg.joint_ids] - asset.data.default_joint_pos[:, asset_cfg.joint_ids]), dim=1
|
||||
)
|
||||
reward = torch.where(
|
||||
torch.logical_or(cmd > command_threshold, body_vel > velocity_threshold),
|
||||
running_reward,
|
||||
stand_still_scale * running_reward,
|
||||
)
|
||||
reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
|
||||
return reward
|
||||
|
||||
|
||||
def joint_mirror(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg, mirror_joints: list[list[str]]) -> torch.Tensor:
|
||||
# extract the used quantities (to enable type-hinting)
|
||||
asset: Articulation = env.scene[asset_cfg.name]
|
||||
if not hasattr(env, "joint_mirror_joints_cache") or env.joint_mirror_joints_cache is None:
|
||||
# Cache joint positions for all pairs
|
||||
env.joint_mirror_joints_cache = [
|
||||
[asset.find_joints(joint_name) for joint_name in joint_pair] for joint_pair in mirror_joints
|
||||
]
|
||||
reward = torch.zeros(env.num_envs, device=env.device)
|
||||
# Iterate over all joint pairs
|
||||
for joint_pair in env.joint_mirror_joints_cache:
|
||||
# Calculate the difference for each pair and add to the total reward
|
||||
diff = torch.sum(
|
||||
torch.square(asset.data.joint_pos[:, joint_pair[0][0]] - asset.data.joint_pos[:, joint_pair[1][0]]),
|
||||
dim=-1,
|
||||
)
|
||||
reward += diff
|
||||
reward *= 1 / len(mirror_joints) if len(mirror_joints) > 0 else 0
|
||||
reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
|
||||
return reward
|
||||
|
||||
|
||||
def action_mirror(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg, mirror_joints: list[list[str]]) -> torch.Tensor:
|
||||
# extract the used quantities (to enable type-hinting)
|
||||
asset: Articulation = env.scene[asset_cfg.name]
|
||||
if not hasattr(env, "action_mirror_joints_cache") or env.action_mirror_joints_cache is None:
|
||||
# Cache joint positions for all pairs
|
||||
env.action_mirror_joints_cache = [
|
||||
[asset.find_joints(joint_name) for joint_name in joint_pair] for joint_pair in mirror_joints
|
||||
]
|
||||
reward = torch.zeros(env.num_envs, device=env.device)
|
||||
# Iterate over all joint pairs
|
||||
for joint_pair in env.action_mirror_joints_cache:
|
||||
# Calculate the difference for each pair and add to the total reward
|
||||
diff = torch.sum(
|
||||
torch.square(
|
||||
torch.abs(env.action_manager.action[:, joint_pair[0][0]])
|
||||
- torch.abs(env.action_manager.action[:, joint_pair[1][0]])
|
||||
),
|
||||
dim=-1,
|
||||
)
|
||||
reward += diff
|
||||
reward *= 1 / len(mirror_joints) if len(mirror_joints) > 0 else 0
|
||||
reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
|
||||
return reward
|
||||
|
||||
|
||||
def action_sync(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg, joint_groups: list[list[str]]) -> torch.Tensor:
|
||||
# extract the used quantities (to enable type-hinting)
|
||||
asset: Articulation = env.scene[asset_cfg.name]
|
||||
|
||||
# Cache joint indices if not already done
|
||||
if not hasattr(env, "action_sync_joint_cache") or env.action_sync_joint_cache is None:
|
||||
env.action_sync_joint_cache = [
|
||||
[asset.find_joints(joint_name) for joint_name in joint_group] for joint_group in joint_groups
|
||||
]
|
||||
|
||||
reward = torch.zeros(env.num_envs, device=env.device)
|
||||
# Iterate over each joint group
|
||||
for joint_group in env.action_sync_joint_cache:
|
||||
if len(joint_group) < 2:
|
||||
continue # need at least 2 joints to compare
|
||||
|
||||
# Get absolute actions for all joints in this group
|
||||
actions = torch.stack(
|
||||
[torch.abs(env.action_manager.action[:, joint[0]]) for joint in joint_group], dim=1
|
||||
) # shape: (num_envs, num_joints_in_group)
|
||||
|
||||
# Calculate mean action for each environment
|
||||
mean_actions = torch.mean(actions, dim=1, keepdim=True)
|
||||
|
||||
# Calculate variance from mean for each joint
|
||||
variance = torch.mean(torch.square(actions - mean_actions), dim=1)
|
||||
|
||||
# Add to reward (we want to minimize this variance)
|
||||
reward += variance.squeeze()
|
||||
reward *= 1 / len(joint_groups) if len(joint_groups) > 0 else 0
|
||||
reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
|
||||
return reward
|
||||
|
||||
|
||||
def feet_air_time(
|
||||
env: ManagerBasedRLEnv, command_name: str, sensor_cfg: SceneEntityCfg, threshold: float
|
||||
) -> torch.Tensor:
|
||||
"""Reward long steps taken by the feet using L2-kernel.
|
||||
|
||||
This function rewards the agent for taking steps that are longer than a threshold. This helps ensure
|
||||
that the robot lifts its feet off the ground and takes steps. The reward is computed as the sum of
|
||||
the time for which the feet are in the air.
|
||||
|
||||
If the commands are small (i.e. the agent is not supposed to take a step), then the reward is zero.
|
||||
"""
|
||||
# extract the used quantities (to enable type-hinting)
|
||||
contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name]
|
||||
# compute the reward
|
||||
first_contact = contact_sensor.compute_first_contact(env.step_dt)[:, sensor_cfg.body_ids]
|
||||
last_air_time = contact_sensor.data.last_air_time[:, sensor_cfg.body_ids]
|
||||
reward = torch.sum((last_air_time - threshold) * first_contact, dim=1)
|
||||
# no reward for zero command
|
||||
reward *= torch.norm(env.command_manager.get_command(command_name), dim=1) > 0.1
|
||||
reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
|
||||
return reward
|
||||
|
||||
|
||||
def feet_air_time_positive_biped(env, command_name: str, threshold: float, sensor_cfg: SceneEntityCfg) -> torch.Tensor:
|
||||
"""Reward long steps taken by the feet for bipeds.
|
||||
|
||||
This function rewards the agent for taking steps up to a specified threshold and also keep one foot at
|
||||
a time in the air.
|
||||
|
||||
If the commands are small (i.e. the agent is not supposed to take a step), then the reward is zero.
|
||||
"""
|
||||
contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name]
|
||||
# compute the reward
|
||||
air_time = contact_sensor.data.current_air_time[:, sensor_cfg.body_ids]
|
||||
contact_time = contact_sensor.data.current_contact_time[:, sensor_cfg.body_ids]
|
||||
in_contact = contact_time > 0.0
|
||||
in_mode_time = torch.where(in_contact, contact_time, air_time)
|
||||
single_stance = torch.sum(in_contact.int(), dim=1) == 1
|
||||
reward = torch.min(torch.where(single_stance.unsqueeze(-1), in_mode_time, 0.0), dim=1)[0]
|
||||
reward = torch.clamp(reward, max=threshold)
|
||||
# no reward for zero command
|
||||
reward *= torch.norm(env.command_manager.get_command(command_name), dim=1) > 0.1
|
||||
reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
|
||||
return reward
|
||||
|
||||
|
||||
def feet_air_time_variance_penalty(env: ManagerBasedRLEnv, sensor_cfg: SceneEntityCfg) -> torch.Tensor:
|
||||
"""Penalize variance in the amount of time each foot spends in the air/on the ground relative to each other"""
|
||||
# extract the used quantities (to enable type-hinting)
|
||||
contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name]
|
||||
# compute the reward
|
||||
last_air_time = contact_sensor.data.last_air_time[:, sensor_cfg.body_ids]
|
||||
last_contact_time = contact_sensor.data.last_contact_time[:, sensor_cfg.body_ids]
|
||||
reward = torch.var(torch.clip(last_air_time, max=0.5), dim=1) + torch.var(
|
||||
torch.clip(last_contact_time, max=0.5), dim=1
|
||||
)
|
||||
reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
|
||||
return reward
|
||||
|
||||
|
||||
def feet_contact(
|
||||
env: ManagerBasedRLEnv, command_name: str, expect_contact_num: int, sensor_cfg: SceneEntityCfg
|
||||
) -> torch.Tensor:
|
||||
"""Reward feet contact"""
|
||||
# extract the used quantities (to enable type-hinting)
|
||||
contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name]
|
||||
# compute the reward
|
||||
contact = contact_sensor.compute_first_contact(env.step_dt)[:, sensor_cfg.body_ids]
|
||||
contact_num = torch.sum(contact, dim=1)
|
||||
reward = (contact_num != expect_contact_num).float()
|
||||
# no reward for zero command
|
||||
reward *= torch.linalg.norm(env.command_manager.get_command(command_name), dim=1) > 0.1
|
||||
reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
|
||||
return reward
|
||||
|
||||
|
||||
def feet_contact_without_cmd(env: ManagerBasedRLEnv, command_name: str, sensor_cfg: SceneEntityCfg) -> torch.Tensor:
|
||||
"""Reward feet contact"""
|
||||
# extract the used quantities (to enable type-hinting)
|
||||
contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name]
|
||||
# compute the reward
|
||||
contact = contact_sensor.compute_first_contact(env.step_dt)[:, sensor_cfg.body_ids]
|
||||
reward = torch.sum(contact, dim=-1).float()
|
||||
reward *= torch.linalg.norm(env.command_manager.get_command(command_name), dim=1) < 0.1
|
||||
reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
|
||||
return reward
|
||||
|
||||
|
||||
def feet_stumble(env: ManagerBasedRLEnv, sensor_cfg: SceneEntityCfg) -> torch.Tensor:
|
||||
# extract the used quantities (to enable type-hinting)
|
||||
contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name]
|
||||
forces_z = torch.abs(contact_sensor.data.net_forces_w[:, sensor_cfg.body_ids, 2])
|
||||
forces_xy = torch.linalg.norm(contact_sensor.data.net_forces_w[:, sensor_cfg.body_ids, :2], dim=2)
|
||||
# Penalize feet hitting vertical surfaces
|
||||
reward = torch.any(forces_xy > 4 * forces_z, dim=1).float()
|
||||
reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
|
||||
return reward
|
||||
|
||||
|
||||
def feet_distance_y_exp(
|
||||
env: ManagerBasedRLEnv, stance_width: float, std: float, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")
|
||||
) -> torch.Tensor:
|
||||
asset: RigidObject = env.scene[asset_cfg.name]
|
||||
cur_footsteps_translated = asset.data.body_link_pos_w[:, asset_cfg.body_ids, :] - asset.data.root_link_pos_w[
|
||||
:, :
|
||||
].unsqueeze(1)
|
||||
n_feet = len(asset_cfg.body_ids)
|
||||
footsteps_in_body_frame = torch.zeros(env.num_envs, n_feet, 3, device=env.device)
|
||||
for i in range(n_feet):
|
||||
footsteps_in_body_frame[:, i, :] = math_utils.quat_apply(
|
||||
math_utils.quat_conjugate(asset.data.root_link_quat_w), cur_footsteps_translated[:, i, :]
|
||||
)
|
||||
side_sign = torch.tensor(
|
||||
[1.0 if i % 2 == 0 else -1.0 for i in range(n_feet)],
|
||||
device=env.device,
|
||||
)
|
||||
stance_width_tensor = stance_width * torch.ones([env.num_envs, 1], device=env.device)
|
||||
desired_ys = stance_width_tensor / 2 * side_sign.unsqueeze(0)
|
||||
stance_diff = torch.square(desired_ys - footsteps_in_body_frame[:, :, 1])
|
||||
reward = torch.exp(-torch.sum(stance_diff, dim=1) / (std**2))
|
||||
reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
|
||||
return reward
|
||||
|
||||
|
||||
def feet_distance_xy_exp(
|
||||
env: ManagerBasedRLEnv,
|
||||
stance_width: float,
|
||||
stance_length: float,
|
||||
std: float,
|
||||
asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"),
|
||||
) -> torch.Tensor:
|
||||
asset: RigidObject = env.scene[asset_cfg.name]
|
||||
|
||||
# Compute the current footstep positions relative to the root
|
||||
cur_footsteps_translated = asset.data.body_link_pos_w[:, asset_cfg.body_ids, :] - asset.data.root_link_pos_w[
|
||||
:, :
|
||||
].unsqueeze(1)
|
||||
|
||||
footsteps_in_body_frame = torch.zeros(env.num_envs, 4, 3, device=env.device)
|
||||
for i in range(4):
|
||||
footsteps_in_body_frame[:, i, :] = math_utils.quat_apply(
|
||||
math_utils.quat_conjugate(asset.data.root_link_quat_w), cur_footsteps_translated[:, i, :]
|
||||
)
|
||||
|
||||
# Desired x and y positions for each foot
|
||||
stance_width_tensor = stance_width * torch.ones([env.num_envs, 1], device=env.device)
|
||||
stance_length_tensor = stance_length * torch.ones([env.num_envs, 1], device=env.device)
|
||||
|
||||
desired_xs = torch.cat(
|
||||
[stance_length_tensor / 2, stance_length_tensor / 2, -stance_length_tensor / 2, -stance_length_tensor / 2],
|
||||
dim=1,
|
||||
)
|
||||
desired_ys = torch.cat(
|
||||
[stance_width_tensor / 2, -stance_width_tensor / 2, stance_width_tensor / 2, -stance_width_tensor / 2], dim=1
|
||||
)
|
||||
|
||||
# Compute differences in x and y
|
||||
stance_diff_x = torch.square(desired_xs - footsteps_in_body_frame[:, :, 0])
|
||||
stance_diff_y = torch.square(desired_ys - footsteps_in_body_frame[:, :, 1])
|
||||
|
||||
# Combine x and y differences and compute the exponential penalty
|
||||
stance_diff = stance_diff_x + stance_diff_y
|
||||
reward = torch.exp(-torch.sum(stance_diff, dim=1) / std**2)
|
||||
reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
|
||||
return reward
|
||||
|
||||
|
||||
def feet_height(
|
||||
env: ManagerBasedRLEnv,
|
||||
command_name: str,
|
||||
asset_cfg: SceneEntityCfg,
|
||||
target_height: float,
|
||||
tanh_mult: float,
|
||||
) -> torch.Tensor:
|
||||
"""Reward the swinging feet for clearing a specified height off the ground"""
|
||||
asset: RigidObject = env.scene[asset_cfg.name]
|
||||
foot_z_target_error = torch.square(asset.data.body_pos_w[:, asset_cfg.body_ids, 2] - target_height)
|
||||
foot_velocity_tanh = torch.tanh(
|
||||
tanh_mult * torch.linalg.norm(asset.data.body_lin_vel_w[:, asset_cfg.body_ids, :2], dim=2)
|
||||
)
|
||||
reward = torch.sum(foot_z_target_error * foot_velocity_tanh, dim=1)
|
||||
# no reward for zero command
|
||||
reward *= torch.linalg.norm(env.command_manager.get_command(command_name), dim=1) > 0.1
|
||||
reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
|
||||
return reward
|
||||
|
||||
|
||||
def feet_height_body(
|
||||
env: ManagerBasedRLEnv,
|
||||
command_name: str,
|
||||
asset_cfg: SceneEntityCfg,
|
||||
target_height: float,
|
||||
tanh_mult: float,
|
||||
) -> torch.Tensor:
|
||||
"""Reward the swinging feet for clearing a specified height off the ground"""
|
||||
asset: RigidObject = env.scene[asset_cfg.name]
|
||||
cur_footpos_translated = asset.data.body_pos_w[:, asset_cfg.body_ids, :] - asset.data.root_pos_w[:, :].unsqueeze(1)
|
||||
footpos_in_body_frame = torch.zeros(env.num_envs, len(asset_cfg.body_ids), 3, device=env.device)
|
||||
cur_footvel_translated = asset.data.body_lin_vel_w[:, asset_cfg.body_ids, :] - asset.data.root_lin_vel_w[
|
||||
:, :
|
||||
].unsqueeze(1)
|
||||
footvel_in_body_frame = torch.zeros(env.num_envs, len(asset_cfg.body_ids), 3, device=env.device)
|
||||
for i in range(len(asset_cfg.body_ids)):
|
||||
footpos_in_body_frame[:, i, :] = math_utils.quat_apply_inverse(
|
||||
asset.data.root_quat_w, cur_footpos_translated[:, i, :]
|
||||
)
|
||||
footvel_in_body_frame[:, i, :] = math_utils.quat_apply_inverse(
|
||||
asset.data.root_quat_w, cur_footvel_translated[:, i, :]
|
||||
)
|
||||
foot_z_target_error = torch.square(footpos_in_body_frame[:, :, 2] - target_height).view(env.num_envs, -1)
|
||||
foot_velocity_tanh = torch.tanh(tanh_mult * torch.norm(footvel_in_body_frame[:, :, :2], dim=2))
|
||||
reward = torch.sum(foot_z_target_error * foot_velocity_tanh, dim=1)
|
||||
reward *= torch.linalg.norm(env.command_manager.get_command(command_name), dim=1) > 0.1
|
||||
reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
|
||||
return reward
|
||||
|
||||
|
||||
def feet_slide(
|
||||
env: ManagerBasedRLEnv, sensor_cfg: SceneEntityCfg, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")
|
||||
) -> torch.Tensor:
|
||||
"""Penalize feet sliding.
|
||||
|
||||
This function penalizes the agent for sliding its feet on the ground. The reward is computed as the
|
||||
norm of the linear velocity of the feet multiplied by a binary contact sensor. This ensures that the
|
||||
agent is penalized only when the feet are in contact with the ground.
|
||||
"""
|
||||
# Penalize feet sliding
|
||||
contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name]
|
||||
contacts = contact_sensor.data.net_forces_w_history[:, :, sensor_cfg.body_ids, :].norm(dim=-1).max(dim=1)[0] > 1.0
|
||||
asset: RigidObject = env.scene[asset_cfg.name]
|
||||
|
||||
# feet_vel = asset.data.body_lin_vel_w[:, asset_cfg.body_ids, :2]
|
||||
# reward = torch.sum(feet_vel.norm(dim=-1) * contacts, dim=1)
|
||||
|
||||
cur_footvel_translated = asset.data.body_lin_vel_w[:, asset_cfg.body_ids, :] - asset.data.root_lin_vel_w[
|
||||
:, :
|
||||
].unsqueeze(1)
|
||||
footvel_in_body_frame = torch.zeros(env.num_envs, len(asset_cfg.body_ids), 3, device=env.device)
|
||||
for i in range(len(asset_cfg.body_ids)):
|
||||
footvel_in_body_frame[:, i, :] = math_utils.quat_apply_inverse(
|
||||
asset.data.root_quat_w, cur_footvel_translated[:, i, :]
|
||||
)
|
||||
foot_leteral_vel = torch.sqrt(torch.sum(torch.square(footvel_in_body_frame[:, :, :2]), dim=2)).view(
|
||||
env.num_envs, -1
|
||||
)
|
||||
reward = torch.sum(foot_leteral_vel * contacts, dim=1)
|
||||
reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
|
||||
return reward
|
||||
|
||||
|
||||
# def smoothness_1(env: ManagerBasedRLEnv) -> torch.Tensor:
|
||||
# # Penalize changes in actions
|
||||
# diff = torch.square(env.action_manager.action - env.action_manager.prev_action)
|
||||
# diff = diff * (env.action_manager.prev_action[:, :] != 0) # ignore first step
|
||||
# return torch.sum(diff, dim=1)
|
||||
|
||||
|
||||
# def smoothness_2(env: ManagerBasedRLEnv) -> torch.Tensor:
|
||||
# # Penalize changes in actions
|
||||
# diff = torch.square(env.action_manager.action - 2 * env.action_manager.prev_action + env.action_manager.prev_prev_action)
|
||||
# diff = diff * (env.action_manager.prev_action[:, :] != 0) # ignore first step
|
||||
# diff = diff * (env.action_manager.prev_prev_action[:, :] != 0) # ignore second step
|
||||
# return torch.sum(diff, dim=1)
|
||||
|
||||
def action_smoothness_l2(env: ManagerBasedRLEnv) -> torch.Tensor:
|
||||
# Penalize changes in actions
|
||||
diff = torch.square(env.action_manager.action - 2 * env.action_manager.prev_action + env.action_manager.prev_prev_action)
|
||||
diff = diff * (env.action_manager.prev_action[:, :] != 0) # ignore first step
|
||||
diff = diff * (env.action_manager.prev_prev_action[:, :] != 0) # ignore second step
|
||||
return torch.sum(diff, dim=1)
|
||||
|
||||
|
||||
def upward(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
|
||||
"""Penalize z-axis base linear velocity using L2 squared kernel."""
|
||||
# extract the used quantities (to enable type-hinting)
|
||||
asset: RigidObject = env.scene[asset_cfg.name]
|
||||
reward = torch.square(1 - asset.data.projected_gravity_b[:, 2])
|
||||
return reward
|
||||
|
||||
|
||||
def base_height_l2(
|
||||
env: ManagerBasedRLEnv,
|
||||
target_height: float,
|
||||
asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"),
|
||||
sensor_cfg: SceneEntityCfg | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""Penalize asset height from its target using L2 squared kernel.
|
||||
|
||||
Note:
|
||||
For flat terrain, target height is in the world frame. For rough terrain,
|
||||
sensor readings can adjust the target height to account for the terrain.
|
||||
"""
|
||||
# extract the used quantities (to enable type-hinting)
|
||||
asset: RigidObject = env.scene[asset_cfg.name]
|
||||
if sensor_cfg is not None:
|
||||
sensor: RayCaster = env.scene[sensor_cfg.name]
|
||||
# Adjust the target height using the sensor data
|
||||
ray_hits = sensor.data.ray_hits_w[..., 2]
|
||||
if torch.isnan(ray_hits).any() or torch.isinf(ray_hits).any() or torch.max(torch.abs(ray_hits)) > 1e6:
|
||||
adjusted_target_height = asset.data.root_link_pos_w[:, 2]
|
||||
else:
|
||||
adjusted_target_height = target_height + torch.mean(ray_hits, dim=1)
|
||||
else:
|
||||
# Use the provided target height directly for flat terrain
|
||||
adjusted_target_height = target_height
|
||||
# Compute the L2 squared penalty
|
||||
reward = torch.square(asset.data.root_pos_w[:, 2] - adjusted_target_height)
|
||||
reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
|
||||
return reward
|
||||
|
||||
|
||||
def lin_vel_z_l2(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
|
||||
"""Penalize z-axis base linear velocity using L2 squared kernel."""
|
||||
# extract the used quantities (to enable type-hinting)
|
||||
asset: RigidObject = env.scene[asset_cfg.name]
|
||||
reward = torch.square(asset.data.root_lin_vel_b[:, 2])
|
||||
reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
|
||||
return reward
|
||||
|
||||
|
||||
def ang_vel_xy_l2(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
|
||||
"""Penalize xy-axis base angular velocity using L2 squared kernel."""
|
||||
# extract the used quantities (to enable type-hinting)
|
||||
asset: RigidObject = env.scene[asset_cfg.name]
|
||||
reward = torch.sum(torch.square(asset.data.root_ang_vel_b[:, :2]), dim=1)
|
||||
reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
|
||||
return reward
|
||||
|
||||
|
||||
def undesired_contacts(env: ManagerBasedRLEnv, threshold: float, sensor_cfg: SceneEntityCfg) -> torch.Tensor:
|
||||
"""Penalize undesired contacts as the number of violations that are above a threshold."""
|
||||
# extract the used quantities (to enable type-hinting)
|
||||
contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name]
|
||||
# check if contact force is above threshold
|
||||
net_contact_forces = contact_sensor.data.net_forces_w_history
|
||||
is_contact = torch.max(torch.norm(net_contact_forces[:, :, sensor_cfg.body_ids], dim=-1), dim=1)[0] > threshold
|
||||
# sum over contacts for each environment
|
||||
reward = torch.sum(is_contact, dim=1).float()
|
||||
reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
|
||||
return reward
|
||||
|
||||
|
||||
def flat_orientation_l2(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
|
||||
"""Penalize non-flat base orientation using L2 squared kernel.
|
||||
|
||||
This is computed by penalizing the xy-components of the projected gravity vector.
|
||||
"""
|
||||
# extract the used quantities (to enable type-hinting)
|
||||
asset: RigidObject = env.scene[asset_cfg.name]
|
||||
reward = torch.sum(torch.square(asset.data.projected_gravity_b[:, :2]), dim=1)
|
||||
reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
|
||||
return reward
|
||||
|
||||
|
||||
def hip_pos_penalty(
|
||||
env: ManagerBasedRLEnv,
|
||||
command_name: str,
|
||||
asset_cfg: SceneEntityCfg,
|
||||
stand_still_scale: float,
|
||||
command_threshold: float,
|
||||
) -> torch.Tensor:
|
||||
"""Penalize joint position error from default on the articulation."""
|
||||
# extract the used quantities (to enable type-hinting)
|
||||
asset: Articulation = env.scene[asset_cfg.name]
|
||||
command = env.command_manager.get_command(command_name)[:, [1, 2]]
|
||||
cmd_large = torch.any(torch.abs(command) > command_threshold, dim=1)
|
||||
running_reward = torch.linalg.norm(
|
||||
(asset.data.joint_pos[:, asset_cfg.joint_ids] - asset.data.default_joint_pos[:, asset_cfg.joint_ids]), dim=1
|
||||
)
|
||||
reward = torch.where(
|
||||
cmd_large,
|
||||
running_reward,
|
||||
stand_still_scale * running_reward
|
||||
)
|
||||
return reward
|
||||
126
source/robot_lab/robot_lab/tasks/go2/mdp/utils.py
Normal file
126
source/robot_lab/robot_lab/tasks/go2/mdp/utils.py
Normal file
@@ -0,0 +1,126 @@
|
||||
# Copyright (c) 2024-2025 Ziqi Fan
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
"""Utility functions for terrain-aware operations."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from isaaclab.envs import ManagerBasedEnv
|
||||
|
||||
|
||||
def _get_terrain_column_range(terrain_cfg, terrain_name: str, device) -> tuple[int, int] | None:
|
||||
"""Helper function to calculate column range for a terrain type.
|
||||
|
||||
Args:
|
||||
terrain_cfg: The terrain generator configuration.
|
||||
terrain_name: Name of the terrain.
|
||||
device: Torch device.
|
||||
|
||||
Returns:
|
||||
Tuple of (col_start, col_end) or None if terrain not found.
|
||||
"""
|
||||
if terrain_cfg.sub_terrains is None or terrain_name not in terrain_cfg.sub_terrains:
|
||||
return None
|
||||
|
||||
sub_terrain_names = list(terrain_cfg.sub_terrains.keys())
|
||||
proportions = torch.tensor([sub_cfg.proportion for sub_cfg in terrain_cfg.sub_terrains.values()], device=device)
|
||||
proportions = proportions / proportions.sum()
|
||||
cumsum_props = torch.cumsum(proportions, dim=0)
|
||||
|
||||
terrain_idx = sub_terrain_names.index(terrain_name)
|
||||
# Use round() instead of int() to properly allocate columns
|
||||
col_start = round((0.0 if terrain_idx == 0 else cumsum_props[terrain_idx - 1].item()) * terrain_cfg.num_cols)
|
||||
col_end = round(cumsum_props[terrain_idx].item() * terrain_cfg.num_cols)
|
||||
|
||||
return (col_start, col_end)
|
||||
|
||||
|
||||
def is_env_assigned_to_terrain(env: ManagerBasedEnv, terrain_name: str) -> torch.Tensor:
|
||||
"""Check which environments are initially assigned to the specified terrain type.
|
||||
|
||||
Each environment is assigned to a specific terrain cell at initialization.
|
||||
This function returns a mask indicating which environments were assigned to the given terrain type.
|
||||
|
||||
Args:
|
||||
env: The environment instance.
|
||||
terrain_name: Name of the terrain to check (e.g., "pits", "stairs").
|
||||
|
||||
Returns:
|
||||
Boolean tensor of shape (num_envs,) where True means the environment is assigned to this terrain.
|
||||
"""
|
||||
# Check if terrain and terrain generator are available
|
||||
terrain = getattr(env.scene, "terrain", None)
|
||||
if terrain is None or not hasattr(terrain, "terrain_types"):
|
||||
return torch.zeros(env.num_envs, dtype=torch.bool, device=env.device)
|
||||
if terrain.cfg.terrain_type != "generator" or terrain.cfg.terrain_generator is None:
|
||||
return torch.zeros(env.num_envs, dtype=torch.bool, device=env.device)
|
||||
|
||||
terrain_cfg = terrain.cfg.terrain_generator
|
||||
col_range = _get_terrain_column_range(terrain_cfg, terrain_name, env.device)
|
||||
if col_range is None:
|
||||
return torch.zeros(env.num_envs, dtype=torch.bool, device=env.device)
|
||||
|
||||
col_start, col_end = col_range
|
||||
# terrain_types directly stores column indices, so just check if they're in range
|
||||
return (terrain.terrain_types >= col_start) & (terrain.terrain_types < col_end)
|
||||
|
||||
|
||||
def is_robot_on_terrain(env: ManagerBasedEnv, terrain_name: str, asset_name: str = "robot") -> torch.Tensor:
|
||||
"""Check which robots are currently standing on the specified terrain type.
|
||||
|
||||
This function calculates which terrain grid cell each robot is on based on its world position,
|
||||
then checks if that cell's terrain type matches the specified terrain.
|
||||
|
||||
Args:
|
||||
env: The environment instance.
|
||||
terrain_name: Name of the terrain to check (e.g., "pits", "stairs").
|
||||
asset_name: Name of the robot asset. Defaults to "robot".
|
||||
|
||||
Returns:
|
||||
Boolean tensor of shape (num_envs,) where True means the robot is currently on this terrain.
|
||||
"""
|
||||
# Check if terrain and terrain generator are available
|
||||
terrain = getattr(env.scene, "terrain", None)
|
||||
if terrain is None or not hasattr(terrain, "terrain_types"):
|
||||
return torch.zeros(env.num_envs, dtype=torch.bool, device=env.device)
|
||||
if terrain.cfg.terrain_type != "generator" or terrain.cfg.terrain_generator is None:
|
||||
return torch.zeros(env.num_envs, dtype=torch.bool, device=env.device)
|
||||
|
||||
terrain_cfg = terrain.cfg.terrain_generator
|
||||
col_range = _get_terrain_column_range(terrain_cfg, terrain_name, env.device)
|
||||
if col_range is None:
|
||||
return torch.zeros(env.num_envs, dtype=torch.bool, device=env.device)
|
||||
|
||||
col_start, col_end = col_range
|
||||
|
||||
# Get robot positions in world frame
|
||||
asset = env.scene[asset_name]
|
||||
robot_pos_w = asset.data.root_pos_w[:, :2] # [num_envs, 2] (x, y)
|
||||
|
||||
# Get terrain grid information
|
||||
terrain_origins = terrain.terrain_origins # [num_rows, num_cols, 3]
|
||||
num_rows, num_cols, _ = terrain_origins.shape
|
||||
|
||||
# Use terrain_origins to directly compute which cell each robot is in
|
||||
# terrain_origins[r, c, :2] is the center of cell (r, c)
|
||||
# We need to find the closest terrain origin for each robot
|
||||
|
||||
# Reshape terrain_origins for distance calculation
|
||||
terrain_origins_2d = terrain_origins[:, :, :2].reshape(num_rows * num_cols, 2) # [num_rows*num_cols, 2]
|
||||
|
||||
# Calculate distances from each robot to all terrain origins
|
||||
distances = torch.cdist(robot_pos_w, terrain_origins_2d) # [num_envs, num_rows*num_cols]
|
||||
|
||||
# Find the closest terrain origin for each robot
|
||||
closest_flat_idx = torch.argmin(distances, dim=1) # [num_envs]
|
||||
|
||||
# Convert flat index to column index
|
||||
# flat_idx = row * num_cols + col
|
||||
col_idx = closest_flat_idx % num_cols # [num_envs]
|
||||
|
||||
# Check if the robot's current terrain column is in the specified terrain's range
|
||||
return (col_idx >= col_start) & (col_idx < col_end)
|
||||
86
source/robot_lab/robot_lab/tasks/go2/rsl_rl_cfg.py
Normal file
86
source/robot_lab/robot_lab/tasks/go2/rsl_rl_cfg.py
Normal file
@@ -0,0 +1,86 @@
|
||||
from isaaclab.utils import configclass
|
||||
from isaaclab_rl.rsl_rl import RslRlOnPolicyRunnerCfg, RslRlPpoActorCriticCfg, RslRlPpoAlgorithmCfg
|
||||
|
||||
@configclass
|
||||
class PPORunnerCfg(RslRlOnPolicyRunnerCfg):
|
||||
num_steps_per_env = 24
|
||||
max_iterations = 150000
|
||||
save_interval = 500
|
||||
experiment_name = "go2_rough"
|
||||
|
||||
policy = RslRlPpoActorCriticCfg(
|
||||
init_noise_std=1.0,
|
||||
actor_obs_normalization=False,
|
||||
critic_obs_normalization=False,
|
||||
actor_hidden_dims=[512, 256, 128],
|
||||
critic_hidden_dims=[512, 256, 128],
|
||||
activation="elu",
|
||||
)
|
||||
algorithm = RslRlPpoAlgorithmCfg(
|
||||
value_loss_coef=1.0,
|
||||
use_clipped_value_loss=True,
|
||||
clip_param=0.2,
|
||||
entropy_coef=0.01,
|
||||
num_learning_epochs=5,
|
||||
num_mini_batches=4,
|
||||
learning_rate=1.0e-3,
|
||||
schedule="adaptive",
|
||||
gamma=0.99,
|
||||
lam=0.95,
|
||||
desired_kl=0.01,
|
||||
max_grad_norm=1.0,
|
||||
)
|
||||
|
||||
|
||||
@configclass
|
||||
class RslRlMoeCtsActorCriticCfg(RslRlPpoActorCriticCfg):
|
||||
class_name = "ActorCriticMoECTS"
|
||||
init_noise_std = 1.0
|
||||
expert_num = 8 # number of experts in the student model
|
||||
latent_dim = 32
|
||||
norm_type = 'l2norm' # normalization type for encoders: l2norm, simnorm
|
||||
teacher_encoder_hidden_dims = [512, 256, 128]
|
||||
student_encoder_hidden_dims = [512, 256, 128]
|
||||
actor_hidden_dims=[512, 256, 128]
|
||||
critic_hidden_dims=[512, 256, 128]
|
||||
activation="elu"
|
||||
actor_obs_normalization = False
|
||||
critic_obs_normalization = False
|
||||
|
||||
@configclass
|
||||
class RslRlMoeCtsAlgorithmCfg(RslRlPpoAlgorithmCfg):
|
||||
class_name = "MoECTS"
|
||||
value_loss_coef = 1.0
|
||||
load_balance_coef = 0.01 # coefficient for load balance loss
|
||||
use_clipped_value_loss = True
|
||||
clip_param = 0.2
|
||||
entropy_coef = 0.01
|
||||
num_learning_epochs = 5
|
||||
num_mini_batches = 4
|
||||
learning_rate = 1e-3
|
||||
student_encoder_learning_rate = 1e-3
|
||||
schedule = "adaptive"
|
||||
gamma = 0.99
|
||||
lam = 0.95
|
||||
betas = (0.9, 0.999)
|
||||
weight_decay = 0.0
|
||||
desired_kl = 0.01
|
||||
max_grad_norm = 1.0
|
||||
teacher_env_ratio = 0.75 # percentage of envs assigned to teacher
|
||||
|
||||
@configclass
|
||||
class MoECTSRunnerCfg(RslRlOnPolicyRunnerCfg):
|
||||
experiment_name = "go2_moe_cts"
|
||||
class_name = "OnPolicyRunnerCTS"
|
||||
num_steps_per_env = 24
|
||||
max_iterations = 150000
|
||||
save_interval = 500
|
||||
policy = RslRlMoeCtsActorCriticCfg()
|
||||
algorithm = RslRlMoeCtsAlgorithmCfg()
|
||||
|
||||
# concat elu inspired by concat relu from https://arxiv.org/pdf/2303.07507
|
||||
@configclass
|
||||
class MoECTSCatELURunnerCfg(MoECTSRunnerCfg):
|
||||
def __post_init__(self):
|
||||
super().__post_init__()
|
||||
self.policy.activation = 'cat_elu'
|
||||
49
source/robot_lab/robot_lab/ui_extension_example.py
Normal file
49
source/robot_lab/robot_lab/ui_extension_example.py
Normal file
@@ -0,0 +1,49 @@
|
||||
# Copyright (c) 2024-2025 Ziqi Fan
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# Copyright (c) 2022-2025, The Isaac Lab Project Developers.
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
|
||||
import omni.ext
|
||||
|
||||
|
||||
# Functions and vars are available to other extension as usual in python: `example.python_ext.some_public_function(x)`
|
||||
def some_public_function(x: int):
|
||||
print("[robot_lab] some_public_function was called with x: ", x)
|
||||
return x**x
|
||||
|
||||
|
||||
# Any class derived from `omni.ext.IExt` in top level module (defined in `python.modules` of `extension.toml`) will be
|
||||
# instantiated when extension gets enabled and `on_startup(ext_id)` will be called. Later when extension gets disabled
|
||||
# on_shutdown() is called.
|
||||
class ExampleExtension(omni.ext.IExt):
|
||||
# ext_id is current extension id. It can be used with extension manager to query additional information, like where
|
||||
# this extension is located on filesystem.
|
||||
def on_startup(self, ext_id):
|
||||
print("[robot_lab] startup")
|
||||
|
||||
self._count = 0
|
||||
|
||||
self._window = omni.ui.Window("My Window", width=300, height=300)
|
||||
with self._window.frame:
|
||||
with omni.ui.VStack():
|
||||
label = omni.ui.Label("")
|
||||
|
||||
def on_click():
|
||||
self._count += 1
|
||||
label.text = f"count: {self._count}"
|
||||
|
||||
def on_reset():
|
||||
self._count = 0
|
||||
label.text = "empty"
|
||||
|
||||
on_reset()
|
||||
|
||||
with omni.ui.HStack():
|
||||
omni.ui.Button("Add", clicked_fn=on_click)
|
||||
omni.ui.Button("Reset", clicked_fn=on_reset)
|
||||
|
||||
def on_shutdown(self):
|
||||
print("[robot_lab] shutdown")
|
||||
53
source/robot_lab/setup.py
Normal file
53
source/robot_lab/setup.py
Normal file
@@ -0,0 +1,53 @@
|
||||
# Copyright (c) 2024-2025 Ziqi Fan
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
"""Installation script for the 'robot_lab' python package."""
|
||||
|
||||
import os
|
||||
import toml
|
||||
|
||||
from setuptools import setup
|
||||
|
||||
# Obtain the extension data from the extension.toml file
|
||||
EXTENSION_PATH = os.path.dirname(os.path.realpath(__file__))
|
||||
# Read the extension.toml file
|
||||
EXTENSION_TOML_DATA = toml.load(os.path.join(EXTENSION_PATH, "config", "extension.toml"))
|
||||
|
||||
# Minimum dependencies required prior to installation
|
||||
INSTALL_REQUIRES = [
|
||||
# base
|
||||
"psutil",
|
||||
"colorama",
|
||||
"xacrodoc",
|
||||
# amp
|
||||
"numpy",
|
||||
"pandas",
|
||||
"pinocchio",
|
||||
# rl
|
||||
"cusrl[all]",
|
||||
]
|
||||
|
||||
# Installation operation
|
||||
setup(
|
||||
name="robot_lab",
|
||||
packages=["robot_lab"],
|
||||
author=EXTENSION_TOML_DATA["package"]["author"],
|
||||
maintainer=EXTENSION_TOML_DATA["package"]["maintainer"],
|
||||
url=EXTENSION_TOML_DATA["package"]["repository"],
|
||||
version=EXTENSION_TOML_DATA["package"]["version"],
|
||||
description=EXTENSION_TOML_DATA["package"]["description"],
|
||||
keywords=EXTENSION_TOML_DATA["package"]["keywords"],
|
||||
install_requires=INSTALL_REQUIRES,
|
||||
license="Apache License 2.0",
|
||||
include_package_data=True,
|
||||
python_requires=">=3.10",
|
||||
classifiers=[
|
||||
"Natural Language :: English",
|
||||
"Programming Language :: Python :: 3.10",
|
||||
"Programming Language :: Python :: 3.11",
|
||||
"Isaac Sim :: 4.5.0",
|
||||
"Isaac Sim :: 5.0.0",
|
||||
"Isaac Sim :: 5.1.0",
|
||||
],
|
||||
zip_safe=False,
|
||||
)
|
||||
14
source/rsl_rl/.gitignore
vendored
Normal file
14
source/rsl_rl/.gitignore
vendored
Normal file
@@ -0,0 +1,14 @@
|
||||
# IDEs
|
||||
.idea
|
||||
|
||||
# builds
|
||||
*.egg-info
|
||||
build/*
|
||||
dist/*
|
||||
|
||||
# cache
|
||||
__pycache__
|
||||
.pytest_cache
|
||||
|
||||
# vs code
|
||||
.vscode
|
||||
33
source/rsl_rl/.pre-commit-config.yaml
Normal file
33
source/rsl_rl/.pre-commit-config.yaml
Normal file
@@ -0,0 +1,33 @@
|
||||
repos:
|
||||
- repo: https://github.com/astral-sh/ruff-pre-commit
|
||||
rev: v0.14.0
|
||||
hooks:
|
||||
- id: ruff-check
|
||||
- id: ruff-format
|
||||
- repo: https://github.com/pre-commit/pre-commit-hooks
|
||||
rev: v4.5.0
|
||||
hooks:
|
||||
- id: check-symlinks
|
||||
- id: destroyed-symlinks
|
||||
- id: check-yaml
|
||||
- id: check-toml
|
||||
- id: check-merge-conflict
|
||||
- id: check-case-conflict
|
||||
- id: check-executables-have-shebangs
|
||||
- id: check-shebang-scripts-are-executable
|
||||
- id: detect-private-key
|
||||
- repo: https://github.com/codespell-project/codespell
|
||||
rev: v2.2.6
|
||||
hooks:
|
||||
- id: codespell
|
||||
additional_dependencies:
|
||||
- tomli
|
||||
- repo: https://github.com/Lucas-C/pre-commit-hooks
|
||||
rev: v1.5.1
|
||||
hooks:
|
||||
- id: insert-license
|
||||
files: \.py$
|
||||
args:
|
||||
# - --remove-header # Remove existing license headers. Useful when updating license.
|
||||
- --license-filepath
|
||||
- .github/LICENSE_HEADER.txt
|
||||
38
source/rsl_rl/CITATION.cff
Normal file
38
source/rsl_rl/CITATION.cff
Normal file
@@ -0,0 +1,38 @@
|
||||
cff-version: 1.2.0
|
||||
title: "RSL-RL: A Learning Library for Robotics Research"
|
||||
message: "If you use this work, please cite the following paper."
|
||||
repository-code: "https://github.com/leggedrobotics/rsl_rl"
|
||||
license: BSD-3-Clause
|
||||
version: 3.3.0
|
||||
type: software
|
||||
authors:
|
||||
- family-names: Schwarke
|
||||
given-names: Clemens
|
||||
- family-names: Mittal
|
||||
given-names: Mayank
|
||||
- family-names: Rudin
|
||||
given-names: Nikita
|
||||
- family-names: Hoeller
|
||||
given-names: David
|
||||
keywords:
|
||||
- reinforcement learning
|
||||
- robotics
|
||||
- control
|
||||
- RSL-RL
|
||||
preferred-citation:
|
||||
type: article
|
||||
authors:
|
||||
- family-names: Schwarke
|
||||
given-names: Clemens
|
||||
- family-names: Mittal
|
||||
given-names: Mayank
|
||||
- family-names: Rudin
|
||||
given-names: Nikita
|
||||
- family-names: Hoeller
|
||||
given-names: David
|
||||
- family-names: Hutter
|
||||
given-names: Marco
|
||||
title: "RSL-RL: A Learning Library for Robotics Research"
|
||||
journal: "arXiv preprint"
|
||||
doi: 10.48550/arXiv.2509.10771
|
||||
url: "https://arxiv.org/abs/2509.10771"
|
||||
42
source/rsl_rl/CONTRIBUTORS.md
Normal file
42
source/rsl_rl/CONTRIBUTORS.md
Normal file
@@ -0,0 +1,42 @@
|
||||
# RSL-RL Maintainers and Contributors
|
||||
|
||||
This is the official list of developers and contributors.
|
||||
|
||||
To see the full list of contributors, see the revision history in the source control.
|
||||
|
||||
Names should be added to this file as: individual names or organizations.
|
||||
|
||||
Email addresses are tracked elsewhere to avoid spam.
|
||||
|
||||
Please keep the lists sorted alphabetically.
|
||||
|
||||
## Maintainers
|
||||
|
||||
* Robotic Systems Lab, ETH Zurich
|
||||
* NVIDIA Corporation
|
||||
|
||||
---
|
||||
|
||||
* Clemens Schwarke
|
||||
* Mayank Mittal
|
||||
|
||||
## Authors
|
||||
|
||||
* Clemens Schwarke
|
||||
* David Hoeller
|
||||
* Mayank Mittal
|
||||
* Nikita Rudin
|
||||
|
||||
## Contributors
|
||||
|
||||
* Bikram Pandit
|
||||
* Eric Vollenweider
|
||||
* Fabian Jenelten
|
||||
* Lorenzo Terenzi
|
||||
* Marko Bjelonic
|
||||
* Matthijs van der Boon
|
||||
* Özhan Özen
|
||||
* Pascal Roth
|
||||
* Shaoshu Su
|
||||
* Zhang Chong
|
||||
* Ziqi Fan
|
||||
30
source/rsl_rl/LICENSE
Normal file
30
source/rsl_rl/LICENSE
Normal file
@@ -0,0 +1,30 @@
|
||||
Copyright (c) 2026, ETH Zurich
|
||||
Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES
|
||||
All rights reserved.
|
||||
|
||||
Redistribution and use in source and binary forms, with or without modification,
|
||||
are permitted provided that the following conditions are met:
|
||||
|
||||
1. Redistributions of source code must retain the above copyright notice,
|
||||
this list of conditions and the following disclaimer.
|
||||
|
||||
2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
this list of conditions and the following disclaimer in the documentation
|
||||
and/or other materials provided with the distribution.
|
||||
|
||||
3. Neither the name of the copyright holder nor the names of its contributors
|
||||
may be used to endorse or promote products derived from this software without
|
||||
specific prior written permission.
|
||||
|
||||
THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
|
||||
ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
|
||||
WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR
|
||||
ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
|
||||
(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
|
||||
LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
|
||||
ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
|
||||
(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
|
||||
SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
|
||||
See licenses/dependencies for license information of dependencies of this package.
|
||||
109
source/rsl_rl/README.md
Normal file
109
source/rsl_rl/README.md
Normal file
@@ -0,0 +1,109 @@
|
||||
# RSL-RL
|
||||
|
||||
A fast and simple implementation of learning algorithms for robotics. For an overview of the library please have a look at https://arxiv.org/pdf/2509.10771.
|
||||
|
||||
Environment repositories using the framework:
|
||||
|
||||
* **`Isaac Lab`** (built on top of NVIDIA Isaac Sim): https://github.com/isaac-sim/IsaacLab
|
||||
* **`Legged Gym`** (built on top of NVIDIA Isaac Gym): https://leggedrobotics.github.io/legged_gym/
|
||||
* **`MuJoCo Playground`** (built on top of MuJoCo MJX and Warp): https://github.com/google-deepmind/mujoco_playground/
|
||||
* **`mjlab`** (built on top of MuJoCo Warp): https://github.com/mujocolab/mjlab
|
||||
|
||||
The library currently supports **PPO** and **Student-Teacher Distillation** with additional features from our research. These include:
|
||||
|
||||
* [Random Network Distillation (RND)](https://proceedings.mlr.press/v229/schwarke23a.html) - Encourages exploration by adding
|
||||
a curiosity driven intrinsic reward.
|
||||
* [Symmetry-based Augmentation](https://arxiv.org/abs/2403.04359) - Makes the learned behaviors more symmetrical.
|
||||
|
||||
We welcome contributions from the community. Please check our contribution guidelines for more
|
||||
information.
|
||||
|
||||
**Maintainer**: Mayank Mittal and Clemens Schwarke <br/>
|
||||
**Affiliation**: Robotic Systems Lab, ETH Zurich & NVIDIA <br/>
|
||||
**Contact**: cschwarke@ethz.ch
|
||||
|
||||
|
||||
## Setup
|
||||
|
||||
The package can be installed via PyPI with:
|
||||
|
||||
```bash
|
||||
pip install rsl-rl-lib
|
||||
```
|
||||
|
||||
or by cloning this repository and installing it with:
|
||||
|
||||
```bash
|
||||
git clone https://github.com/leggedrobotics/rsl_rl
|
||||
cd rsl_rl
|
||||
pip install -e .
|
||||
```
|
||||
|
||||
The package supports the following logging frameworks which can be configured through `logger`:
|
||||
|
||||
* Tensorboard: https://www.tensorflow.org/tensorboard/
|
||||
* Weights & Biases: https://wandb.ai/site
|
||||
* Neptune: https://docs.neptune.ai/
|
||||
|
||||
For a demo configuration of PPO, please check the [example_config.yaml](config/example_config.yaml) file.
|
||||
|
||||
|
||||
## Contribution Guidelines
|
||||
|
||||
For documentation, we adopt the [Google Style Guide](https://sphinxcontrib-napoleon.readthedocs.io/en/latest/example_google.html) for docstrings. Please make sure that your code is well-documented and follows the guidelines.
|
||||
|
||||
We use the following tools for maintaining code quality:
|
||||
|
||||
- [pre-commit](https://pre-commit.com/): Runs a list of formatters and linters over the codebase.
|
||||
- [ruff](https://github.com/astral-sh/ruff): An extremely fast Python linter and code formatter, written in Rust.
|
||||
|
||||
Please check [here](https://pre-commit.com/#install) for instructions to set these up. To run over the entire repository, please execute the following command in the terminal:
|
||||
|
||||
```bash
|
||||
# for installation (only once)
|
||||
pre-commit install
|
||||
# for running
|
||||
pre-commit run --all-files
|
||||
```
|
||||
|
||||
## Citing
|
||||
|
||||
If you use this library for your research, please cite the following work:
|
||||
|
||||
```text
|
||||
@article{schwarke2025rslrl,
|
||||
title={RSL-RL: A Learning Library for Robotics Research},
|
||||
author={Schwarke, Clemens and Mittal, Mayank and Rudin, Nikita and Hoeller, David and Hutter, Marco},
|
||||
journal={arXiv preprint arXiv:2509.10771},
|
||||
year={2025}
|
||||
}
|
||||
```
|
||||
|
||||
If you use the library with curiosity-driven exploration (random network distillation), please cite:
|
||||
|
||||
```text
|
||||
@InProceedings{schwarke2023curiosity,
|
||||
title = {Curiosity-Driven Learning of Joint Locomotion and Manipulation Tasks},
|
||||
author = {Schwarke, Clemens and Klemm, Victor and Boon, Matthijs van der and Bjelonic, Marko and Hutter, Marco},
|
||||
booktitle = {Proceedings of The 7th Conference on Robot Learning},
|
||||
pages = {2594--2610},
|
||||
year = {2023},
|
||||
volume = {229},
|
||||
series = {Proceedings of Machine Learning Research},
|
||||
publisher = {PMLR},
|
||||
url = {https://proceedings.mlr.press/v229/schwarke23a.html},
|
||||
}
|
||||
```
|
||||
|
||||
If you use the library with symmetry augmentation, please cite:
|
||||
|
||||
```text
|
||||
@InProceedings{mittal2024symmetry,
|
||||
author={Mittal, Mayank and Rudin, Nikita and Klemm, Victor and Allshire, Arthur and Hutter, Marco},
|
||||
booktitle={2024 IEEE International Conference on Robotics and Automation (ICRA)},
|
||||
title={Symmetry Considerations for Learning Task Symmetric Robot Policies},
|
||||
year={2024},
|
||||
pages={7433-7439},
|
||||
doi={10.1109/ICRA57147.2024.10611493}
|
||||
}
|
||||
```
|
||||
77
source/rsl_rl/config/example_config.yaml
Normal file
77
source/rsl_rl/config/example_config.yaml
Normal file
@@ -0,0 +1,77 @@
|
||||
runner:
|
||||
class_name: OnPolicyRunner
|
||||
# General
|
||||
num_steps_per_env: 24 # Number of steps per environment per iteration
|
||||
max_iterations: 1500 # Number of policy updates
|
||||
seed: 1
|
||||
# Observations
|
||||
obs_groups: {"policy": ["policy"], "critic": ["policy", "privileged"]} # Maps observation groups to sets. See `vec_env.py` for more information
|
||||
# Logging parameters
|
||||
save_interval: 50 # Check for potential saves every `save_interval` iterations
|
||||
experiment_name: walking_experiment
|
||||
run_name: ""
|
||||
# Logging writer
|
||||
logger: tensorboard # tensorboard, neptune, wandb
|
||||
neptune_project: legged_gym
|
||||
wandb_project: legged_gym
|
||||
|
||||
# Policy
|
||||
policy:
|
||||
class_name: ActorCritic
|
||||
activation: elu
|
||||
actor_obs_normalization: false
|
||||
critic_obs_normalization: false
|
||||
actor_hidden_dims: [256, 256, 256]
|
||||
critic_hidden_dims: [256, 256, 256]
|
||||
init_noise_std: 1.0
|
||||
noise_std_type: "scalar" # 'scalar' or 'log'
|
||||
state_dependent_std: false
|
||||
|
||||
# Algorithm
|
||||
algorithm:
|
||||
class_name: PPO
|
||||
# Training
|
||||
learning_rate: 0.001
|
||||
num_learning_epochs: 5
|
||||
num_mini_batches: 4 # mini batch size = num_envs * num_steps / num_mini_batches
|
||||
schedule: adaptive # adaptive, fixed
|
||||
# Value function
|
||||
value_loss_coef: 1.0
|
||||
clip_param: 0.2
|
||||
use_clipped_value_loss: true
|
||||
# Surrogate loss
|
||||
desired_kl: 0.01
|
||||
entropy_coef: 0.01
|
||||
gamma: 0.99
|
||||
lam: 0.95
|
||||
max_grad_norm: 1.0
|
||||
# Miscellaneous
|
||||
normalize_advantage_per_mini_batch: false
|
||||
|
||||
# Random network distillation
|
||||
rnd_cfg:
|
||||
weight: 0.0 # Initial weight of the RND reward
|
||||
weight_schedule: null # This is a dictionary with a required key called "mode". Please check the RND module for more information
|
||||
reward_normalization: false # Whether to normalize RND reward
|
||||
# Learning parameters
|
||||
learning_rate: 0.001 # Learning rate for RND
|
||||
# Network parameters
|
||||
num_outputs: 1 # Number of outputs of RND network. Note: if -1, then the network will use dimensions of the observation
|
||||
predictor_hidden_dims: [-1] # Hidden dimensions of predictor network
|
||||
target_hidden_dims: [-1] # Hidden dimensions of target network
|
||||
|
||||
# Symmetry augmentation
|
||||
symmetry_cfg:
|
||||
use_data_augmentation: true # This adds symmetric trajectories to the batch
|
||||
use_mirror_loss: false # This adds symmetry loss term to the loss function
|
||||
data_augmentation_func: null # String containing the module and function name to import
|
||||
# Example: "legged_gym.envs.locomotion.anymal_c.symmetry:get_symmetric_states"
|
||||
#
|
||||
# .. code-block:: python
|
||||
#
|
||||
# @torch.no_grad()
|
||||
# def get_symmetric_states(
|
||||
# obs: Optional[torch.Tensor] = None, actions: Optional[torch.Tensor] = None, cfg: "BaseEnvCfg" = None, obs_type: str = "policy"
|
||||
# ) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
#
|
||||
mirror_loss_coeff: 0.0 # Coefficient for symmetry loss term. If 0, no symmetry loss is used
|
||||
339
source/rsl_rl/licenses/dependencies/codespell-license.txt
Normal file
339
source/rsl_rl/licenses/dependencies/codespell-license.txt
Normal file
@@ -0,0 +1,339 @@
|
||||
GNU GENERAL PUBLIC LICENSE
|
||||
Version 2, June 1991
|
||||
|
||||
Copyright (C) 1989, 1991 Free Software Foundation, Inc.,
|
||||
51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
|
||||
Everyone is permitted to copy and distribute verbatim copies
|
||||
of this license document, but changing it is not allowed.
|
||||
|
||||
Preamble
|
||||
|
||||
The licenses for most software are designed to take away your
|
||||
freedom to share and change it. By contrast, the GNU General Public
|
||||
License is intended to guarantee your freedom to share and change free
|
||||
software--to make sure the software is free for all its users. This
|
||||
General Public License applies to most of the Free Software
|
||||
Foundation's software and to any other program whose authors commit to
|
||||
using it. (Some other Free Software Foundation software is covered by
|
||||
the GNU Lesser General Public License instead.) You can apply it to
|
||||
your programs, too.
|
||||
|
||||
When we speak of free software, we are referring to freedom, not
|
||||
price. Our General Public Licenses are designed to make sure that you
|
||||
have the freedom to distribute copies of free software (and charge for
|
||||
this service if you wish), that you receive source code or can get it
|
||||
if you want it, that you can change the software or use pieces of it
|
||||
in new free programs; and that you know you can do these things.
|
||||
|
||||
To protect your rights, we need to make restrictions that forbid
|
||||
anyone to deny you these rights or to ask you to surrender the rights.
|
||||
These restrictions translate to certain responsibilities for you if you
|
||||
distribute copies of the software, or if you modify it.
|
||||
|
||||
For example, if you distribute copies of such a program, whether
|
||||
gratis or for a fee, you must give the recipients all the rights that
|
||||
you have. You must make sure that they, too, receive or can get the
|
||||
source code. And you must show them these terms so they know their
|
||||
rights.
|
||||
|
||||
We protect your rights with two steps: (1) copyright the software, and
|
||||
(2) offer you this license which gives you legal permission to copy,
|
||||
distribute and/or modify the software.
|
||||
|
||||
Also, for each author's protection and ours, we want to make certain
|
||||
that everyone understands that there is no warranty for this free
|
||||
software. If the software is modified by someone else and passed on, we
|
||||
want its recipients to know that what they have is not the original, so
|
||||
that any problems introduced by others will not reflect on the original
|
||||
authors' reputations.
|
||||
|
||||
Finally, any free program is threatened constantly by software
|
||||
patents. We wish to avoid the danger that redistributors of a free
|
||||
program will individually obtain patent licenses, in effect making the
|
||||
program proprietary. To prevent this, we have made it clear that any
|
||||
patent must be licensed for everyone's free use or not licensed at all.
|
||||
|
||||
The precise terms and conditions for copying, distribution and
|
||||
modification follow.
|
||||
|
||||
GNU GENERAL PUBLIC LICENSE
|
||||
TERMS AND CONDITIONS FOR COPYING, DISTRIBUTION AND MODIFICATION
|
||||
|
||||
0. This License applies to any program or other work which contains
|
||||
a notice placed by the copyright holder saying it may be distributed
|
||||
under the terms of this General Public License. The "Program", below,
|
||||
refers to any such program or work, and a "work based on the Program"
|
||||
means either the Program or any derivative work under copyright law:
|
||||
that is to say, a work containing the Program or a portion of it,
|
||||
either verbatim or with modifications and/or translated into another
|
||||
language. (Hereinafter, translation is included without limitation in
|
||||
the term "modification".) Each licensee is addressed as "you".
|
||||
|
||||
Activities other than copying, distribution and modification are not
|
||||
covered by this License; they are outside its scope. The act of
|
||||
running the Program is not restricted, and the output from the Program
|
||||
is covered only if its contents constitute a work based on the
|
||||
Program (independent of having been made by running the Program).
|
||||
Whether that is true depends on what the Program does.
|
||||
|
||||
1. You may copy and distribute verbatim copies of the Program's
|
||||
source code as you receive it, in any medium, provided that you
|
||||
conspicuously and appropriately publish on each copy an appropriate
|
||||
copyright notice and disclaimer of warranty; keep intact all the
|
||||
notices that refer to this License and to the absence of any warranty;
|
||||
and give any other recipients of the Program a copy of this License
|
||||
along with the Program.
|
||||
|
||||
You may charge a fee for the physical act of transferring a copy, and
|
||||
you may at your option offer warranty protection in exchange for a fee.
|
||||
|
||||
2. You may modify your copy or copies of the Program or any portion
|
||||
of it, thus forming a work based on the Program, and copy and
|
||||
distribute such modifications or work under the terms of Section 1
|
||||
above, provided that you also meet all of these conditions:
|
||||
|
||||
a) You must cause the modified files to carry prominent notices
|
||||
stating that you changed the files and the date of any change.
|
||||
|
||||
b) You must cause any work that you distribute or publish, that in
|
||||
whole or in part contains or is derived from the Program or any
|
||||
part thereof, to be licensed as a whole at no charge to all third
|
||||
parties under the terms of this License.
|
||||
|
||||
c) If the modified program normally reads commands interactively
|
||||
when run, you must cause it, when started running for such
|
||||
interactive use in the most ordinary way, to print or display an
|
||||
announcement including an appropriate copyright notice and a
|
||||
notice that there is no warranty (or else, saying that you provide
|
||||
a warranty) and that users may redistribute the program under
|
||||
these conditions, and telling the user how to view a copy of this
|
||||
License. (Exception: if the Program itself is interactive but
|
||||
does not normally print such an announcement, your work based on
|
||||
the Program is not required to print an announcement.)
|
||||
|
||||
These requirements apply to the modified work as a whole. If
|
||||
identifiable sections of that work are not derived from the Program,
|
||||
and can be reasonably considered independent and separate works in
|
||||
themselves, then this License, and its terms, do not apply to those
|
||||
sections when you distribute them as separate works. But when you
|
||||
distribute the same sections as part of a whole which is a work based
|
||||
on the Program, the distribution of the whole must be on the terms of
|
||||
this License, whose permissions for other licensees extend to the
|
||||
entire whole, and thus to each and every part regardless of who wrote it.
|
||||
|
||||
Thus, it is not the intent of this section to claim rights or contest
|
||||
your rights to work written entirely by you; rather, the intent is to
|
||||
exercise the right to control the distribution of derivative or
|
||||
collective works based on the Program.
|
||||
|
||||
In addition, mere aggregation of another work not based on the Program
|
||||
with the Program (or with a work based on the Program) on a volume of
|
||||
a storage or distribution medium does not bring the other work under
|
||||
the scope of this License.
|
||||
|
||||
3. You may copy and distribute the Program (or a work based on it,
|
||||
under Section 2) in object code or executable form under the terms of
|
||||
Sections 1 and 2 above provided that you also do one of the following:
|
||||
|
||||
a) Accompany it with the complete corresponding machine-readable
|
||||
source code, which must be distributed under the terms of Sections
|
||||
1 and 2 above on a medium customarily used for software interchange; or,
|
||||
|
||||
b) Accompany it with a written offer, valid for at least three
|
||||
years, to give any third party, for a charge no more than your
|
||||
cost of physically performing source distribution, a complete
|
||||
machine-readable copy of the corresponding source code, to be
|
||||
distributed under the terms of Sections 1 and 2 above on a medium
|
||||
customarily used for software interchange; or,
|
||||
|
||||
c) Accompany it with the information you received as to the offer
|
||||
to distribute corresponding source code. (This alternative is
|
||||
allowed only for noncommercial distribution and only if you
|
||||
received the program in object code or executable form with such
|
||||
an offer, in accord with Subsection b above.)
|
||||
|
||||
The source code for a work means the preferred form of the work for
|
||||
making modifications to it. For an executable work, complete source
|
||||
code means all the source code for all modules it contains, plus any
|
||||
associated interface definition files, plus the scripts used to
|
||||
control compilation and installation of the executable. However, as a
|
||||
special exception, the source code distributed need not include
|
||||
anything that is normally distributed (in either source or binary
|
||||
form) with the major components (compiler, kernel, and so on) of the
|
||||
operating system on which the executable runs, unless that component
|
||||
itself accompanies the executable.
|
||||
|
||||
If distribution of executable or object code is made by offering
|
||||
access to copy from a designated place, then offering equivalent
|
||||
access to copy the source code from the same place counts as
|
||||
distribution of the source code, even though third parties are not
|
||||
compelled to copy the source along with the object code.
|
||||
|
||||
4. You may not copy, modify, sublicense, or distribute the Program
|
||||
except as expressly provided under this License. Any attempt
|
||||
otherwise to copy, modify, sublicense or distribute the Program is
|
||||
void, and will automatically terminate your rights under this License.
|
||||
However, parties who have received copies, or rights, from you under
|
||||
this License will not have their licenses terminated so long as such
|
||||
parties remain in full compliance.
|
||||
|
||||
5. You are not required to accept this License, since you have not
|
||||
signed it. However, nothing else grants you permission to modify or
|
||||
distribute the Program or its derivative works. These actions are
|
||||
prohibited by law if you do not accept this License. Therefore, by
|
||||
modifying or distributing the Program (or any work based on the
|
||||
Program), you indicate your acceptance of this License to do so, and
|
||||
all its terms and conditions for copying, distributing or modifying
|
||||
the Program or works based on it.
|
||||
|
||||
6. Each time you redistribute the Program (or any work based on the
|
||||
Program), the recipient automatically receives a license from the
|
||||
original licensor to copy, distribute or modify the Program subject to
|
||||
these terms and conditions. You may not impose any further
|
||||
restrictions on the recipients' exercise of the rights granted herein.
|
||||
You are not responsible for enforcing compliance by third parties to
|
||||
this License.
|
||||
|
||||
7. If, as a consequence of a court judgment or allegation of patent
|
||||
infringement or for any other reason (not limited to patent issues),
|
||||
conditions are imposed on you (whether by court order, agreement or
|
||||
otherwise) that contradict the conditions of this License, they do not
|
||||
excuse you from the conditions of this License. If you cannot
|
||||
distribute so as to satisfy simultaneously your obligations under this
|
||||
License and any other pertinent obligations, then as a consequence you
|
||||
may not distribute the Program at all. For example, if a patent
|
||||
license would not permit royalty-free redistribution of the Program by
|
||||
all those who receive copies directly or indirectly through you, then
|
||||
the only way you could satisfy both it and this License would be to
|
||||
refrain entirely from distribution of the Program.
|
||||
|
||||
If any portion of this section is held invalid or unenforceable under
|
||||
any particular circumstance, the balance of the section is intended to
|
||||
apply and the section as a whole is intended to apply in other
|
||||
circumstances.
|
||||
|
||||
It is not the purpose of this section to induce you to infringe any
|
||||
patents or other property right claims or to contest validity of any
|
||||
such claims; this section has the sole purpose of protecting the
|
||||
integrity of the free software distribution system, which is
|
||||
implemented by public license practices. Many people have made
|
||||
generous contributions to the wide range of software distributed
|
||||
through that system in reliance on consistent application of that
|
||||
system; it is up to the author/donor to decide if he or she is willing
|
||||
to distribute software through any other system and a licensee cannot
|
||||
impose that choice.
|
||||
|
||||
This section is intended to make thoroughly clear what is believed to
|
||||
be a consequence of the rest of this License.
|
||||
|
||||
8. If the distribution and/or use of the Program is restricted in
|
||||
certain countries either by patents or by copyrighted interfaces, the
|
||||
original copyright holder who places the Program under this License
|
||||
may add an explicit geographical distribution limitation excluding
|
||||
those countries, so that distribution is permitted only in or among
|
||||
countries not thus excluded. In such case, this License incorporates
|
||||
the limitation as if written in the body of this License.
|
||||
|
||||
9. The Free Software Foundation may publish revised and/or new versions
|
||||
of the General Public License from time to time. Such new versions will
|
||||
be similar in spirit to the present version, but may differ in detail to
|
||||
address new problems or concerns.
|
||||
|
||||
Each version is given a distinguishing version number. If the Program
|
||||
specifies a version number of this License which applies to it and "any
|
||||
later version", you have the option of following the terms and conditions
|
||||
either of that version or of any later version published by the Free
|
||||
Software Foundation. If the Program does not specify a version number of
|
||||
this License, you may choose any version ever published by the Free Software
|
||||
Foundation.
|
||||
|
||||
10. If you wish to incorporate parts of the Program into other free
|
||||
programs whose distribution conditions are different, write to the author
|
||||
to ask for permission. For software which is copyrighted by the Free
|
||||
Software Foundation, write to the Free Software Foundation; we sometimes
|
||||
make exceptions for this. Our decision will be guided by the two goals
|
||||
of preserving the free status of all derivatives of our free software and
|
||||
of promoting the sharing and reuse of software generally.
|
||||
|
||||
NO WARRANTY
|
||||
|
||||
11. BECAUSE THE PROGRAM IS LICENSED FREE OF CHARGE, THERE IS NO WARRANTY
|
||||
FOR THE PROGRAM, TO THE EXTENT PERMITTED BY APPLICABLE LAW. EXCEPT WHEN
|
||||
OTHERWISE STATED IN WRITING THE COPYRIGHT HOLDERS AND/OR OTHER PARTIES
|
||||
PROVIDE THE PROGRAM "AS IS" WITHOUT WARRANTY OF ANY KIND, EITHER EXPRESSED
|
||||
OR IMPLIED, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
|
||||
MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE. THE ENTIRE RISK AS
|
||||
TO THE QUALITY AND PERFORMANCE OF THE PROGRAM IS WITH YOU. SHOULD THE
|
||||
PROGRAM PROVE DEFECTIVE, YOU ASSUME THE COST OF ALL NECESSARY SERVICING,
|
||||
REPAIR OR CORRECTION.
|
||||
|
||||
12. IN NO EVENT UNLESS REQUIRED BY APPLICABLE LAW OR AGREED TO IN WRITING
|
||||
WILL ANY COPYRIGHT HOLDER, OR ANY OTHER PARTY WHO MAY MODIFY AND/OR
|
||||
REDISTRIBUTE THE PROGRAM AS PERMITTED ABOVE, BE LIABLE TO YOU FOR DAMAGES,
|
||||
INCLUDING ANY GENERAL, SPECIAL, INCIDENTAL OR CONSEQUENTIAL DAMAGES ARISING
|
||||
OUT OF THE USE OR INABILITY TO USE THE PROGRAM (INCLUDING BUT NOT LIMITED
|
||||
TO LOSS OF DATA OR DATA BEING RENDERED INACCURATE OR LOSSES SUSTAINED BY
|
||||
YOU OR THIRD PARTIES OR A FAILURE OF THE PROGRAM TO OPERATE WITH ANY OTHER
|
||||
PROGRAMS), EVEN IF SUCH HOLDER OR OTHER PARTY HAS BEEN ADVISED OF THE
|
||||
POSSIBILITY OF SUCH DAMAGES.
|
||||
|
||||
END OF TERMS AND CONDITIONS
|
||||
|
||||
How to Apply These Terms to Your New Programs
|
||||
|
||||
If you develop a new program, and you want it to be of the greatest
|
||||
possible use to the public, the best way to achieve this is to make it
|
||||
free software which everyone can redistribute and change under these terms.
|
||||
|
||||
To do so, attach the following notices to the program. It is safest
|
||||
to attach them to the start of each source file to most effectively
|
||||
convey the exclusion of warranty; and each file should have at least
|
||||
the "copyright" line and a pointer to where the full notice is found.
|
||||
|
||||
<one line to give the program's name and a brief idea of what it does.>
|
||||
Copyright (C) <year> <name of author>
|
||||
|
||||
This program is free software; you can redistribute it and/or modify
|
||||
it under the terms of the GNU General Public License as published by
|
||||
the Free Software Foundation; either version 2 of the License, or
|
||||
(at your option) any later version.
|
||||
|
||||
This program is distributed in the hope that it will be useful,
|
||||
but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||||
GNU General Public License for more details.
|
||||
|
||||
You should have received a copy of the GNU General Public License along
|
||||
with this program; if not, write to the Free Software Foundation, Inc.,
|
||||
51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
|
||||
|
||||
Also add information on how to contact you by electronic and paper mail.
|
||||
|
||||
If the program is interactive, make it output a short notice like this
|
||||
when it starts in an interactive mode:
|
||||
|
||||
Gnomovision version 69, Copyright (C) year name of author
|
||||
Gnomovision comes with ABSOLUTELY NO WARRANTY; for details type `show w'.
|
||||
This is free software, and you are welcome to redistribute it
|
||||
under certain conditions; type `show c' for details.
|
||||
|
||||
The hypothetical commands `show w' and `show c' should show the appropriate
|
||||
parts of the General Public License. Of course, the commands you use may
|
||||
be called something other than `show w' and `show c'; they could even be
|
||||
mouse-clicks or menu items--whatever suits your program.
|
||||
|
||||
You should also get your employer (if you work as a programmer) or your
|
||||
school, if any, to sign a "copyright disclaimer" for the program, if
|
||||
necessary. Here is a sample; alter the names:
|
||||
|
||||
Yoyodyne, Inc., hereby disclaims all copyright interest in the program
|
||||
`Gnomovision' (which makes passes at compilers) written by James Hacker.
|
||||
|
||||
<signature of Ty Coon>, 1 April 1989
|
||||
Ty Coon, President of Vice
|
||||
|
||||
This General Public License does not permit incorporating your program into
|
||||
proprietary programs. If your program is a subroutine library, you may
|
||||
consider it more useful to permit linking proprietary applications with the
|
||||
library. If this is what you want to do, use the GNU Lesser General
|
||||
Public License instead of this License.
|
||||
30
source/rsl_rl/licenses/dependencies/numpy-license.txt
Normal file
30
source/rsl_rl/licenses/dependencies/numpy-license.txt
Normal file
@@ -0,0 +1,30 @@
|
||||
Copyright (c) 2005-2021, NumPy Developers.
|
||||
All rights reserved.
|
||||
|
||||
Redistribution and use in source and binary forms, with or without
|
||||
modification, are permitted provided that the following conditions are
|
||||
met:
|
||||
|
||||
* Redistributions of source code must retain the above copyright
|
||||
notice, this list of conditions and the following disclaimer.
|
||||
|
||||
* Redistributions in binary form must reproduce the above
|
||||
copyright notice, this list of conditions and the following
|
||||
disclaimer in the documentation and/or other materials provided
|
||||
with the distribution.
|
||||
|
||||
* Neither the name of the NumPy Developers nor the names of any
|
||||
contributors may be used to endorse or promote products derived
|
||||
from this software without specific prior written permission.
|
||||
|
||||
THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
|
||||
"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
|
||||
LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
|
||||
A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
|
||||
OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
|
||||
SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
|
||||
LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
|
||||
DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
|
||||
THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
|
||||
(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
202
source/rsl_rl/licenses/dependencies/onnx-license.txt
Normal file
202
source/rsl_rl/licenses/dependencies/onnx-license.txt
Normal file
@@ -0,0 +1,202 @@
|
||||
|
||||
Apache License
|
||||
Version 2.0, January 2004
|
||||
http://www.apache.org/licenses/
|
||||
|
||||
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
|
||||
|
||||
1. Definitions.
|
||||
|
||||
"License" shall mean the terms and conditions for use, reproduction,
|
||||
and distribution as defined by Sections 1 through 9 of this document.
|
||||
|
||||
"Licensor" shall mean the copyright owner or entity authorized by
|
||||
the copyright owner that is granting the License.
|
||||
|
||||
"Legal Entity" shall mean the union of the acting entity and all
|
||||
other entities that control, are controlled by, or are under common
|
||||
control with that entity. For the purposes of this definition,
|
||||
"control" means (i) the power, direct or indirect, to cause the
|
||||
direction or management of such entity, whether by contract or
|
||||
otherwise, or (ii) ownership of fifty percent (50%) or more of the
|
||||
outstanding shares, or (iii) beneficial ownership of such entity.
|
||||
|
||||
"You" (or "Your") shall mean an individual or Legal Entity
|
||||
exercising permissions granted by this License.
|
||||
|
||||
"Source" form shall mean the preferred form for making modifications,
|
||||
including but not limited to software source code, documentation
|
||||
source, and configuration files.
|
||||
|
||||
"Object" form shall mean any form resulting from mechanical
|
||||
transformation or translation of a Source form, including but
|
||||
not limited to compiled object code, generated documentation,
|
||||
and conversions to other media types.
|
||||
|
||||
"Work" shall mean the work of authorship, whether in Source or
|
||||
Object form, made available under the License, as indicated by a
|
||||
copyright notice that is included in or attached to the work
|
||||
(an example is provided in the Appendix below).
|
||||
|
||||
"Derivative Works" shall mean any work, whether in Source or Object
|
||||
form, that is based on (or derived from) the Work and for which the
|
||||
editorial revisions, annotations, elaborations, or other modifications
|
||||
represent, as a whole, an original work of authorship. For the purposes
|
||||
of this License, Derivative Works shall not include works that remain
|
||||
separable from, or merely link (or bind by name) to the interfaces of,
|
||||
the Work and Derivative Works thereof.
|
||||
|
||||
"Contribution" shall mean any work of authorship, including
|
||||
the original version of the Work and any modifications or additions
|
||||
to that Work or Derivative Works thereof, that is intentionally
|
||||
submitted to Licensor for inclusion in the Work by the copyright owner
|
||||
or by an individual or Legal Entity authorized to submit on behalf of
|
||||
the copyright owner. For the purposes of this definition, "submitted"
|
||||
means any form of electronic, verbal, or written communication sent
|
||||
to the Licensor or its representatives, including but not limited to
|
||||
communication on electronic mailing lists, source code control systems,
|
||||
and issue tracking systems that are managed by, or on behalf of, the
|
||||
Licensor for the purpose of discussing and improving the Work, but
|
||||
excluding communication that is conspicuously marked or otherwise
|
||||
designated in writing by the copyright owner as "Not a Contribution."
|
||||
|
||||
"Contributor" shall mean Licensor and any individual or Legal Entity
|
||||
on behalf of whom a Contribution has been received by Licensor and
|
||||
subsequently incorporated within the Work.
|
||||
|
||||
2. Grant of Copyright License. Subject to the terms and conditions of
|
||||
this License, each Contributor hereby grants to You a perpetual,
|
||||
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||
copyright license to reproduce, prepare Derivative Works of,
|
||||
publicly display, publicly perform, sublicense, and distribute the
|
||||
Work and such Derivative Works in Source or Object form.
|
||||
|
||||
3. Grant of Patent License. Subject to the terms and conditions of
|
||||
this License, each Contributor hereby grants to You a perpetual,
|
||||
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||
(except as stated in this section) patent license to make, have made,
|
||||
use, offer to sell, sell, import, and otherwise transfer the Work,
|
||||
where such license applies only to those patent claims licensable
|
||||
by such Contributor that are necessarily infringed by their
|
||||
Contribution(s) alone or by combination of their Contribution(s)
|
||||
with the Work to which such Contribution(s) was submitted. If You
|
||||
institute patent litigation against any entity (including a
|
||||
cross-claim or counterclaim in a lawsuit) alleging that the Work
|
||||
or a Contribution incorporated within the Work constitutes direct
|
||||
or contributory patent infringement, then any patent licenses
|
||||
granted to You under this License for that Work shall terminate
|
||||
as of the date such litigation is filed.
|
||||
|
||||
4. Redistribution. You may reproduce and distribute copies of the
|
||||
Work or Derivative Works thereof in any medium, with or without
|
||||
modifications, and in Source or Object form, provided that You
|
||||
meet the following conditions:
|
||||
|
||||
(a) You must give any other recipients of the Work or
|
||||
Derivative Works a copy of this License; and
|
||||
|
||||
(b) You must cause any modified files to carry prominent notices
|
||||
stating that You changed the files; and
|
||||
|
||||
(c) You must retain, in the Source form of any Derivative Works
|
||||
that You distribute, all copyright, patent, trademark, and
|
||||
attribution notices from the Source form of the Work,
|
||||
excluding those notices that do not pertain to any part of
|
||||
the Derivative Works; and
|
||||
|
||||
(d) If the Work includes a "NOTICE" text file as part of its
|
||||
distribution, then any Derivative Works that You distribute must
|
||||
include a readable copy of the attribution notices contained
|
||||
within such NOTICE file, excluding those notices that do not
|
||||
pertain to any part of the Derivative Works, in at least one
|
||||
of the following places: within a NOTICE text file distributed
|
||||
as part of the Derivative Works; within the Source form or
|
||||
documentation, if provided along with the Derivative Works; or,
|
||||
within a display generated by the Derivative Works, if and
|
||||
wherever such third-party notices normally appear. The contents
|
||||
of the NOTICE file are for informational purposes only and
|
||||
do not modify the License. You may add Your own attribution
|
||||
notices within Derivative Works that You distribute, alongside
|
||||
or as an addendum to the NOTICE text from the Work, provided
|
||||
that such additional attribution notices cannot be construed
|
||||
as modifying the License.
|
||||
|
||||
You may add Your own copyright statement to Your modifications and
|
||||
may provide additional or different license terms and conditions
|
||||
for use, reproduction, or distribution of Your modifications, or
|
||||
for any such Derivative Works as a whole, provided Your use,
|
||||
reproduction, and distribution of the Work otherwise complies with
|
||||
the conditions stated in this License.
|
||||
|
||||
5. Submission of Contributions. Unless You explicitly state otherwise,
|
||||
any Contribution intentionally submitted for inclusion in the Work
|
||||
by You to the Licensor shall be under the terms and conditions of
|
||||
this License, without any additional terms or conditions.
|
||||
Notwithstanding the above, nothing herein shall supersede or modify
|
||||
the terms of any separate license agreement you may have executed
|
||||
with Licensor regarding such Contributions.
|
||||
|
||||
6. Trademarks. This License does not grant permission to use the trade
|
||||
names, trademarks, service marks, or product names of the Licensor,
|
||||
except as required for reasonable and customary use in describing the
|
||||
origin of the Work and reproducing the content of the NOTICE file.
|
||||
|
||||
7. Disclaimer of Warranty. Unless required by applicable law or
|
||||
agreed to in writing, Licensor provides the Work (and each
|
||||
Contributor provides its Contributions) on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
|
||||
implied, including, without limitation, any warranties or conditions
|
||||
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
|
||||
PARTICULAR PURPOSE. You are solely responsible for determining the
|
||||
appropriateness of using or redistributing the Work and assume any
|
||||
risks associated with Your exercise of permissions under this License.
|
||||
|
||||
8. Limitation of Liability. In no event and under no legal theory,
|
||||
whether in tort (including negligence), contract, or otherwise,
|
||||
unless required by applicable law (such as deliberate and grossly
|
||||
negligent acts) or agreed to in writing, shall any Contributor be
|
||||
liable to You for damages, including any direct, indirect, special,
|
||||
incidental, or consequential damages of any character arising as a
|
||||
result of this License or out of the use or inability to use the
|
||||
Work (including but not limited to damages for loss of goodwill,
|
||||
work stoppage, computer failure or malfunction, or any and all
|
||||
other commercial damages or losses), even if such Contributor
|
||||
has been advised of the possibility of such damages.
|
||||
|
||||
9. Accepting Warranty or Additional Liability. While redistributing
|
||||
the Work or Derivative Works thereof, You may choose to offer,
|
||||
and charge a fee for, acceptance of support, warranty, indemnity,
|
||||
or other liability obligations and/or rights consistent with this
|
||||
License. However, in accepting such obligations, You may act only
|
||||
on Your own behalf and on Your sole responsibility, not on behalf
|
||||
of any other Contributor, and only if You agree to indemnify,
|
||||
defend, and hold each Contributor harmless for any liability
|
||||
incurred by, or claims asserted against, such Contributor by reason
|
||||
of your accepting any such warranty or additional liability.
|
||||
|
||||
END OF TERMS AND CONDITIONS
|
||||
|
||||
APPENDIX: How to apply the Apache License to your work.
|
||||
|
||||
To apply the Apache License to your work, attach the following
|
||||
boilerplate notice, with the fields enclosed by brackets "[]"
|
||||
replaced with your own identifying information. (Don't include
|
||||
the brackets!) The text should be enclosed in the appropriate
|
||||
comment syntax for the file format. We also recommend that a
|
||||
file or class name and description of purpose be included on the
|
||||
same "printed page" as the copyright notice for easier
|
||||
identification within third-party archives.
|
||||
|
||||
Copyright [yyyy] [name of copyright owner]
|
||||
|
||||
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.
|
||||
21
source/rsl_rl/licenses/dependencies/onnxscript-license.txt
Normal file
21
source/rsl_rl/licenses/dependencies/onnxscript-license.txt
Normal file
@@ -0,0 +1,21 @@
|
||||
MIT License
|
||||
|
||||
Copyright (c) Microsoft Corporation
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE.
|
||||
@@ -0,0 +1,19 @@
|
||||
Copyright (c) 2014 pre-commit dev team: Anthony Sottile, Ken Struys
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in
|
||||
all copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
|
||||
THE SOFTWARE.
|
||||
19
source/rsl_rl/licenses/dependencies/pre-commit-license.txt
Normal file
19
source/rsl_rl/licenses/dependencies/pre-commit-license.txt
Normal file
@@ -0,0 +1,19 @@
|
||||
Copyright (c) 2014 pre-commit dev team: Anthony Sottile, Ken Struys
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in
|
||||
all copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
|
||||
THE SOFTWARE.
|
||||
47
source/rsl_rl/licenses/dependencies/pyright-license.txt
Normal file
47
source/rsl_rl/licenses/dependencies/pyright-license.txt
Normal file
@@ -0,0 +1,47 @@
|
||||
MIT License
|
||||
|
||||
Copyright (c) 2021 Robert Craigie
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE.
|
||||
|
||||
|
||||
===============================================================================
|
||||
|
||||
MIT License
|
||||
|
||||
Pyright - A static type checker for the Python language
|
||||
Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE
|
||||
430
source/rsl_rl/licenses/dependencies/ruff-license.txt
Normal file
430
source/rsl_rl/licenses/dependencies/ruff-license.txt
Normal file
@@ -0,0 +1,430 @@
|
||||
MIT License
|
||||
|
||||
Copyright (c) 2022 Charles Marsh
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE.
|
||||
|
||||
end of terms and conditions
|
||||
|
||||
The externally maintained libraries from which parts of the Software is derived
|
||||
are:
|
||||
|
||||
- autoflake, licensed as follows:
|
||||
"""
|
||||
Copyright (C) 2012-2018 Steven Myint
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy of
|
||||
this software and associated documentation files (the "Software"), to deal in
|
||||
the Software without restriction, including without limitation the rights to
|
||||
use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies
|
||||
of the Software, and to permit persons to whom the Software is furnished to do
|
||||
so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE.
|
||||
"""
|
||||
|
||||
- autotyping, licensed as follows:
|
||||
"""
|
||||
MIT License
|
||||
|
||||
Copyright (c) 2023 Jelle Zijlstra
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE.
|
||||
"""
|
||||
|
||||
- Flake8, licensed as follows:
|
||||
"""
|
||||
== Flake8 License (MIT) ==
|
||||
|
||||
Copyright (C) 2011-2013 Tarek Ziade <tarek@ziade.org>
|
||||
Copyright (C) 2012-2016 Ian Cordasco <graffatcolmingov@gmail.com>
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy of
|
||||
this software and associated documentation files (the "Software"), to deal in
|
||||
the Software without restriction, including without limitation the rights to
|
||||
use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies
|
||||
of the Software, and to permit persons to whom the Software is furnished to do
|
||||
so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE.
|
||||
"""
|
||||
|
||||
- flake8-eradicate, licensed as follows:
|
||||
"""
|
||||
MIT License
|
||||
|
||||
Copyright (c) 2018 Nikita Sobolev
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE.
|
||||
"""
|
||||
|
||||
- flake8-pyi, licensed as follows:
|
||||
"""
|
||||
The MIT License (MIT)
|
||||
|
||||
Copyright (c) 2016 Łukasz Langa
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE.
|
||||
"""
|
||||
|
||||
- flake8-simplify, licensed as follows:
|
||||
"""
|
||||
MIT License
|
||||
|
||||
Copyright (c) 2020 Martin Thoma
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE.
|
||||
"""
|
||||
|
||||
- isort, licensed as follows:
|
||||
"""
|
||||
The MIT License (MIT)
|
||||
|
||||
Copyright (c) 2013 Timothy Edmund Crosley
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in
|
||||
all copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
|
||||
THE SOFTWARE.
|
||||
"""
|
||||
|
||||
- pygrep-hooks, licensed as follows:
|
||||
"""
|
||||
Copyright (c) 2018 Anthony Sottile
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in
|
||||
all copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
|
||||
THE SOFTWARE.
|
||||
"""
|
||||
|
||||
- pycodestyle, licensed as follows:
|
||||
"""
|
||||
Copyright © 2006-2009 Johann C. Rocholl <johann@rocholl.net>
|
||||
Copyright © 2009-2014 Florent Xicluna <florent.xicluna@gmail.com>
|
||||
Copyright © 2014-2020 Ian Lee <IanLee1521@gmail.com>
|
||||
|
||||
Licensed under the terms of the Expat License
|
||||
|
||||
Permission is hereby granted, free of charge, to any person
|
||||
obtaining a copy of this software and associated documentation files
|
||||
(the "Software"), to deal in the Software without restriction,
|
||||
including without limitation the rights to use, copy, modify, merge,
|
||||
publish, distribute, sublicense, and/or sell copies of the Software,
|
||||
and to permit persons to whom the Software is furnished to do so,
|
||||
subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be
|
||||
included in all copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
|
||||
EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
|
||||
MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
|
||||
NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS
|
||||
BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN
|
||||
ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
|
||||
CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE.
|
||||
"""
|
||||
|
||||
- pydocstyle, licensed as follows:
|
||||
"""
|
||||
Copyright (c) 2012 GreenSteam, <http://greensteam.dk/>
|
||||
|
||||
Copyright (c) 2014-2020 Amir Rachum, <http://amir.rachum.com/>
|
||||
|
||||
Copyright (c) 2020 Sambhav Kothari, <https://github.com/samj1912>
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy of
|
||||
this software and associated documentation files (the "Software"), to deal in
|
||||
the Software without restriction, including without limitation the rights to
|
||||
use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies
|
||||
of the Software, and to permit persons to whom the Software is furnished to do
|
||||
so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE.
|
||||
"""
|
||||
|
||||
- Pyflakes, licensed as follows:
|
||||
"""
|
||||
Copyright 2005-2011 Divmod, Inc.
|
||||
Copyright 2013-2014 Florent Xicluna
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining
|
||||
a copy of this software and associated documentation files (the
|
||||
"Software"), to deal in the Software without restriction, including
|
||||
without limitation the rights to use, copy, modify, merge, publish,
|
||||
distribute, sublicense, and/or sell copies of the Software, and to
|
||||
permit persons to whom the Software is furnished to do so, subject to
|
||||
the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be
|
||||
included in all copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
|
||||
EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
|
||||
MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
|
||||
NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
|
||||
LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
|
||||
OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
|
||||
WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
|
||||
"""
|
||||
|
||||
- Pyright, licensed as follows:
|
||||
"""
|
||||
MIT License
|
||||
|
||||
Pyright - A static type checker for the Python language
|
||||
Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE
|
||||
"""
|
||||
|
||||
- pyupgrade, licensed as follows:
|
||||
"""
|
||||
Copyright (c) 2017 Anthony Sottile
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in
|
||||
all copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
|
||||
THE SOFTWARE.
|
||||
"""
|
||||
|
||||
- rome/tools, licensed under the MIT license:
|
||||
"""
|
||||
MIT License
|
||||
|
||||
Copyright (c) Rome Tools, Inc. and its affiliates.
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE.
|
||||
"""
|
||||
|
||||
- RustPython, licensed as follows:
|
||||
"""
|
||||
MIT License
|
||||
|
||||
Copyright (c) 2020 RustPython Team
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE.
|
||||
"""
|
||||
|
||||
- rust-analyzer/text-size, licensed under the MIT license:
|
||||
"""
|
||||
Permission is hereby granted, free of charge, to any
|
||||
person obtaining a copy of this software and associated
|
||||
documentation files (the "Software"), to deal in the
|
||||
Software without restriction, including without
|
||||
limitation the rights to use, copy, modify, merge,
|
||||
publish, distribute, sublicense, and/or sell copies of
|
||||
the Software, and to permit persons to whom the Software
|
||||
is furnished to do so, subject to the following
|
||||
conditions:
|
||||
|
||||
The above copyright notice and this permission notice
|
||||
shall be included in all copies or substantial portions
|
||||
of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF
|
||||
ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED
|
||||
TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A
|
||||
PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT
|
||||
SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
|
||||
CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
|
||||
OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR
|
||||
IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
|
||||
DEALINGS IN THE SOFTWARE.
|
||||
"""
|
||||
21
source/rsl_rl/licenses/dependencies/tensordict-license.txt
Normal file
21
source/rsl_rl/licenses/dependencies/tensordict-license.txt
Normal file
@@ -0,0 +1,21 @@
|
||||
MIT License
|
||||
|
||||
Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE.
|
||||
73
source/rsl_rl/licenses/dependencies/torch-license.txt
Normal file
73
source/rsl_rl/licenses/dependencies/torch-license.txt
Normal file
@@ -0,0 +1,73 @@
|
||||
From PyTorch:
|
||||
|
||||
Copyright (c) 2016- Facebook, Inc (Adam Paszke)
|
||||
Copyright (c) 2014- Facebook, Inc (Soumith Chintala)
|
||||
Copyright (c) 2011-2014 Idiap Research Institute (Ronan Collobert)
|
||||
Copyright (c) 2012-2014 Deepmind Technologies (Koray Kavukcuoglu)
|
||||
Copyright (c) 2011-2012 NEC Laboratories America (Koray Kavukcuoglu)
|
||||
Copyright (c) 2011-2013 NYU (Clement Farabet)
|
||||
Copyright (c) 2006-2010 NEC Laboratories America (Ronan Collobert, Leon Bottou, Iain Melvin, Jason Weston)
|
||||
Copyright (c) 2006 Idiap Research Institute (Samy Bengio)
|
||||
Copyright (c) 2001-2004 Idiap Research Institute (Ronan Collobert, Samy Bengio, Johnny Mariethoz)
|
||||
|
||||
From Caffe2:
|
||||
|
||||
Copyright (c) 2016-present, Facebook Inc. All rights reserved.
|
||||
|
||||
All contributions by Facebook:
|
||||
Copyright (c) 2016 Facebook Inc.
|
||||
|
||||
All contributions by Google:
|
||||
Copyright (c) 2015 Google Inc.
|
||||
All rights reserved.
|
||||
|
||||
All contributions by Yangqing Jia:
|
||||
Copyright (c) 2015 Yangqing Jia
|
||||
All rights reserved.
|
||||
|
||||
All contributions by Kakao Brain:
|
||||
Copyright 2019-2020 Kakao Brain
|
||||
|
||||
All contributions from Caffe:
|
||||
Copyright(c) 2013, 2014, 2015, the respective contributors
|
||||
All rights reserved.
|
||||
|
||||
All other contributions:
|
||||
Copyright(c) 2015, 2016 the respective contributors
|
||||
All rights reserved.
|
||||
|
||||
Caffe2 uses a copyright model similar to Caffe: each contributor holds
|
||||
copyright over their contributions to Caffe2. The project versioning records
|
||||
all such contribution and copyright details. If a contributor wants to further
|
||||
mark their specific copyright on a particular contribution, they should
|
||||
indicate their copyright solely in the commit message of the change when it is
|
||||
committed.
|
||||
|
||||
All rights reserved.
|
||||
|
||||
Redistribution and use in source and binary forms, with or without
|
||||
modification, are permitted provided that the following conditions are met:
|
||||
|
||||
1. Redistributions of source code must retain the above copyright
|
||||
notice, this list of conditions and the following disclaimer.
|
||||
|
||||
2. Redistributions in binary form must reproduce the above copyright
|
||||
notice, this list of conditions and the following disclaimer in the
|
||||
documentation and/or other materials provided with the distribution.
|
||||
|
||||
3. Neither the names of Facebook, Deepmind Technologies, NYU, NEC Laboratories America
|
||||
and IDIAP Research Institute nor the names of its contributors may be
|
||||
used to endorse or promote products derived from this software without
|
||||
specific prior written permission.
|
||||
|
||||
THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
|
||||
ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
|
||||
LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
|
||||
CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
|
||||
SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
|
||||
INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
|
||||
CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
|
||||
ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
|
||||
POSSIBILITY OF SUCH DAMAGE.
|
||||
60
source/rsl_rl/pyproject.toml
Normal file
60
source/rsl_rl/pyproject.toml
Normal file
@@ -0,0 +1,60 @@
|
||||
[build-system]
|
||||
requires = ["setuptools>=61.0", "wheel"]
|
||||
build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "rsl-rl-lib"
|
||||
version = "3.3.0"
|
||||
keywords = ["reinforcement-learning", "robotics"]
|
||||
maintainers = [
|
||||
{ name="Clemens Schwarke", email="cschwarke@ethz.ch" },
|
||||
{ name="Mayank Mittal", email="mittalma@ethz.ch" },
|
||||
]
|
||||
authors = [
|
||||
{ name="Clemens Schwarke", email="cschwarke@ethz.ch" },
|
||||
{ name="Mayank Mittal", email="mittalma@ethz.ch" },
|
||||
{ name="Nikita Rudin", email="rudinn@ethz.ch" },
|
||||
{ name="David Hoeller", email="holler.david78@gmail.com" },
|
||||
]
|
||||
description = "Fast and simple RL algorithms implemented in PyTorch"
|
||||
readme = { file = "README.md", content-type = "text/markdown" }
|
||||
license = "BSD-3-Clause"
|
||||
requires-python = ">=3.9"
|
||||
classifiers = [
|
||||
"Programming Language :: Python :: 3",
|
||||
"Operating System :: OS Independent",
|
||||
]
|
||||
dependencies = [
|
||||
"torch>=2.6.0",
|
||||
"torchvision>=0.5.0",
|
||||
"tensordict>=0.7.0",
|
||||
"numpy>=1.16.4",
|
||||
"GitPython",
|
||||
"onnx",
|
||||
"onnxscript>=0.5.4",
|
||||
]
|
||||
|
||||
[project.urls]
|
||||
Homepage = "https://github.com/leggedrobotics/rsl_rl"
|
||||
Issues = "https://github.com/leggedrobotics/rsl_rl/issues"
|
||||
|
||||
[tool.setuptools.packages.find]
|
||||
where = ["."]
|
||||
include = ["rsl_rl*"]
|
||||
|
||||
[tool.setuptools.package-data]
|
||||
"rsl_rl" = ["config/*", "licenses/*"]
|
||||
|
||||
[tool.pyright]
|
||||
include = ["rsl_rl"]
|
||||
typeCheckingMode = "basic"
|
||||
pythonVersion = "3.9"
|
||||
pythonPlatform = "Linux"
|
||||
enableTypeIgnoreComments = true
|
||||
# This is required as the CI pre-commit does not download the module (i.e. numpy, torch, prettytable)
|
||||
reportMissingImports = "none"
|
||||
# This is required to ignore type checks of modules with stubs missing.
|
||||
reportMissingModuleSource = "none" # -> most common: prettytable in mdp managers
|
||||
reportGeneralTypeIssues = "none" # -> usage of literal MISSING in dataclasses
|
||||
reportOptionalMemberAccess = "none"
|
||||
reportPrivateUsage = "warning"
|
||||
6
source/rsl_rl/rsl_rl/__init__.py
Normal file
6
source/rsl_rl/rsl_rl/__init__.py
Normal file
@@ -0,0 +1,6 @@
|
||||
# Copyright (c) 2021-2026, ETH Zurich and NVIDIA CORPORATION
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
|
||||
"""Main module for the rsl_rl package."""
|
||||
12
source/rsl_rl/rsl_rl/algorithms/__init__.py
Normal file
12
source/rsl_rl/rsl_rl/algorithms/__init__.py
Normal file
@@ -0,0 +1,12 @@
|
||||
# Copyright (c) 2021-2026, ETH Zurich and NVIDIA CORPORATION
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
|
||||
"""Implementation of different learning algorithms."""
|
||||
|
||||
from .distillation import Distillation
|
||||
from .ppo import PPO
|
||||
from .moe_cts import MoECTS
|
||||
|
||||
__all__ = ["PPO", "Distillation", "MoECTS"]
|
||||
179
source/rsl_rl/rsl_rl/algorithms/distillation.py
Normal file
179
source/rsl_rl/rsl_rl/algorithms/distillation.py
Normal file
@@ -0,0 +1,179 @@
|
||||
# Copyright (c) 2021-2026, 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 tensordict import TensorDict
|
||||
|
||||
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: StudentTeacher | StudentTeacherRecurrent,
|
||||
storage: RolloutStorage,
|
||||
num_learning_epochs: int = 1,
|
||||
gradient_length: int = 15,
|
||||
learning_rate: float = 1e-3,
|
||||
max_grad_norm: float | None = None,
|
||||
loss_type: str = "mse",
|
||||
optimizer: str = "adam",
|
||||
device: str = "cpu",
|
||||
# Distributed training parameters
|
||||
multi_gpu_cfg: dict | None = 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)
|
||||
|
||||
# Create the optimizer
|
||||
self.optimizer = resolve_optimizer(optimizer)(self.policy.parameters(), lr=learning_rate)
|
||||
|
||||
# Add storage
|
||||
self.storage = storage
|
||||
self.transition = RolloutStorage.Transition()
|
||||
self.last_hidden_states = (None, 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 act(self, obs: TensorDict) -> torch.Tensor:
|
||||
# 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: TensorDict, rewards: torch.Tensor, dones: torch.Tensor, extras: dict[str, torch.Tensor]
|
||||
) -> None:
|
||||
# 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_transition(self.transition)
|
||||
self.transition.clear()
|
||||
self.policy.reset(dones)
|
||||
|
||||
def compute_returns(self, obs: TensorDict) -> None:
|
||||
# Not needed for distillation
|
||||
pass
|
||||
|
||||
def update(self) -> dict[str, float]:
|
||||
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 of 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
|
||||
|
||||
def broadcast_parameters(self) -> None:
|
||||
"""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) -> None:
|
||||
"""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
|
||||
484
source/rsl_rl/rsl_rl/algorithms/moe_cts.py
Normal file
484
source/rsl_rl/rsl_rl/algorithms/moe_cts.py
Normal file
@@ -0,0 +1,484 @@
|
||||
# Copyright (c) 2021-2026, 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 tensordict import TensorDict
|
||||
import itertools
|
||||
|
||||
from rsl_rl.modules import ActorCriticMoECTS
|
||||
from rsl_rl.modules.rnd import RandomNetworkDistillation
|
||||
from rsl_rl.storage import RolloutStorageCTS
|
||||
|
||||
|
||||
class MoECTS:
|
||||
"""Concurrent Teacher-Student algorithm (https://arxiv.org/abs/2405.10830) with MoE."""
|
||||
|
||||
policy: ActorCriticMoECTS
|
||||
"""The actor critic module."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
policy: ActorCriticMoECTS,
|
||||
storage: RolloutStorageCTS,
|
||||
num_envs: int,
|
||||
num_learning_epochs: int = 5,
|
||||
num_mini_batches: int = 4,
|
||||
clip_param: float = 0.2,
|
||||
gamma: float = 0.99,
|
||||
lam: float = 0.95,
|
||||
betas: tuple = (0.9, 0.999),
|
||||
weight_decay: float = 0.0,
|
||||
value_loss_coef: float = 1.0,
|
||||
entropy_coef: float = 0.01,
|
||||
load_balance_coef: float = 0.01,
|
||||
learning_rate: float = 0.001,
|
||||
student_encoder_learning_rate: float = 0.001,
|
||||
max_grad_norm: float = 1.0,
|
||||
use_clipped_value_loss: bool = True,
|
||||
schedule: str = "adaptive",
|
||||
desired_kl: float = 0.01,
|
||||
teacher_env_ratio: float = 0.75,
|
||||
normalize_advantage_per_mini_batch: bool = False,
|
||||
device: str = "cpu",
|
||||
# RND parameters
|
||||
rnd_cfg: dict | None = None,
|
||||
# Symmetry parameters
|
||||
symmetry_cfg: dict | None = None,
|
||||
# Distributed training parameters
|
||||
multi_gpu_cfg: dict | None = None,
|
||||
) -> None:
|
||||
assert isinstance(policy, ActorCriticMoECTS), "Policy must be an instance of ActorCriticMoECTS."
|
||||
assert not policy.is_recurrent, "Recurrent policies are not supported yet for MoECTS."
|
||||
# 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:
|
||||
# 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:
|
||||
print("[WARNING] `symmetry_cfg` detected, but MoECTS does not currently support symmetry; the configuration will be ignored.")
|
||||
self.symmetry = None
|
||||
|
||||
# PPO components
|
||||
self.policy = policy
|
||||
self.policy.to(self.device)
|
||||
|
||||
# Create the optimizer
|
||||
params1 = [
|
||||
{"params": self.policy.teacher_encoder.parameters()},
|
||||
{"params": self.policy.critic.parameters()},
|
||||
{"params": self.policy.actor.parameters()},
|
||||
{"params": getattr(self.policy, 'std', getattr(self.policy, 'log_std', []))}
|
||||
]
|
||||
self.optimizer = optim.Adam(params1, lr=learning_rate, betas=betas, weight_decay=weight_decay)
|
||||
self.optimizer_stu_enc = optim.Adam(self.policy.student_moe_encoder.parameters(), lr=student_encoder_learning_rate, betas=betas, weight_decay=weight_decay)
|
||||
|
||||
# Add storage
|
||||
self.storage = storage
|
||||
self.transition = RolloutStorageCTS.Transition()
|
||||
|
||||
# MoECTS & 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.load_balance_coef = load_balance_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
|
||||
|
||||
# Teacher-student environment split
|
||||
self.teacher_num_envs = max(int(num_envs * teacher_env_ratio), 1)
|
||||
self.student_num_envs = num_envs - self.teacher_num_envs
|
||||
student_env_ratio = 1 - teacher_env_ratio
|
||||
self.teacher_env_idxs = torch.tensor([i for i in range(num_envs) if i % int(1/student_env_ratio) != 0], device=self.device)
|
||||
self.student_env_idxs = torch.tensor([i for i in range(num_envs) if i % int(1/student_env_ratio) == 0], device=self.device)
|
||||
assert len(self.teacher_env_idxs) == self.teacher_num_envs, f"{len(self.teacher_env_idxs)=} != {self.teacher_num_envs=}"
|
||||
assert len(self.student_env_idxs) == self.student_num_envs, f"{len(self.student_env_idxs)=} != {self.student_num_envs=}"
|
||||
|
||||
def act(self, obs: TensorDict) -> torch.Tensor:
|
||||
# Compute the actions and values
|
||||
def _get_results(obs, is_teacher):
|
||||
actions = self.policy.act(obs, is_teacher)
|
||||
return (
|
||||
actions.detach(),
|
||||
self.policy.evaluate(obs, is_teacher).detach(),
|
||||
self.policy.get_actions_log_prob(actions).detach(),
|
||||
self.policy.action_mean.detach(),
|
||||
self.policy.action_std.detach(),
|
||||
)
|
||||
ti, si = self.teacher_env_idxs, self.student_env_idxs
|
||||
teacher_results = _get_results(obs[ti], is_teacher=True)
|
||||
student_results = _get_results(obs[si], is_teacher=False)
|
||||
results = []
|
||||
for x1, x2 in zip(teacher_results, student_results):
|
||||
results.append(torch.cat([x1, x2], dim=0))
|
||||
self.transition.actions = results[0]
|
||||
self.transition.values = results[1]
|
||||
self.transition.actions_log_prob = results[2]
|
||||
self.transition.action_mean = results[3]
|
||||
self.transition.action_sigma = results[4]
|
||||
|
||||
# Record observations before env.step()
|
||||
self.transition.observations = torch.cat([obs[ti], obs[si]], dim=0)
|
||||
|
||||
# Reconstruct the actions in the original order
|
||||
reordered_actions = torch.zeros_like(self.transition.actions)
|
||||
reordered_actions[ti] = self.transition.actions[:self.teacher_num_envs]
|
||||
reordered_actions[si] = self.transition.actions[self.teacher_num_envs:]
|
||||
return reordered_actions
|
||||
|
||||
def process_env_step(
|
||||
self, obs: TensorDict, rewards: torch.Tensor, dones: torch.Tensor, extras: dict[str, torch.Tensor]
|
||||
) -> 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
|
||||
ti, si = self.teacher_env_idxs, self.student_env_idxs
|
||||
rewards = rewards.clone()
|
||||
self.transition.rewards = torch.cat([rewards[ti], rewards[si]], dim=0)
|
||||
self.transition.dones = torch.cat([dones[ti], dones[si]], dim=0)
|
||||
|
||||
# Compute the intrinsic rewards and add to extrinsic rewards
|
||||
if self.rnd:
|
||||
# Compute the intrinsic rewards
|
||||
reordered_obs = torch.cat([obs[ti], obs[si]], dim=0)
|
||||
self.intrinsic_rewards = self.rnd.get_intrinsic_reward(reordered_obs)
|
||||
# Add intrinsic rewards to extrinsic rewards
|
||||
self.transition.rewards += self.intrinsic_rewards
|
||||
|
||||
# Bootstrapping on time outs
|
||||
if "time_outs" in extras:
|
||||
time_outs = extras["time_outs"].to(self.device)
|
||||
reordered_time_outs = torch.cat([time_outs[ti], time_outs[si]], dim=0)
|
||||
self.transition.rewards += self.gamma * torch.squeeze(
|
||||
self.transition.values * reordered_time_outs.unsqueeze(1).to(self.device), 1
|
||||
)
|
||||
|
||||
# Record the transition
|
||||
self.storage.add_transition(self.transition)
|
||||
self.transition.clear()
|
||||
self.policy.reset(dones)
|
||||
|
||||
def compute_returns(self, obs: TensorDict) -> None:
|
||||
st = self.storage
|
||||
# Compute value for the last step
|
||||
ti, si = self.teacher_env_idxs, self.student_env_idxs
|
||||
last_values = torch.cat([
|
||||
self.policy.evaluate(obs[ti], is_teacher=True).detach(),
|
||||
self.policy.evaluate(obs[si], is_teacher=False).detach(),
|
||||
], dim=0)
|
||||
# Compute returns and advantages
|
||||
advantage = 0
|
||||
for step in reversed(range(st.num_transitions_per_env)):
|
||||
# If we are at the last step, bootstrap the return value
|
||||
next_values = last_values if step == st.num_transitions_per_env - 1 else st.values[step + 1]
|
||||
# 1 if we are not in a terminal state, 0 otherwise
|
||||
next_is_not_terminal = 1.0 - st.dones[step].float()
|
||||
# TD error: r_t + gamma * V(s_{t+1}) - V(s_t)
|
||||
delta = st.rewards[step] + next_is_not_terminal * self.gamma * next_values - st.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 * self.gamma * self.lam * advantage
|
||||
# Return: R_t = A(s_t, a_t) + V(s_t)
|
||||
st.returns[step] = advantage + st.values[step]
|
||||
# Compute the advantages
|
||||
st.advantages = st.returns - st.values
|
||||
# Normalize the advantages if per minibatch normalization is not used
|
||||
if not self.normalize_advantage_per_mini_batch:
|
||||
st.advantages = (st.advantages - st.advantages.mean()) / (st.advantages.std() + 1e-8)
|
||||
|
||||
def update(self) -> dict[str, float]:
|
||||
mean_value_loss = 0
|
||||
mean_surrogate_loss = 0
|
||||
mean_entropy = 0
|
||||
mean_latent_loss = 0
|
||||
mean_load_balance_loss = 0
|
||||
# RND loss
|
||||
mean_rnd_loss = 0 if self.rnd else None
|
||||
|
||||
# Get mini batch generator
|
||||
generator = self.storage.mini_batch_generator(self.num_mini_batches, self.num_learning_epochs)
|
||||
data = list(generator)
|
||||
|
||||
# Iterate over batches
|
||||
teacher_samples = self.teacher_num_envs * self.storage.num_transitions_per_env // self.num_mini_batches
|
||||
student_samples = self.student_num_envs * self.storage.num_transitions_per_env // self.num_mini_batches
|
||||
for (
|
||||
obs_batch,
|
||||
actions_batch,
|
||||
target_values_batch,
|
||||
advantages_batch,
|
||||
returns_batch,
|
||||
old_actions_log_prob_batch,
|
||||
old_mu_batch,
|
||||
old_sigma_batch,
|
||||
hidden_states_batch,
|
||||
masks_batch,
|
||||
) in data:
|
||||
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)
|
||||
|
||||
def _get_results(start, end, is_teacher):
|
||||
self.policy.act(obs_batch[start:end], is_teacher)
|
||||
actions_log_prob = self.policy.get_actions_log_prob(actions_batch[start:end])
|
||||
value = self.policy.evaluate(obs_batch[start:end], is_teacher)
|
||||
mu = self.policy.action_mean
|
||||
sigma = self.policy.action_std
|
||||
entropy = self.policy.entropy
|
||||
return actions_log_prob, value, mu, sigma, entropy
|
||||
teacher_results = _get_results(0, teacher_samples, is_teacher=True)
|
||||
student_results = _get_results(teacher_samples, teacher_samples + student_samples, is_teacher=False)
|
||||
results = []
|
||||
for x1, x2 in zip(teacher_results, student_results):
|
||||
results.append(torch.cat([x1, x2], dim=0))
|
||||
actions_log_prob_batch, value_batch, mu_batch, sigma_batch, entropy_batch = results
|
||||
|
||||
# Compute KL divergence and adapt the learning rate
|
||||
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 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_losses = torch.max(surrogate, surrogate_clipped)
|
||||
teacher_surrogate_loss = surrogate_losses[:teacher_samples].mean()
|
||||
student_surrogate_loss = surrogate_losses[teacher_samples:].mean()
|
||||
surrogate_loss = teacher_surrogate_loss + student_surrogate_loss
|
||||
|
||||
# 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()
|
||||
|
||||
# RND 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()
|
||||
# Compute the gradients for RND
|
||||
if self.rnd:
|
||||
self.rnd_optimizer.zero_grad()
|
||||
rnd_loss.backward()
|
||||
|
||||
# Collect gradients from all GPUs
|
||||
if self.is_multi_gpu:
|
||||
self.reduce_parameters()
|
||||
|
||||
# Apply the gradients for PPO
|
||||
params_to_clip = itertools.chain.from_iterable(g['params'] for g in self.optimizer.param_groups)
|
||||
nn.utils.clip_grad_norm_(params_to_clip, self.max_grad_norm)
|
||||
self.optimizer.step()
|
||||
# Apply the gradients 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()
|
||||
|
||||
for (
|
||||
obs_batch,
|
||||
actions_batch,
|
||||
target_values_batch,
|
||||
advantages_batch,
|
||||
returns_batch,
|
||||
old_actions_log_prob_batch,
|
||||
old_mu_batch,
|
||||
old_sigma_batch,
|
||||
hidden_states_batch,
|
||||
masks_batch,
|
||||
) in data:
|
||||
# Student encoder loss
|
||||
obs_a_batch = self.policy.get_actor_obs(obs_batch)
|
||||
obs_a_batch = self.policy.actor_obs_normalizer(obs_a_batch)
|
||||
student_latent, gating_weights = self.policy.student_moe_encoder(obs_a_batch[teacher_samples:])
|
||||
with torch.no_grad():
|
||||
obs_c_batch = self.policy.get_critic_obs(obs_batch)
|
||||
obs_c_batch = self.policy.critic_obs_normalizer(obs_c_batch)
|
||||
teacher_latent = self.policy.teacher_encoder(obs_c_batch[teacher_samples:])
|
||||
latent_loss = (teacher_latent - student_latent).pow(2).mean()
|
||||
|
||||
# Load balance loss
|
||||
mean_usage = torch.mean(gating_weights, dim=0)
|
||||
target_usage = torch.full_like(mean_usage, 1.0 / gating_weights.shape[1])
|
||||
load_balance_loss = torch.mean((mean_usage - target_usage).pow(2))
|
||||
# load_balance_loss = torch.sum(mean_usage.pow(2)) * gating_weights.shape[1] # Switch Transformer style
|
||||
student_loss = latent_loss + self.load_balance_coef * load_balance_loss
|
||||
|
||||
self.optimizer_stu_enc.zero_grad()
|
||||
student_loss.backward()
|
||||
nn.utils.clip_grad_norm_(self.policy.student_moe_encoder.parameters(), self.max_grad_norm)
|
||||
self.optimizer_stu_enc.step()
|
||||
|
||||
mean_latent_loss += latent_loss.item()
|
||||
mean_load_balance_loss += load_balance_loss.item()
|
||||
|
||||
# Divide the losses by the number of updates
|
||||
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_latent_loss /= num_updates
|
||||
mean_load_balance_loss /= num_updates
|
||||
if mean_rnd_loss is not None:
|
||||
mean_rnd_loss /= num_updates
|
||||
|
||||
# Clear the storage
|
||||
self.storage.clear()
|
||||
|
||||
# Construct the loss dictionary
|
||||
loss_dict = {
|
||||
"value": mean_value_loss,
|
||||
"surrogate": mean_surrogate_loss,
|
||||
"entropy": mean_entropy,
|
||||
"mean_latent_loss": mean_latent_loss,
|
||||
"mean_load_balance_loss": mean_load_balance_loss
|
||||
}
|
||||
if self.rnd:
|
||||
loss_dict["rnd"] = mean_rnd_loss
|
||||
|
||||
return loss_dict
|
||||
|
||||
def broadcast_parameters(self) -> None:
|
||||
"""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) -> None:
|
||||
"""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
|
||||
458
source/rsl_rl/rsl_rl/algorithms/ppo.py
Normal file
458
source/rsl_rl/rsl_rl/algorithms/ppo.py
Normal file
@@ -0,0 +1,458 @@
|
||||
# Copyright (c) 2021-2026, 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 tensordict import TensorDict
|
||||
|
||||
from rsl_rl.modules import ActorCritic, ActorCriticCNN, ActorCriticRecurrent
|
||||
from rsl_rl.modules.rnd import RandomNetworkDistillation
|
||||
from rsl_rl.storage import RolloutStorage
|
||||
from rsl_rl.utils import resolve_callable
|
||||
|
||||
|
||||
class PPO:
|
||||
"""Proximal Policy Optimization algorithm (https://arxiv.org/abs/1707.06347)."""
|
||||
|
||||
policy: ActorCritic | ActorCriticRecurrent | ActorCriticCNN
|
||||
"""The actor critic module."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
policy: ActorCritic | ActorCriticRecurrent | ActorCriticCNN,
|
||||
storage: RolloutStorage,
|
||||
num_learning_epochs: int = 5,
|
||||
num_mini_batches: int = 4,
|
||||
clip_param: float = 0.2,
|
||||
gamma: float = 0.99,
|
||||
lam: float = 0.95,
|
||||
value_loss_coef: float = 1.0,
|
||||
entropy_coef: float = 0.01,
|
||||
learning_rate: float = 0.001,
|
||||
max_grad_norm: float = 1.0,
|
||||
use_clipped_value_loss: bool = True,
|
||||
schedule: str = "adaptive",
|
||||
desired_kl: float = 0.01,
|
||||
normalize_advantage_per_mini_batch: bool = False,
|
||||
device: str = "cpu",
|
||||
# RND parameters
|
||||
rnd_cfg: dict | None = None,
|
||||
# Symmetry parameters
|
||||
symmetry_cfg: dict | None = None,
|
||||
# Distributed training parameters
|
||||
multi_gpu_cfg: dict | None = 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:
|
||||
# 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.")
|
||||
# Resolve the data augmentation function (supports string names or direct callables)
|
||||
symmetry_cfg["data_augmentation_func"] = resolve_callable(symmetry_cfg["data_augmentation_func"])
|
||||
# Check valid configuration
|
||||
if not callable(symmetry_cfg["data_augmentation_func"]):
|
||||
raise ValueError(
|
||||
f"Symmetry configuration exists but the function is not callable: "
|
||||
f"{symmetry_cfg['data_augmentation_func']}"
|
||||
)
|
||||
# Check if the policy is compatible with symmetry
|
||||
if isinstance(policy, ActorCriticRecurrent):
|
||||
raise ValueError("Symmetry augmentation is not supported for recurrent policies.")
|
||||
# Store symmetry configuration
|
||||
self.symmetry = symmetry_cfg
|
||||
else:
|
||||
self.symmetry = None
|
||||
|
||||
# PPO components
|
||||
self.policy = policy
|
||||
self.policy.to(self.device)
|
||||
|
||||
# Create the optimizer
|
||||
self.optimizer = optim.Adam(self.policy.parameters(), lr=learning_rate)
|
||||
|
||||
# Add storage
|
||||
self.storage = storage
|
||||
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 act(self, obs: TensorDict) -> torch.Tensor:
|
||||
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()
|
||||
# Record observations before env.step()
|
||||
self.transition.observations = obs
|
||||
return self.transition.actions
|
||||
|
||||
def process_env_step(
|
||||
self, obs: TensorDict, rewards: torch.Tensor, dones: torch.Tensor, extras: dict[str, torch.Tensor]
|
||||
) -> 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
|
||||
)
|
||||
|
||||
# Record the transition
|
||||
self.storage.add_transition(self.transition)
|
||||
self.transition.clear()
|
||||
self.policy.reset(dones)
|
||||
|
||||
def compute_returns(self, obs: TensorDict) -> None:
|
||||
st = self.storage
|
||||
# Compute value for the last step
|
||||
last_values = self.policy.evaluate(obs).detach()
|
||||
# Compute returns and advantages
|
||||
advantage = 0
|
||||
for step in reversed(range(st.num_transitions_per_env)):
|
||||
# If we are at the last step, bootstrap the return value
|
||||
next_values = last_values if step == st.num_transitions_per_env - 1 else st.values[step + 1]
|
||||
# 1 if we are not in a terminal state, 0 otherwise
|
||||
next_is_not_terminal = 1.0 - st.dones[step].float()
|
||||
# TD error: r_t + gamma * V(s_{t+1}) - V(s_t)
|
||||
delta = st.rewards[step] + next_is_not_terminal * self.gamma * next_values - st.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 * self.gamma * self.lam * advantage
|
||||
# Return: R_t = A(s_t, a_t) + V(s_t)
|
||||
st.returns[step] = advantage + st.values[step]
|
||||
# Compute the advantages
|
||||
st.advantages = st.returns - st.values
|
||||
# Normalize the advantages if per minibatch normalization is not used
|
||||
if not self.normalize_advantage_per_mini_batch:
|
||||
st.advantages = (st.advantages - st.advantages.mean()) / (st.advantages.std() + 1e-8)
|
||||
|
||||
def update(self) -> dict[str, float]:
|
||||
mean_value_loss = 0
|
||||
mean_surrogate_loss = 0
|
||||
mean_entropy = 0
|
||||
# RND loss
|
||||
mean_rnd_loss = 0 if self.rnd else None
|
||||
# Symmetry loss
|
||||
mean_symmetry_loss = 0 if self.symmetry else None
|
||||
|
||||
# Get mini batch generator
|
||||
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,
|
||||
hidden_states_batch,
|
||||
masks_batch,
|
||||
) in generator:
|
||||
num_aug = 1 # Number of augmentations per sample. Starts at 1 for no 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
|
||||
num_aug = int(obs_batch.batch_size[0] / original_batch_size)
|
||||
# Repeat the rest of the batch
|
||||
old_actions_log_prob_batch = old_actions_log_prob_batch.repeat(num_aug, 1)
|
||||
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
|
||||
self.policy.act(obs_batch, masks=masks_batch, hidden_state=hidden_states_batch[0])
|
||||
actions_log_prob_batch = self.policy.get_actions_log_prob(actions_batch)
|
||||
value_batch = self.policy.evaluate(obs_batch, masks=masks_batch, hidden_state=hidden_states_batch[1])
|
||||
# Note: 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]
|
||||
|
||||
# Compute KL divergence and adapt the learning rate
|
||||
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 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
|
||||
# Note: 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
|
||||
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()
|
||||
|
||||
# RND 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()
|
||||
# Compute the gradients for RND
|
||||
if self.rnd:
|
||||
self.rnd_optimizer.zero_grad()
|
||||
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()
|
||||
# Apply the gradients 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()
|
||||
|
||||
# Divide the losses by the number of updates
|
||||
num_updates = self.num_learning_epochs * self.num_mini_batches
|
||||
mean_value_loss /= num_updates
|
||||
mean_surrogate_loss /= num_updates
|
||||
mean_entropy /= num_updates
|
||||
if mean_rnd_loss is not None:
|
||||
mean_rnd_loss /= num_updates
|
||||
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": 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
|
||||
|
||||
def broadcast_parameters(self) -> None:
|
||||
"""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) -> None:
|
||||
"""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
|
||||
10
source/rsl_rl/rsl_rl/env/__init__.py
vendored
Normal file
10
source/rsl_rl/rsl_rl/env/__init__.py
vendored
Normal file
@@ -0,0 +1,10 @@
|
||||
# Copyright (c) 2021-2026, 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"]
|
||||
91
source/rsl_rl/rsl_rl/env/vec_env.py
vendored
Normal file
91
source/rsl_rl/rsl_rl/env/vec_env.py
vendored
Normal file
@@ -0,0 +1,91 @@
|
||||
# Copyright (c) 2021-2026, 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
|
||||
"""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."""
|
||||
|
||||
@abstractmethod
|
||||
def get_observations(self) -> TensorDict:
|
||||
"""Return the current observations.
|
||||
|
||||
Returns:
|
||||
The 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: Input actions to apply. Shape: (num_envs, num_actions)
|
||||
|
||||
Returns:
|
||||
observations: Observations from the environment.
|
||||
rewards: Rewards from the environment. Shape: (num_envs,)
|
||||
dones: Done flags from the environment. Shape: (num_envs,)
|
||||
extras: 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
|
||||
27
source/rsl_rl/rsl_rl/modules/__init__.py
Normal file
27
source/rsl_rl/rsl_rl/modules/__init__.py
Normal file
@@ -0,0 +1,27 @@
|
||||
# Copyright (c) 2021-2026, 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_cnn import ActorCriticCNN
|
||||
from .actor_critic_recurrent import ActorCriticRecurrent
|
||||
from .actor_critic_moe_cts import ActorCriticMoECTS
|
||||
from .rnd import RandomNetworkDistillation, resolve_rnd_config
|
||||
from .student_teacher import StudentTeacher
|
||||
from .student_teacher_recurrent import StudentTeacherRecurrent
|
||||
from .symmetry import resolve_symmetry_config
|
||||
|
||||
__all__ = [
|
||||
"ActorCritic",
|
||||
"ActorCriticCNN",
|
||||
"ActorCriticRecurrent",
|
||||
"ActorCriticMoECTS",
|
||||
"RandomNetworkDistillation",
|
||||
"StudentTeacher",
|
||||
"StudentTeacherRecurrent",
|
||||
"resolve_rnd_config",
|
||||
"resolve_symmetry_config",
|
||||
]
|
||||
198
source/rsl_rl/rsl_rl/modules/actor_critic.py
Normal file
198
source/rsl_rl/rsl_rl/modules/actor_critic.py
Normal file
@@ -0,0 +1,198 @@
|
||||
# Copyright (c) 2021-2026, 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 tensordict import TensorDict
|
||||
from torch.distributions import Normal
|
||||
from typing import Any, NoReturn
|
||||
|
||||
from rsl_rl.networks import MLP, EmpiricalNormalization
|
||||
|
||||
|
||||
class ActorCritic(nn.Module):
|
||||
is_recurrent: bool = False
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
obs: TensorDict,
|
||||
obs_groups: dict[str, list[str]],
|
||||
num_actions: int,
|
||||
actor_obs_normalization: bool = False,
|
||||
critic_obs_normalization: bool = False,
|
||||
actor_hidden_dims: tuple[int] | list[int] = [256, 256, 256],
|
||||
critic_hidden_dims: tuple[int] | list[int] = [256, 256, 256],
|
||||
activation: str = "elu",
|
||||
init_noise_std: float = 1.0,
|
||||
noise_std_type: str = "scalar",
|
||||
state_dependent_std: bool = False,
|
||||
**kwargs: dict[str, Any],
|
||||
) -> None:
|
||||
if kwargs:
|
||||
print(
|
||||
"ActorCritic.__init__ got unexpected arguments, which will be ignored: " + str([key for key in kwargs])
|
||||
)
|
||||
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]
|
||||
|
||||
# Actor
|
||||
self.state_dependent_std = state_dependent_std
|
||||
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)
|
||||
print(f"Actor MLP: {self.actor}")
|
||||
|
||||
# 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()
|
||||
|
||||
# Critic
|
||||
self.critic = MLP(num_critic_obs, 1, critic_hidden_dims, activation)
|
||||
print(f"Critic MLP: {self.critic}")
|
||||
|
||||
# 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()
|
||||
|
||||
# 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
|
||||
# Note: Populated in update_distribution
|
||||
self.distribution = None
|
||||
|
||||
# Disable args validation for speedup
|
||||
Normal.set_default_validate_args(False)
|
||||
|
||||
def reset(self, dones: torch.Tensor | None = None) -> None:
|
||||
pass
|
||||
|
||||
def forward(self) -> NoReturn:
|
||||
raise NotImplementedError
|
||||
|
||||
@property
|
||||
def action_mean(self) -> torch.Tensor:
|
||||
return self.distribution.mean
|
||||
|
||||
@property
|
||||
def action_std(self) -> torch.Tensor:
|
||||
return self.distribution.stddev
|
||||
|
||||
@property
|
||||
def entropy(self) -> torch.Tensor:
|
||||
return self.distribution.entropy().sum(dim=-1)
|
||||
|
||||
def _update_distribution(self, obs: torch.Tensor) -> None:
|
||||
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: TensorDict, **kwargs: dict[str, Any]) -> torch.Tensor:
|
||||
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: TensorDict) -> torch.Tensor:
|
||||
obs = self.get_actor_obs(obs)
|
||||
obs = self.actor_obs_normalizer(obs)
|
||||
if self.state_dependent_std:
|
||||
return self.actor(obs)[..., 0, :]
|
||||
else:
|
||||
return self.actor(obs)
|
||||
|
||||
def evaluate(self, obs: TensorDict, **kwargs: dict[str, Any]) -> torch.Tensor:
|
||||
obs = self.get_critic_obs(obs)
|
||||
obs = self.critic_obs_normalizer(obs)
|
||||
return self.critic(obs)
|
||||
|
||||
def get_actor_obs(self, obs: TensorDict) -> torch.Tensor:
|
||||
obs_list = [obs[obs_group] for obs_group in self.obs_groups["policy"]]
|
||||
return torch.cat(obs_list, dim=-1)
|
||||
|
||||
def get_critic_obs(self, obs: TensorDict) -> torch.Tensor:
|
||||
obs_list = [obs[obs_group] for obs_group in self.obs_groups["critic"]]
|
||||
return torch.cat(obs_list, dim=-1)
|
||||
|
||||
def get_actions_log_prob(self, actions: torch.Tensor) -> torch.Tensor:
|
||||
return self.distribution.log_prob(actions).sum(dim=-1)
|
||||
|
||||
def update_normalization(self, obs: TensorDict) -> None:
|
||||
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: dict, strict: bool = True) -> bool:
|
||||
"""Load the parameters of the actor-critic model.
|
||||
|
||||
Args:
|
||||
state_dict: State dictionary of the model.
|
||||
strict: Whether to strictly enforce that the keys in `state_dict` match the keys returned by this module's
|
||||
:meth:`state_dict` function.
|
||||
|
||||
Returns:
|
||||
Whether this training resumes a previous training. This flag is used by the :func:`load` function of
|
||||
:class:`OnPolicyRunner` to determine how to load further parameters (relevant for, e.g., distillation).
|
||||
"""
|
||||
super().load_state_dict(state_dict, strict=strict)
|
||||
return True
|
||||
262
source/rsl_rl/rsl_rl/modules/actor_critic_cnn.py
Normal file
262
source/rsl_rl/rsl_rl/modules/actor_critic_cnn.py
Normal file
@@ -0,0 +1,262 @@
|
||||
# Copyright (c) 2021-2026, 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 tensordict import TensorDict
|
||||
from torch.distributions import Normal
|
||||
from typing import Any
|
||||
|
||||
from rsl_rl.networks import CNN, MLP, EmpiricalNormalization
|
||||
|
||||
from .actor_critic import ActorCritic
|
||||
|
||||
|
||||
class ActorCriticCNN(ActorCritic):
|
||||
def __init__(
|
||||
self,
|
||||
obs: TensorDict,
|
||||
obs_groups: dict[str, list[str]],
|
||||
num_actions: int,
|
||||
actor_obs_normalization: bool = False,
|
||||
critic_obs_normalization: bool = False,
|
||||
actor_hidden_dims: tuple[int] | list[int] = [256, 256, 256],
|
||||
critic_hidden_dims: tuple[int] | list[int] = [256, 256, 256],
|
||||
actor_cnn_cfg: dict[str, dict] | dict | None = None,
|
||||
critic_cnn_cfg: dict[str, dict] | dict | None = None,
|
||||
activation: str = "elu",
|
||||
init_noise_std: float = 1.0,
|
||||
noise_std_type: str = "scalar",
|
||||
state_dependent_std: bool = False,
|
||||
**kwargs: dict[str, Any],
|
||||
) -> None:
|
||||
if kwargs:
|
||||
print(
|
||||
"ActorCriticCNN.__init__ got unexpected arguments, which will be ignored: "
|
||||
+ str([key for key in kwargs])
|
||||
)
|
||||
super(ActorCritic, self).__init__()
|
||||
|
||||
# Get the observation dimensions
|
||||
self.obs_groups = obs_groups
|
||||
num_actor_obs_1d = 0
|
||||
self.actor_obs_groups_1d = []
|
||||
actor_in_dims_2d = []
|
||||
actor_in_channels_2d = []
|
||||
self.actor_obs_groups_2d = []
|
||||
for obs_group in obs_groups["policy"]:
|
||||
if len(obs[obs_group].shape) == 4: # B, C, H, W
|
||||
self.actor_obs_groups_2d.append(obs_group)
|
||||
actor_in_dims_2d.append(obs[obs_group].shape[2:4])
|
||||
actor_in_channels_2d.append(obs[obs_group].shape[1])
|
||||
elif len(obs[obs_group].shape) == 2: # B, C
|
||||
self.actor_obs_groups_1d.append(obs_group)
|
||||
num_actor_obs_1d += obs[obs_group].shape[-1]
|
||||
else:
|
||||
raise ValueError(f"Invalid observation shape for {obs_group}: {obs[obs_group].shape}")
|
||||
num_critic_obs_1d = 0
|
||||
self.critic_obs_groups_1d = []
|
||||
critic_in_dims_2d = []
|
||||
critic_in_channels_2d = []
|
||||
self.critic_obs_groups_2d = []
|
||||
for obs_group in obs_groups["critic"]:
|
||||
if len(obs[obs_group].shape) == 4: # B, C, H, W
|
||||
self.critic_obs_groups_2d.append(obs_group)
|
||||
critic_in_dims_2d.append(obs[obs_group].shape[2:4])
|
||||
critic_in_channels_2d.append(obs[obs_group].shape[1])
|
||||
elif len(obs[obs_group].shape) == 2: # B, C
|
||||
self.critic_obs_groups_1d.append(obs_group)
|
||||
num_critic_obs_1d += obs[obs_group].shape[-1]
|
||||
else:
|
||||
raise ValueError(f"Invalid observation shape for {obs_group}: {obs[obs_group].shape}")
|
||||
|
||||
# Assert that there are 2D observations
|
||||
assert self.actor_obs_groups_2d or self.critic_obs_groups_2d, (
|
||||
"No 2D observations are provided. If this is intentional, use the ActorCritic module instead."
|
||||
)
|
||||
|
||||
# Actor CNN
|
||||
if self.actor_obs_groups_2d:
|
||||
# Resolve the actor CNN configuration
|
||||
assert actor_cnn_cfg is not None, "An actor CNN configuration is required for 2D actor observations."
|
||||
# If a single configuration dictionary is provided, create a dictionary for each 2D observation group
|
||||
if not all(isinstance(v, dict) for v in actor_cnn_cfg.values()):
|
||||
actor_cnn_cfg = {group: actor_cnn_cfg for group in self.actor_obs_groups_2d}
|
||||
# Check that the number of configs matches the number of observation groups
|
||||
assert len(actor_cnn_cfg) == len(self.actor_obs_groups_2d), (
|
||||
"The number of CNN configurations must match the number of 2D actor observations."
|
||||
)
|
||||
|
||||
# Create CNNs for each 2D actor observation
|
||||
self.actor_cnns = nn.ModuleDict()
|
||||
encoding_dim = 0
|
||||
for idx, obs_group in enumerate(self.actor_obs_groups_2d):
|
||||
self.actor_cnns[obs_group] = CNN(
|
||||
input_dim=actor_in_dims_2d[idx],
|
||||
input_channels=actor_in_channels_2d[idx],
|
||||
**actor_cnn_cfg[obs_group],
|
||||
)
|
||||
print(f"Actor CNN for {obs_group}: {self.actor_cnns[obs_group]}")
|
||||
# Get the output dimension of the CNN
|
||||
if self.actor_cnns[obs_group].output_channels is None:
|
||||
encoding_dim += int(self.actor_cnns[obs_group].output_dim)
|
||||
else:
|
||||
raise ValueError("The output of the actor CNN must be flattened before passing it to the MLP.")
|
||||
else:
|
||||
self.actor_cnns = None
|
||||
encoding_dim = 0
|
||||
|
||||
# Actor MLP
|
||||
self.state_dependent_std = state_dependent_std
|
||||
if self.state_dependent_std:
|
||||
self.actor = MLP(num_actor_obs_1d + encoding_dim, [2, num_actions], actor_hidden_dims, activation)
|
||||
else:
|
||||
self.actor = MLP(num_actor_obs_1d + encoding_dim, num_actions, actor_hidden_dims, activation)
|
||||
print(f"Actor MLP: {self.actor}")
|
||||
|
||||
# Actor observation normalization (only for 1D actor observations)
|
||||
self.actor_obs_normalization = actor_obs_normalization
|
||||
if actor_obs_normalization:
|
||||
self.actor_obs_normalizer = EmpiricalNormalization(num_actor_obs_1d)
|
||||
else:
|
||||
self.actor_obs_normalizer = torch.nn.Identity()
|
||||
|
||||
# Critic CNN
|
||||
if self.critic_obs_groups_2d:
|
||||
# Resolve the critic CNN configuration
|
||||
assert critic_cnn_cfg is not None, "A critic CNN configuration is required for 2D critic observations."
|
||||
# If a single configuration dictionary is provided, create a dictionary for each 2D observation group
|
||||
if not all(isinstance(v, dict) for v in critic_cnn_cfg.values()):
|
||||
critic_cnn_cfg = {group: critic_cnn_cfg for group in self.critic_obs_groups_2d}
|
||||
# Check that the number of configs matches the number of observation groups
|
||||
assert len(critic_cnn_cfg) == len(self.critic_obs_groups_2d), (
|
||||
"The number of CNN configurations must match the number of 2D critic observations."
|
||||
)
|
||||
|
||||
# Create CNNs for each 2D critic observation
|
||||
self.critic_cnns = nn.ModuleDict()
|
||||
encoding_dim = 0
|
||||
for idx, obs_group in enumerate(self.critic_obs_groups_2d):
|
||||
self.critic_cnns[obs_group] = CNN(
|
||||
input_dim=critic_in_dims_2d[idx],
|
||||
input_channels=critic_in_channels_2d[idx],
|
||||
**critic_cnn_cfg[obs_group],
|
||||
)
|
||||
print(f"Critic CNN for {obs_group}: {self.critic_cnns[obs_group]}")
|
||||
# Get the output dimension of the CNN
|
||||
if self.critic_cnns[obs_group].output_channels is None:
|
||||
encoding_dim += int(self.critic_cnns[obs_group].output_dim)
|
||||
else:
|
||||
raise ValueError("The output of the critic CNN must be flattened before passing it to the MLP.")
|
||||
else:
|
||||
self.critic_cnns = None
|
||||
encoding_dim = 0
|
||||
|
||||
# Critic MLP
|
||||
self.critic = MLP(num_critic_obs_1d + encoding_dim, 1, critic_hidden_dims, activation)
|
||||
print(f"Critic MLP: {self.critic}")
|
||||
|
||||
# Critic observation normalization (only for 1D critic observations)
|
||||
self.critic_obs_normalization = critic_obs_normalization
|
||||
if critic_obs_normalization:
|
||||
self.critic_obs_normalizer = EmpiricalNormalization(num_critic_obs_1d)
|
||||
else:
|
||||
self.critic_obs_normalizer = torch.nn.Identity()
|
||||
|
||||
# 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
|
||||
# Note: Populated in update_distribution
|
||||
self.distribution = None
|
||||
|
||||
# Disable args validation for speedup
|
||||
Normal.set_default_validate_args(False)
|
||||
|
||||
def _update_distribution(self, mlp_obs: torch.Tensor, cnn_obs: dict[str, torch.Tensor]) -> None:
|
||||
if self.actor_cnns is not None:
|
||||
# Encode the 2D actor observations
|
||||
cnn_enc_list = [self.actor_cnns[obs_group](cnn_obs[obs_group]) for obs_group in self.actor_obs_groups_2d]
|
||||
cnn_enc = torch.cat(cnn_enc_list, dim=-1)
|
||||
# Concatenate to the MLP observations
|
||||
mlp_obs = torch.cat([mlp_obs, cnn_enc], dim=-1)
|
||||
|
||||
super()._update_distribution(mlp_obs)
|
||||
|
||||
def act(self, obs: TensorDict, **kwargs: dict[str, Any]) -> torch.Tensor:
|
||||
mlp_obs, cnn_obs = self.get_actor_obs(obs)
|
||||
mlp_obs = self.actor_obs_normalizer(mlp_obs)
|
||||
self._update_distribution(mlp_obs, cnn_obs)
|
||||
return self.distribution.sample() # type: ignore
|
||||
|
||||
def act_inference(self, obs: TensorDict) -> torch.Tensor:
|
||||
mlp_obs, cnn_obs = self.get_actor_obs(obs)
|
||||
mlp_obs = self.actor_obs_normalizer(mlp_obs)
|
||||
|
||||
if self.actor_cnns is not None:
|
||||
# Encode the 2D actor observations
|
||||
cnn_enc_list = [self.actor_cnns[obs_group](cnn_obs[obs_group]) for obs_group in self.actor_obs_groups_2d]
|
||||
cnn_enc = torch.cat(cnn_enc_list, dim=-1)
|
||||
# Concatenate to the MLP observations
|
||||
mlp_obs = torch.cat([mlp_obs, cnn_enc], dim=-1)
|
||||
|
||||
if self.state_dependent_std:
|
||||
return self.actor(mlp_obs)[..., 0, :]
|
||||
else:
|
||||
return self.actor(mlp_obs)
|
||||
|
||||
def evaluate(self, obs: TensorDict, **kwargs: dict[str, Any]) -> torch.Tensor:
|
||||
mlp_obs, cnn_obs = self.get_critic_obs(obs)
|
||||
mlp_obs = self.critic_obs_normalizer(mlp_obs)
|
||||
|
||||
if self.critic_cnns is not None:
|
||||
# Encode the 2D critic observations
|
||||
cnn_enc_list = [self.critic_cnns[obs_group](cnn_obs[obs_group]) for obs_group in self.critic_obs_groups_2d]
|
||||
cnn_enc = torch.cat(cnn_enc_list, dim=-1)
|
||||
# Concatenate to the MLP observations
|
||||
mlp_obs = torch.cat([mlp_obs, cnn_enc], dim=-1)
|
||||
|
||||
return self.critic(mlp_obs)
|
||||
|
||||
def get_actor_obs(self, obs: TensorDict) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
|
||||
obs_list_1d = [obs[obs_group] for obs_group in self.actor_obs_groups_1d]
|
||||
obs_dict_2d = {}
|
||||
for obs_group in self.actor_obs_groups_2d:
|
||||
obs_dict_2d[obs_group] = obs[obs_group]
|
||||
return torch.cat(obs_list_1d, dim=-1), obs_dict_2d
|
||||
|
||||
def get_critic_obs(self, obs: TensorDict) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
|
||||
obs_list_1d = [obs[obs_group] for obs_group in self.critic_obs_groups_1d]
|
||||
obs_dict_2d = {}
|
||||
for obs_group in self.critic_obs_groups_2d:
|
||||
obs_dict_2d[obs_group] = obs[obs_group]
|
||||
return torch.cat(obs_list_1d, dim=-1), obs_dict_2d
|
||||
|
||||
def update_normalization(self, obs: TensorDict) -> None:
|
||||
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)
|
||||
281
source/rsl_rl/rsl_rl/modules/actor_critic_moe_cts.py
Normal file
281
source/rsl_rl/rsl_rl/modules/actor_critic_moe_cts.py
Normal file
@@ -0,0 +1,281 @@
|
||||
# Copyright (c) 2021-2026, 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 tensordict import TensorDict
|
||||
from torch.distributions import Normal
|
||||
from typing import Any, NoReturn
|
||||
from rsl_rl.networks.moe import MLP
|
||||
|
||||
from rsl_rl.networks import EmpiricalNormalization, L2Norm, SimNorm, MoE
|
||||
|
||||
class StudentMoEEncoder(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
expert_num,
|
||||
input_dim,
|
||||
hidden_dims,
|
||||
output_dim,
|
||||
activation='elu',
|
||||
norm_type='l2norm',
|
||||
):
|
||||
super().__init__()
|
||||
self.norm_layer = L2Norm() if norm_type == 'l2norm' else SimNorm()
|
||||
self.moe = MoE(
|
||||
expert_num=expert_num,
|
||||
input_dim=input_dim,
|
||||
hidden_dims=hidden_dims,
|
||||
output_dim=output_dim,
|
||||
activation=activation,
|
||||
)
|
||||
|
||||
def forward(self, obs):
|
||||
latent, weights = self.moe(obs)
|
||||
latent = self.norm_layer(latent)
|
||||
return latent, weights
|
||||
|
||||
class ActorCriticMoECTS(nn.Module):
|
||||
is_recurrent: bool = False
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
obs: TensorDict,
|
||||
obs_groups: dict[str, list[str]],
|
||||
num_actions: int,
|
||||
actor_obs_normalization: bool = False,
|
||||
critic_obs_normalization: bool = False,
|
||||
actor_hidden_dims: tuple[int] | list[int] = [256, 256, 256],
|
||||
critic_hidden_dims: tuple[int] | list[int] = [256, 256, 256],
|
||||
teacher_encoder_hidden_dims: tuple[int] | list[int] = [512, 256],
|
||||
student_encoder_hidden_dims: tuple[int] | list[int] = [512, 256, 128],
|
||||
expert_num: int = 8,
|
||||
activation: str = "elu",
|
||||
init_noise_std: float = 1.0,
|
||||
noise_std_type: str = "scalar",
|
||||
state_dependent_std: bool = False,
|
||||
latent_dim: int = 32,
|
||||
norm_type: str = 'l2norm',
|
||||
**kwargs: dict[str, Any],
|
||||
) -> None:
|
||||
if kwargs:
|
||||
print(
|
||||
"ActorCriticMoECTS.__init__ got unexpected arguments, which will be ignored: " + str([key for key in kwargs])
|
||||
)
|
||||
assert norm_type in ['l2norm', 'simnorm'], f"Normalization type {norm_type} not supported!"
|
||||
assert "policy" in obs.keys() and "critic" in obs.keys() and "single_obs" in obs.keys(), \
|
||||
"obs must contain 'policy', 'critic' and 'single_obs' keys for ActorCriticMoECTS."
|
||||
super().__init__()
|
||||
|
||||
self.num_actions = num_actions
|
||||
|
||||
# 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 ActorCriticMoECTS 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 ActorCriticMoECTS module only supports 1D observations."
|
||||
num_critic_obs += obs[obs_group].shape[-1]
|
||||
|
||||
# MLP input dimensions (teacher, student, actor, critic)
|
||||
self.num_actor_obs = num_actor_obs
|
||||
self.num_single_obs = obs['single_obs'].shape[-1]
|
||||
mlp_input_dim_t = num_critic_obs
|
||||
mlp_input_dim_s = num_actor_obs
|
||||
mlp_input_dim_a = latent_dim + self.num_single_obs
|
||||
mlp_input_dim_c = latent_dim + num_critic_obs
|
||||
|
||||
# Teacher encoder
|
||||
self.teacher_encoder = nn.Sequential(
|
||||
MLP(mlp_input_dim_t, latent_dim, teacher_encoder_hidden_dims, activation=activation),
|
||||
L2Norm() if norm_type == 'l2norm' else SimNorm()
|
||||
)
|
||||
print(f"Teacher Encoder: {self.teacher_encoder}")
|
||||
|
||||
# Student MoE encoder
|
||||
self.student_moe_encoder = StudentMoEEncoder(
|
||||
expert_num=expert_num,
|
||||
input_dim=mlp_input_dim_s,
|
||||
hidden_dims=student_encoder_hidden_dims,
|
||||
output_dim=latent_dim,
|
||||
activation=activation,
|
||||
norm_type=norm_type,
|
||||
)
|
||||
print(f"Student MoE Encoder: {self.student_moe_encoder}")
|
||||
|
||||
# Actor
|
||||
self.state_dependent_std = state_dependent_std
|
||||
if self.state_dependent_std:
|
||||
self.actor = MLP(mlp_input_dim_a, [2, num_actions], actor_hidden_dims, activation)
|
||||
else:
|
||||
self.actor = MLP(mlp_input_dim_a, num_actions, actor_hidden_dims, activation)
|
||||
print(f"Actor MLP: {self.actor}")
|
||||
|
||||
# Actor observation normalization
|
||||
self.actor_obs_normalization = actor_obs_normalization
|
||||
if actor_obs_normalization:
|
||||
self.actor_obs_normalizer = EmpiricalNormalization(self.num_actor_obs)
|
||||
self.single_obs_normalizer = EmpiricalNormalization(self.num_single_obs)
|
||||
else:
|
||||
self.actor_obs_normalizer = torch.nn.Identity()
|
||||
self.single_obs_normalizer = torch.nn.Identity()
|
||||
|
||||
# Critic
|
||||
self.critic = MLP(mlp_input_dim_c, 1, critic_hidden_dims, activation)
|
||||
print(f"Critic MLP: {self.critic}")
|
||||
|
||||
# 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()
|
||||
|
||||
# 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
|
||||
# Note: Populated in update_distribution
|
||||
self.distribution = None
|
||||
|
||||
# Disable args validation for speedup
|
||||
Normal.set_default_validate_args(False)
|
||||
|
||||
def reset(self, dones: torch.Tensor | None = None) -> None:
|
||||
pass
|
||||
|
||||
def forward(self) -> NoReturn:
|
||||
raise NotImplementedError
|
||||
|
||||
@property
|
||||
def action_mean(self) -> torch.Tensor:
|
||||
return self.distribution.mean
|
||||
|
||||
@property
|
||||
def action_std(self) -> torch.Tensor:
|
||||
return self.distribution.stddev
|
||||
|
||||
@property
|
||||
def entropy(self) -> torch.Tensor:
|
||||
return self.distribution.entropy().sum(dim=-1)
|
||||
|
||||
def _update_distribution(self, latent_and_obs: torch.Tensor) -> None:
|
||||
if self.state_dependent_std:
|
||||
# Compute mean and standard deviation
|
||||
mean_and_std = self.actor(latent_and_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(latent_and_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: TensorDict, is_teacher: bool, **kwargs: dict[str, Any]) -> torch.Tensor:
|
||||
single_obs = self.single_obs_normalizer(obs['single_obs'])
|
||||
if is_teacher:
|
||||
obs_c = self.get_critic_obs(obs)
|
||||
obs_c = self.critic_obs_normalizer(obs_c)
|
||||
latent = self.teacher_encoder(obs_c)
|
||||
else:
|
||||
with torch.no_grad():
|
||||
obs_a = self.get_actor_obs(obs)
|
||||
obs_a = self.actor_obs_normalizer(obs_a)
|
||||
latent, _ = self.student_moe_encoder(obs_a)
|
||||
latent_and_obs = torch.cat([latent, single_obs], dim=-1)
|
||||
self._update_distribution(latent_and_obs)
|
||||
return self.distribution.sample()
|
||||
|
||||
def act_inference(self, obs: TensorDict) -> torch.Tensor:
|
||||
single_obs = self.single_obs_normalizer(obs['single_obs'])
|
||||
obs_a = self.get_actor_obs(obs)
|
||||
obs_a = self.actor_obs_normalizer(obs_a)
|
||||
latent, _ = self.student_moe_encoder(obs_a)
|
||||
latent_and_obs = torch.cat([latent, single_obs], dim=-1)
|
||||
if self.state_dependent_std:
|
||||
return self.actor(latent_and_obs)[..., 0, :]
|
||||
else:
|
||||
return self.actor(latent_and_obs)
|
||||
|
||||
def evaluate(self, obs: TensorDict, is_teacher: bool, **kwargs: dict[str, Any]) -> torch.Tensor:
|
||||
obs_c = self.get_critic_obs(obs)
|
||||
obs_c = self.critic_obs_normalizer(obs_c)
|
||||
if is_teacher:
|
||||
latent = self.teacher_encoder(obs_c)
|
||||
else:
|
||||
obs_a = self.get_actor_obs(obs)
|
||||
obs_a = self.actor_obs_normalizer(obs_a)
|
||||
latent, _ = self.student_moe_encoder(obs_a)
|
||||
latent_and_obs = torch.cat([latent.detach(), obs_c], dim=-1)
|
||||
return self.critic(latent_and_obs)
|
||||
|
||||
def get_actor_obs(self, obs: TensorDict) -> torch.Tensor:
|
||||
obs_list = [obs[obs_group] for obs_group in self.obs_groups["policy"]]
|
||||
return torch.cat(obs_list, dim=-1)
|
||||
|
||||
def get_critic_obs(self, obs: TensorDict) -> torch.Tensor:
|
||||
obs_list = [obs[obs_group] for obs_group in self.obs_groups["critic"]]
|
||||
return torch.cat(obs_list, dim=-1)
|
||||
|
||||
def get_actions_log_prob(self, actions: torch.Tensor) -> torch.Tensor:
|
||||
return self.distribution.log_prob(actions).sum(dim=-1)
|
||||
|
||||
def update_normalization(self, obs: TensorDict) -> None:
|
||||
if self.actor_obs_normalization:
|
||||
actor_obs = self.get_actor_obs(obs)
|
||||
self.actor_obs_normalizer.update(actor_obs)
|
||||
self.single_obs_normalizer.update(obs['single_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: dict, strict: bool = True) -> bool:
|
||||
"""Load the parameters of the actor-critic model.
|
||||
|
||||
Args:
|
||||
state_dict: State dictionary of the model.
|
||||
strict: Whether to strictly enforce that the keys in `state_dict` match the keys returned by this module's
|
||||
:meth:`state_dict` function.
|
||||
|
||||
Returns:
|
||||
Whether this training resumes a previous training. This flag is used by the :func:`load` function of
|
||||
:class:`OnPolicyRunner` to determine how to load further parameters (relevant for, e.g., distillation).
|
||||
"""
|
||||
super().load_state_dict(state_dict, strict=strict)
|
||||
return True
|
||||
223
source/rsl_rl/rsl_rl/modules/actor_critic_recurrent.py
Normal file
223
source/rsl_rl/rsl_rl/modules/actor_critic_recurrent.py
Normal file
@@ -0,0 +1,223 @@
|
||||
# Copyright (c) 2021-2026, 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 tensordict import TensorDict
|
||||
from torch.distributions import Normal
|
||||
from typing import Any, NoReturn
|
||||
|
||||
from rsl_rl.networks import MLP, EmpiricalNormalization, HiddenState, Memory
|
||||
|
||||
|
||||
class ActorCriticRecurrent(nn.Module):
|
||||
is_recurrent: bool = True
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
obs: TensorDict,
|
||||
obs_groups: dict[str, list[str]],
|
||||
num_actions: int,
|
||||
actor_obs_normalization: bool = False,
|
||||
critic_obs_normalization: bool = False,
|
||||
actor_hidden_dims: tuple[int] | list[int] = [256, 256, 256],
|
||||
critic_hidden_dims: tuple[int] | list[int] = [256, 256, 256],
|
||||
activation: str = "elu",
|
||||
init_noise_std: float = 1.0,
|
||||
noise_std_type: str = "scalar",
|
||||
state_dependent_std: bool = False,
|
||||
rnn_type: str = "lstm",
|
||||
rnn_hidden_dim: int = 256,
|
||||
rnn_num_layers: int = 1,
|
||||
**kwargs: dict[str, Any],
|
||||
) -> None:
|
||||
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]
|
||||
|
||||
# Actor
|
||||
self.state_dependent_std = state_dependent_std
|
||||
self.memory_a = Memory(num_actor_obs, rnn_hidden_dim, rnn_num_layers, rnn_type)
|
||||
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)
|
||||
print(f"Actor RNN: {self.memory_a}")
|
||||
print(f"Actor MLP: {self.actor}")
|
||||
|
||||
# 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()
|
||||
|
||||
# Critic
|
||||
self.memory_c = Memory(num_critic_obs, rnn_hidden_dim, rnn_num_layers, rnn_type)
|
||||
self.critic = MLP(rnn_hidden_dim, 1, critic_hidden_dims, activation)
|
||||
print(f"Critic RNN: {self.memory_c}")
|
||||
print(f"Critic MLP: {self.critic}")
|
||||
|
||||
# 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()
|
||||
|
||||
# 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
|
||||
# Note: Populated in update_distribution
|
||||
self.distribution = None
|
||||
|
||||
# Disable args validation for speedup
|
||||
Normal.set_default_validate_args(False)
|
||||
|
||||
@property
|
||||
def action_mean(self) -> torch.Tensor:
|
||||
return self.distribution.mean
|
||||
|
||||
@property
|
||||
def action_std(self) -> torch.Tensor:
|
||||
return self.distribution.stddev
|
||||
|
||||
@property
|
||||
def entropy(self) -> torch.Tensor:
|
||||
return self.distribution.entropy().sum(dim=-1)
|
||||
|
||||
def reset(self, dones: torch.Tensor | None = None) -> None:
|
||||
self.memory_a.reset(dones)
|
||||
self.memory_c.reset(dones)
|
||||
|
||||
def forward(self) -> NoReturn:
|
||||
raise NotImplementedError
|
||||
|
||||
def _update_distribution(self, obs: torch.Tensor) -> None:
|
||||
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: TensorDict, masks: torch.Tensor | None = None, hidden_state: HiddenState = None) -> torch.Tensor:
|
||||
obs = self.get_actor_obs(obs)
|
||||
obs = self.actor_obs_normalizer(obs)
|
||||
out_mem = self.memory_a(obs, masks, hidden_state).squeeze(0)
|
||||
self._update_distribution(out_mem)
|
||||
return self.distribution.sample()
|
||||
|
||||
def act_inference(self, obs: TensorDict) -> torch.Tensor:
|
||||
obs = self.get_actor_obs(obs)
|
||||
obs = self.actor_obs_normalizer(obs)
|
||||
out_mem = self.memory_a(obs).squeeze(0)
|
||||
if self.state_dependent_std:
|
||||
return self.actor(out_mem)[..., 0, :]
|
||||
else:
|
||||
return self.actor(out_mem)
|
||||
|
||||
def evaluate(
|
||||
self, obs: TensorDict, masks: torch.Tensor | None = None, hidden_state: HiddenState = None
|
||||
) -> torch.Tensor:
|
||||
obs = self.get_critic_obs(obs)
|
||||
obs = self.critic_obs_normalizer(obs)
|
||||
out_mem = self.memory_c(obs, masks, hidden_state).squeeze(0)
|
||||
return self.critic(out_mem)
|
||||
|
||||
def get_actor_obs(self, obs: TensorDict) -> torch.Tensor:
|
||||
obs_list = [obs[obs_group] for obs_group in self.obs_groups["policy"]]
|
||||
return torch.cat(obs_list, dim=-1)
|
||||
|
||||
def get_critic_obs(self, obs: TensorDict) -> torch.Tensor:
|
||||
obs_list = [obs[obs_group] for obs_group in self.obs_groups["critic"]]
|
||||
return torch.cat(obs_list, dim=-1)
|
||||
|
||||
def get_actions_log_prob(self, actions: torch.Tensor) -> torch.Tensor:
|
||||
return self.distribution.log_prob(actions).sum(dim=-1)
|
||||
|
||||
def get_hidden_states(self) -> tuple[HiddenState, HiddenState]:
|
||||
return self.memory_a.hidden_state, self.memory_c.hidden_state
|
||||
|
||||
def update_normalization(self, obs: TensorDict) -> None:
|
||||
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: dict, strict: bool = True) -> bool:
|
||||
"""Load the parameters of the actor-critic model.
|
||||
|
||||
Args:
|
||||
state_dict: State dictionary of the model.
|
||||
strict: Whether to strictly enforce that the keys in `state_dict` match the keys returned by this module's
|
||||
:meth:`state_dict` function.
|
||||
|
||||
Returns:
|
||||
Whether this training resumes a previous training. This flag is used by the :func:`load` function of
|
||||
:class:`OnPolicyRunner` to determine how to load further parameters (relevant for, e.g., distillation).
|
||||
"""
|
||||
super().load_state_dict(state_dict, strict=strict)
|
||||
return True
|
||||
210
source/rsl_rl/rsl_rl/modules/rnd.py
Normal file
210
source/rsl_rl/rsl_rl/modules/rnd.py
Normal file
@@ -0,0 +1,210 @@
|
||||
# Copyright (c) 2021-2026, 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 tensordict import TensorDict
|
||||
from typing import Any, NoReturn
|
||||
|
||||
from rsl_rl.env import VecEnv
|
||||
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: tuple[int] | list[int],
|
||||
target_hidden_dims: tuple[int] | 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,
|
||||
) -> 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.
|
||||
- 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.
|
||||
obs_groups: Dictionary of observation groups.
|
||||
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.
|
||||
weight: Scaling factor of the intrinsic reward.
|
||||
state_normalization: Whether to normalize the input state.
|
||||
reward_normalization: Whether to normalize the intrinsic reward.
|
||||
device: Device to use.
|
||||
weight_schedule: Type of schedule to use for the RND weight parameter.
|
||||
It is a dictionary with the following keys:
|
||||
|
||||
- "mode": 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": Step at which the weight parameter is set to the final value.
|
||||
- "final_value": Final value of the weight parameter.
|
||||
|
||||
For the "linear" weight schedule, the following parameters are required:
|
||||
- "initial_step": Step at which the weight parameter is set to the initial value.
|
||||
- "final_step": Step at which the weight parameter is set to the final value.
|
||||
- "final_value": 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: TensorDict) -> 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: Any, **kwargs: dict[str, Any]) -> NoReturn:
|
||||
raise RuntimeError("Forward method is not implemented. Use get_intrinsic_reward instead.")
|
||||
|
||||
def train(self, mode: bool = True) -> RandomNetworkDistillation:
|
||||
# Set 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) -> RandomNetworkDistillation:
|
||||
return self.train(False)
|
||||
|
||||
def get_rnd_state(self, obs: TensorDict) -> torch.Tensor:
|
||||
obs_list = [obs[obs_group] for obs_group in self.obs_groups["rnd_state"]]
|
||||
return torch.cat(obs_list, dim=-1)
|
||||
|
||||
def update_normalization(self, obs: TensorDict) -> None:
|
||||
# Normalize the state
|
||||
if self.state_normalization:
|
||||
rnd_state = self.get_rnd_state(obs)
|
||||
self.state_normalizer.update(rnd_state)
|
||||
|
||||
def _constant_weight_schedule(self, step: int, **kwargs: dict[str, Any]) -> float:
|
||||
return self.initial_weight
|
||||
|
||||
def _step_weight_schedule(self, step: int, final_step: int, final_value: float, **kwargs: dict[str, Any]) -> float:
|
||||
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: dict[str, Any]
|
||||
) -> float:
|
||||
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: dict, obs: TensorDict, obs_groups: dict[str, list[str]], env: VecEnv) -> dict:
|
||||
"""Resolve the RND configuration.
|
||||
|
||||
Args:
|
||||
alg_cfg: Algorithm configuration dictionary.
|
||||
obs: Observation dictionary.
|
||||
obs_groups: Observation groups dictionary.
|
||||
env: Environment object.
|
||||
|
||||
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
|
||||
else:
|
||||
alg_cfg["rnd_cfg"] = None
|
||||
return alg_cfg
|
||||
204
source/rsl_rl/rsl_rl/modules/student_teacher.py
Normal file
204
source/rsl_rl/rsl_rl/modules/student_teacher.py
Normal file
@@ -0,0 +1,204 @@
|
||||
# Copyright (c) 2021-2026, 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 tensordict import TensorDict
|
||||
from torch.distributions import Normal
|
||||
from typing import Any, NoReturn
|
||||
|
||||
from rsl_rl.networks import MLP, EmpiricalNormalization, HiddenState
|
||||
|
||||
|
||||
class StudentTeacher(nn.Module):
|
||||
is_recurrent: bool = False
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
obs: TensorDict,
|
||||
obs_groups: dict[str, list[str]],
|
||||
num_actions: int,
|
||||
student_obs_normalization: bool = False,
|
||||
teacher_obs_normalization: bool = False,
|
||||
student_hidden_dims: tuple[int] | list[int] = [256, 256, 256],
|
||||
teacher_hidden_dims: tuple[int] | list[int] = [256, 256, 256],
|
||||
activation: str = "elu",
|
||||
init_noise_std: float = 0.1,
|
||||
noise_std_type: str = "scalar",
|
||||
**kwargs: dict[str, Any],
|
||||
) -> None:
|
||||
if kwargs:
|
||||
print(
|
||||
"StudentTeacher.__init__ got unexpected arguments, which will be ignored: "
|
||||
+ str([key for key in kwargs])
|
||||
)
|
||||
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)
|
||||
print(f"Student MLP: {self.student}")
|
||||
|
||||
# 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()
|
||||
|
||||
# Teacher
|
||||
self.teacher = MLP(num_teacher_obs, num_actions, teacher_hidden_dims, activation)
|
||||
print(f"Teacher MLP: {self.teacher}")
|
||||
|
||||
# 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()
|
||||
|
||||
# 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
|
||||
# Note: Populated in update_distribution
|
||||
self.distribution = None
|
||||
|
||||
# Disable args validation for speedup
|
||||
Normal.set_default_validate_args(False)
|
||||
|
||||
def reset(
|
||||
self, dones: torch.Tensor | None = None, hidden_states: tuple[HiddenState, HiddenState] = (None, None)
|
||||
) -> None:
|
||||
pass
|
||||
|
||||
def forward(self) -> NoReturn:
|
||||
raise NotImplementedError
|
||||
|
||||
@property
|
||||
def action_mean(self) -> torch.Tensor:
|
||||
return self.distribution.mean
|
||||
|
||||
@property
|
||||
def action_std(self) -> torch.Tensor:
|
||||
return self.distribution.stddev
|
||||
|
||||
@property
|
||||
def entropy(self) -> torch.Tensor:
|
||||
return self.distribution.entropy().sum(dim=-1)
|
||||
|
||||
def _update_distribution(self, obs: TensorDict) -> None:
|
||||
# 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: TensorDict) -> torch.Tensor:
|
||||
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: TensorDict) -> torch.Tensor:
|
||||
obs = self.get_student_obs(obs)
|
||||
obs = self.student_obs_normalizer(obs)
|
||||
return self.student(obs)
|
||||
|
||||
def evaluate(self, obs: TensorDict) -> torch.Tensor:
|
||||
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: TensorDict) -> torch.Tensor:
|
||||
obs_list = [obs[obs_group] for obs_group in self.obs_groups["policy"]]
|
||||
return torch.cat(obs_list, dim=-1)
|
||||
|
||||
def get_teacher_obs(self, obs: TensorDict) -> torch.Tensor:
|
||||
obs_list = [obs[obs_group] for obs_group in self.obs_groups["teacher"]]
|
||||
return torch.cat(obs_list, dim=-1)
|
||||
|
||||
def get_hidden_states(self) -> tuple[HiddenState, HiddenState]:
|
||||
return None, None
|
||||
|
||||
def detach_hidden_states(self, dones: torch.Tensor | None = None) -> None:
|
||||
pass
|
||||
|
||||
def train(self, mode: bool = True) -> None:
|
||||
super().train(mode)
|
||||
# Make sure teacher is in eval mode
|
||||
self.teacher.eval()
|
||||
self.teacher_obs_normalizer.eval()
|
||||
|
||||
def update_normalization(self, obs: TensorDict) -> None:
|
||||
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: dict, strict: bool = True) -> bool:
|
||||
"""Load the parameters of the student and teacher networks.
|
||||
|
||||
Args:
|
||||
state_dict: State dictionary of the model.
|
||||
strict: Whether to strictly enforce that the keys in `state_dict` match the keys returned by this module's
|
||||
:meth:`state_dict` function.
|
||||
|
||||
Returns:
|
||||
Whether this training resumes a previous training. This flag is used by the :func:`load` function of
|
||||
:class:`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): # Load 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): # Load 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")
|
||||
244
source/rsl_rl/rsl_rl/modules/student_teacher_recurrent.py
Normal file
244
source/rsl_rl/rsl_rl/modules/student_teacher_recurrent.py
Normal file
@@ -0,0 +1,244 @@
|
||||
# Copyright (c) 2021-2026, 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 tensordict import TensorDict
|
||||
from torch.distributions import Normal
|
||||
from typing import Any, NoReturn
|
||||
|
||||
from rsl_rl.networks import MLP, EmpiricalNormalization, HiddenState, Memory
|
||||
|
||||
|
||||
class StudentTeacherRecurrent(nn.Module):
|
||||
is_recurrent: bool = True
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
obs: TensorDict,
|
||||
obs_groups: dict[str, list[str]],
|
||||
num_actions: int,
|
||||
student_obs_normalization: bool = False,
|
||||
teacher_obs_normalization: bool = False,
|
||||
student_hidden_dims: tuple[int] | list[int] = [256, 256, 256],
|
||||
teacher_hidden_dims: tuple[int] | list[int] = [256, 256, 256],
|
||||
activation: str = "elu",
|
||||
init_noise_std: float = 0.1,
|
||||
noise_std_type: str = "scalar",
|
||||
rnn_type: str = "lstm",
|
||||
rnn_hidden_dim: int = 256,
|
||||
rnn_num_layers: int = 1,
|
||||
teacher_recurrent: bool = False,
|
||||
**kwargs: dict[str, Any],
|
||||
) -> None:
|
||||
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, rnn_hidden_dim, rnn_num_layers, rnn_type)
|
||||
self.student = MLP(rnn_hidden_dim, num_actions, student_hidden_dims, activation)
|
||||
print(f"Student RNN: {self.memory_s}")
|
||||
print(f"Student MLP: {self.student}")
|
||||
|
||||
# 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()
|
||||
|
||||
# Teacher
|
||||
if self.teacher_recurrent:
|
||||
self.memory_t = Memory(num_teacher_obs, rnn_hidden_dim, rnn_num_layers, rnn_type)
|
||||
teacher_input_dim = rnn_hidden_dim if self.teacher_recurrent else num_teacher_obs
|
||||
self.teacher = MLP(teacher_input_dim, num_actions, teacher_hidden_dims, activation)
|
||||
if self.teacher_recurrent:
|
||||
print(f"Teacher RNN: {self.memory_t}")
|
||||
print(f"Teacher MLP: {self.teacher}")
|
||||
|
||||
# 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()
|
||||
|
||||
# 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
|
||||
# Note: Populated in update_distribution
|
||||
self.distribution = None
|
||||
|
||||
# Disable args validation for speedup
|
||||
Normal.set_default_validate_args(False)
|
||||
|
||||
def reset(
|
||||
self, dones: torch.Tensor | None = None, hidden_states: tuple[HiddenState, HiddenState] = (None, 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) -> NoReturn:
|
||||
raise NotImplementedError
|
||||
|
||||
@property
|
||||
def action_mean(self) -> torch.Tensor:
|
||||
return self.distribution.mean
|
||||
|
||||
@property
|
||||
def action_std(self) -> torch.Tensor:
|
||||
return self.distribution.stddev
|
||||
|
||||
@property
|
||||
def entropy(self) -> torch.Tensor:
|
||||
return self.distribution.entropy().sum(dim=-1)
|
||||
|
||||
def _update_distribution(self, obs: TensorDict) -> None:
|
||||
# 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: TensorDict) -> torch.Tensor:
|
||||
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: TensorDict) -> torch.Tensor:
|
||||
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: TensorDict) -> torch.Tensor:
|
||||
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: TensorDict) -> torch.Tensor:
|
||||
obs_list = [obs[obs_group] for obs_group in self.obs_groups["policy"]]
|
||||
return torch.cat(obs_list, dim=-1)
|
||||
|
||||
def get_teacher_obs(self, obs: TensorDict) -> torch.Tensor:
|
||||
obs_list = [obs[obs_group] for obs_group in self.obs_groups["teacher"]]
|
||||
return torch.cat(obs_list, dim=-1)
|
||||
|
||||
def get_hidden_states(self) -> tuple[HiddenState, HiddenState]:
|
||||
if self.teacher_recurrent:
|
||||
return self.memory_s.hidden_state, self.memory_t.hidden_state
|
||||
else:
|
||||
return self.memory_s.hidden_state, None
|
||||
|
||||
def detach_hidden_states(self, dones: torch.Tensor | None = None) -> None:
|
||||
self.memory_s.detach_hidden_state(dones)
|
||||
if self.teacher_recurrent:
|
||||
self.memory_t.detach_hidden_state(dones)
|
||||
|
||||
def train(self, mode: bool = True) -> None:
|
||||
super().train(mode)
|
||||
# Make sure teacher is in eval mode
|
||||
self.teacher.eval()
|
||||
self.teacher_obs_normalizer.eval()
|
||||
|
||||
def update_normalization(self, obs: TensorDict) -> None:
|
||||
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: dict, strict: bool = True) -> bool:
|
||||
"""Load the parameters of the student and teacher networks.
|
||||
|
||||
Args:
|
||||
state_dict: State dictionary of the model.
|
||||
strict: Whether to strictly enforce that the keys in `state_dict` match the keys returned by this module's
|
||||
:meth:`state_dict` function.
|
||||
|
||||
Returns:
|
||||
Whether this training resumes a previous training. This flag is used by the :func:`load` function of
|
||||
:class:`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): # Load 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): # Load 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")
|
||||
27
source/rsl_rl/rsl_rl/modules/symmetry.py
Normal file
27
source/rsl_rl/rsl_rl/modules/symmetry.py
Normal file
@@ -0,0 +1,27 @@
|
||||
# Copyright (c) 2021-2026, ETH Zurich and NVIDIA CORPORATION
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from rsl_rl.env import VecEnv
|
||||
|
||||
|
||||
def resolve_symmetry_config(alg_cfg: dict, env: VecEnv) -> dict:
|
||||
"""Resolve the symmetry configuration.
|
||||
|
||||
Args:
|
||||
alg_cfg: Algorithm configuration dictionary.
|
||||
env: Environment object.
|
||||
|
||||
Returns:
|
||||
The resolved algorithm configuration dictionary.
|
||||
"""
|
||||
# If using symmetry then pass the environment config object
|
||||
# Note: This is used by the symmetry function for handling different observation terms
|
||||
if "symmetry_cfg" in alg_cfg and alg_cfg["symmetry_cfg"] is not None:
|
||||
alg_cfg["symmetry_cfg"]["_env"] = env
|
||||
else:
|
||||
alg_cfg["symmetry_cfg"] = None
|
||||
return alg_cfg
|
||||
23
source/rsl_rl/rsl_rl/networks/__init__.py
Normal file
23
source/rsl_rl/rsl_rl/networks/__init__.py
Normal file
@@ -0,0 +1,23 @@
|
||||
# Copyright (c) 2021-2026, ETH Zurich and NVIDIA CORPORATION
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
|
||||
"""Definitions for components of modules."""
|
||||
|
||||
from .cnn import CNN
|
||||
from .memory import HiddenState, Memory
|
||||
from .mlp import MLP
|
||||
from .normalization import EmpiricalDiscountedVariationNormalization, EmpiricalNormalization
|
||||
from .moe import L2Norm, SimNorm, MoE
|
||||
__all__ = [
|
||||
"CNN",
|
||||
"MLP",
|
||||
"EmpiricalDiscountedVariationNormalization",
|
||||
"EmpiricalNormalization",
|
||||
"HiddenState",
|
||||
"Memory",
|
||||
"L2Norm",
|
||||
"SimNorm",
|
||||
"MoE",
|
||||
]
|
||||
192
source/rsl_rl/rsl_rl/networks/cnn.py
Normal file
192
source/rsl_rl/rsl_rl/networks/cnn.py
Normal file
@@ -0,0 +1,192 @@
|
||||
# Copyright (c) 2021-2026, ETH Zurich and NVIDIA CORPORATION
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
import torch
|
||||
from torch import nn as nn
|
||||
|
||||
from rsl_rl.utils import get_param, resolve_nn_activation
|
||||
|
||||
|
||||
class CNN(nn.Sequential):
|
||||
"""Convolutional Neural Network (CNN).
|
||||
|
||||
The CNN network is a sequence of convolutional layers, optional normalization layers, optional activation functions,
|
||||
and optional pooling. The final output can be flattened.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
input_dim: tuple[int, int],
|
||||
input_channels: int,
|
||||
output_channels: tuple[int] | list[int],
|
||||
kernel_size: int | tuple[int] | list[int],
|
||||
stride: int | tuple[int] | list[int] = 1,
|
||||
dilation: int | tuple[int] | list[int] = 1,
|
||||
padding: str = "none",
|
||||
norm: str | tuple[str] | list[str] = "none",
|
||||
activation: str = "elu",
|
||||
max_pool: bool | tuple[bool] | list[bool] = False,
|
||||
global_pool: str = "none",
|
||||
flatten: bool = True,
|
||||
) -> None:
|
||||
"""Initialize the CNN.
|
||||
|
||||
Args:
|
||||
input_dim: Height and width of the input.
|
||||
input_channels: Number of input channels.
|
||||
output_channels: List of output channels for each convolutional layer.
|
||||
kernel_size: List of kernel sizes for each convolutional layer or a single kernel size for all layers.
|
||||
stride: List of strides for each convolutional layer or a single stride for all layers.
|
||||
dilation: List of dilations for each convolutional layer or a single dilation for all layers.
|
||||
padding: Padding type to use. Either 'none', 'zeros', 'reflect', 'replicate', or 'circular'.
|
||||
norm: List of normalization types for each convolutional layer or a single type for all layers. Either
|
||||
'none', 'batch', or 'layer'.
|
||||
activation: Activation function to use.
|
||||
max_pool: List of booleans indicating whether to apply max pooling after each convolutional layer or a
|
||||
single boolean for all layers.
|
||||
global_pool: Global pooling type to apply at the end. Either 'none', 'max', or 'avg'.
|
||||
flatten: Whether to flatten the output tensor.
|
||||
"""
|
||||
super().__init__()
|
||||
|
||||
# Resolve activation function
|
||||
activation_function = resolve_nn_activation(activation)
|
||||
|
||||
# Create layers sequentially
|
||||
layers = []
|
||||
last_channels = input_channels
|
||||
last_dim = input_dim
|
||||
for idx in range(len(output_channels)):
|
||||
# Get parameters for the current layer
|
||||
k = get_param(kernel_size, idx)
|
||||
s = get_param(stride, idx)
|
||||
d = get_param(dilation, idx)
|
||||
p = (
|
||||
_compute_padding(last_dim, k, s, d)
|
||||
if padding in ["zeros", "reflect", "replicate", "circular"]
|
||||
else (0, 0)
|
||||
)
|
||||
|
||||
# Append convolutional layer
|
||||
layers.append(
|
||||
nn.Conv2d(
|
||||
in_channels=last_channels,
|
||||
out_channels=output_channels[idx],
|
||||
kernel_size=k,
|
||||
stride=s,
|
||||
padding=p,
|
||||
dilation=d,
|
||||
padding_mode=padding if padding in ["zeros", "reflect", "replicate", "circular"] else "zeros",
|
||||
)
|
||||
)
|
||||
|
||||
# Append normalization layer if specified
|
||||
n = get_param(norm, idx)
|
||||
if n == "none":
|
||||
pass
|
||||
elif n == "batch":
|
||||
layers.append(nn.BatchNorm2d(output_channels[idx]))
|
||||
elif n == "layer":
|
||||
norm_input_dim = _compute_output_dim(last_dim, k, s, d, p)
|
||||
layers.append(nn.LayerNorm([output_channels[idx], norm_input_dim[0], norm_input_dim[1]]))
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unsupported normalization type: {n}. Supported types are 'none', 'batch', and 'layer'."
|
||||
)
|
||||
|
||||
# Append activation function
|
||||
layers.append(activation_function)
|
||||
|
||||
# Apply max pooling if specified
|
||||
if get_param(max_pool, idx):
|
||||
layers.append(nn.MaxPool2d(kernel_size=3, stride=2, padding=1))
|
||||
|
||||
# Update last channels and dimensions
|
||||
last_channels = output_channels[idx]
|
||||
last_dim = _compute_output_dim(last_dim, k, s, d, p, is_max_pool=get_param(max_pool, idx))
|
||||
|
||||
# Apply global pooling if specified
|
||||
if global_pool == "none":
|
||||
pass
|
||||
elif global_pool == "max":
|
||||
layers.append(nn.AdaptiveMaxPool2d((1, 1)))
|
||||
last_dim = (1, 1)
|
||||
elif global_pool == "avg":
|
||||
layers.append(nn.AdaptiveAvgPool2d((1, 1)))
|
||||
last_dim = (1, 1)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unsupported global pooling type: {global_pool}. Supported types are 'none', 'max', and 'avg'."
|
||||
)
|
||||
|
||||
# Apply flattening if specified
|
||||
if flatten:
|
||||
layers.append(nn.Flatten(start_dim=1))
|
||||
|
||||
# Store final output dimension
|
||||
self._output_channels = last_channels if not flatten else None
|
||||
self._output_dim = last_dim if not flatten else last_channels * last_dim[0] * last_dim[1]
|
||||
|
||||
# Register the layers
|
||||
for idx, layer in enumerate(layers):
|
||||
self.add_module(f"{idx}", layer)
|
||||
|
||||
@property
|
||||
def output_channels(self) -> int | None:
|
||||
"""Get the number of output channels or None if output is flattened."""
|
||||
return self._output_channels
|
||||
|
||||
@property
|
||||
def output_dim(self) -> tuple[int, int] | int:
|
||||
"""Get the output height and width or total output dimension if output is flattened."""
|
||||
return self._output_dim
|
||||
|
||||
def init_weights(self) -> None:
|
||||
"""Initialize the weights of the CNN with Kaiming initialization."""
|
||||
for idx, module in enumerate(self):
|
||||
if isinstance(module, nn.Conv2d):
|
||||
torch.nn.init.kaiming_normal_(module.weight)
|
||||
torch.nn.init.zeros_(module.bias)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
"""Forward pass of the CNN."""
|
||||
for layer in self:
|
||||
x = layer(x)
|
||||
return x
|
||||
|
||||
|
||||
def _compute_padding(input_hw: tuple[int, int], kernel: int, stride: int, dilation: int) -> tuple[int, int]:
|
||||
"""Compute the optimal padding for the current layer.
|
||||
|
||||
Reference: https://pytorch.org/docs/stable/generated/torch.nn.Conv2d.html
|
||||
"""
|
||||
h = math.ceil((stride * math.floor(input_hw[0] / stride) - input_hw[0] - stride + dilation * (kernel - 1) + 1) / 2)
|
||||
w = math.ceil((stride * math.floor(input_hw[1] / stride) - input_hw[1] - stride + dilation * (kernel - 1) + 1) / 2)
|
||||
return (h, w)
|
||||
|
||||
|
||||
def _compute_output_dim(
|
||||
input_hw: tuple[int, int],
|
||||
kernel: int,
|
||||
stride: int,
|
||||
dilation: int,
|
||||
padding: tuple[int, int],
|
||||
is_max_pool: bool = False,
|
||||
) -> tuple[int, int]:
|
||||
"""Compute the output height and width of the current layer.
|
||||
|
||||
Reference: https://pytorch.org/docs/stable/generated/torch.nn.Conv2d.html
|
||||
"""
|
||||
h = math.floor((input_hw[0] + 2 * padding[0] - dilation * (kernel - 1) - 1) / stride + 1)
|
||||
w = math.floor((input_hw[1] + 2 * padding[1] - dilation * (kernel - 1) - 1) / stride + 1)
|
||||
|
||||
if is_max_pool:
|
||||
h = math.ceil(h / 2)
|
||||
w = math.ceil(w / 2)
|
||||
|
||||
return (h, w)
|
||||
81
source/rsl_rl/rsl_rl/networks/memory.py
Normal file
81
source/rsl_rl/rsl_rl/networks/memory.py
Normal file
@@ -0,0 +1,81 @@
|
||||
# Copyright (c) 2021-2026, 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 typing import Union
|
||||
|
||||
from rsl_rl.utils import unpad_trajectories
|
||||
|
||||
HiddenState = Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor], None] # Using Union due to Python <3.10
|
||||
"""Type alias for the hidden state of RNNs (GRU/LSTM).
|
||||
|
||||
For GRUs, this is a single tensor while for LSTMs, this is a tuple of two tensors (hidden state and cell state).
|
||||
"""
|
||||
|
||||
|
||||
class Memory(nn.Module):
|
||||
"""Memory network for recurrent architectures.
|
||||
|
||||
This network is used to store the hidden state of the policy. It currently supports GRU and LSTM.
|
||||
"""
|
||||
|
||||
def __init__(self, input_size: int, hidden_dim: int = 256, num_layers: int = 1, type: str = "lstm") -> None:
|
||||
super().__init__()
|
||||
rnn_cls = nn.GRU if type.lower() == "gru" else nn.LSTM
|
||||
self.rnn = rnn_cls(input_size=input_size, hidden_size=hidden_dim, num_layers=num_layers)
|
||||
self.hidden_state = None
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input: torch.Tensor,
|
||||
masks: torch.Tensor | None = None,
|
||||
hidden_state: HiddenState = None,
|
||||
) -> torch.Tensor:
|
||||
batch_mode = masks is not None
|
||||
if batch_mode:
|
||||
# Batch mode needs saved hidden states
|
||||
if hidden_state is None:
|
||||
raise ValueError("Hidden states not passed to memory module during policy update")
|
||||
out, _ = self.rnn(input, hidden_state)
|
||||
out = unpad_trajectories(out, masks)
|
||||
else:
|
||||
# Inference/distillation mode uses hidden state of last step
|
||||
out, self.hidden_state = self.rnn(input.unsqueeze(0), self.hidden_state)
|
||||
return out
|
||||
|
||||
def reset(self, dones: torch.Tensor | None = None, hidden_state: HiddenState = None) -> None:
|
||||
if dones is None: # Reset hidden state
|
||||
if hidden_state is None:
|
||||
self.hidden_state = None
|
||||
else:
|
||||
self.hidden_state = hidden_state
|
||||
elif self.hidden_state is not None: # Reset hidden state of done environments
|
||||
if hidden_state is None:
|
||||
if isinstance(self.hidden_state, tuple): # Tuple in case of LSTM
|
||||
for hidden_state in self.hidden_state:
|
||||
hidden_state[..., dones == 1, :] = 0.0
|
||||
else:
|
||||
self.hidden_state[..., dones == 1, :] = 0.0
|
||||
else:
|
||||
NotImplementedError(
|
||||
"Resetting the hidden state of done environments with a custom hidden state is not implemented"
|
||||
)
|
||||
|
||||
def detach_hidden_state(self, dones: torch.Tensor | None = None) -> None:
|
||||
if self.hidden_state is not None:
|
||||
if dones is None: # Detach hidden state
|
||||
if isinstance(self.hidden_state, tuple): # Tuple in case of LSTM
|
||||
self.hidden_state = tuple(hidden_state.detach() for hidden_state in self.hidden_state)
|
||||
else:
|
||||
self.hidden_state = self.hidden_state.detach()
|
||||
else: # Detach hidden state of done environments
|
||||
if isinstance(self.hidden_state, tuple): # Tuple in case of LSTM
|
||||
for hidden_state in self.hidden_state:
|
||||
hidden_state[..., dones == 1, :] = hidden_state[..., dones == 1, :].detach()
|
||||
else:
|
||||
self.hidden_state[..., dones == 1, :] = self.hidden_state[..., dones == 1, :].detach()
|
||||
94
source/rsl_rl/rsl_rl/networks/mlp.py
Normal file
94
source/rsl_rl/rsl_rl/networks/mlp.py
Normal file
@@ -0,0 +1,94 @@
|
||||
# Copyright (c) 2021-2026, 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 get_param, 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,
|
||||
) -> 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.
|
||||
last_activation: Activation function of the last layer. None results in a linear last layer.
|
||||
"""
|
||||
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]) -> None:
|
||||
"""Initialize the weights of the MLP.
|
||||
|
||||
Args:
|
||||
scales: Scale factor for the weights.
|
||||
"""
|
||||
for idx, module in enumerate(self):
|
||||
if isinstance(module, nn.Linear):
|
||||
nn.init.orthogonal_(module.weight, gain=get_param(scales, idx))
|
||||
nn.init.zeros_(module.bias)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
"""Forward pass of the MLP."""
|
||||
for layer in self:
|
||||
x = layer(x)
|
||||
return x
|
||||
170
source/rsl_rl/rsl_rl/networks/moe.py
Normal file
170
source/rsl_rl/rsl_rl/networks/moe.py
Normal file
@@ -0,0 +1,170 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import numpy as np
|
||||
|
||||
class CatELU(nn.Module):
|
||||
"""
|
||||
CatELU activation (feature-doubling version).
|
||||
|
||||
Applies ELU to the input and its negation,
|
||||
doubling the feature dimension.
|
||||
|
||||
Output: [..., 2 * D] given input [..., D]
|
||||
|
||||
NOTE:
|
||||
This is a structural activation and NOT element-wise.
|
||||
Assumes the last dimension is the feature dimension.
|
||||
"""
|
||||
def __init__(self, inplace: bool = False):
|
||||
super().__init__()
|
||||
self.elu = nn.ELU(inplace=inplace)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
assert x.dim() >= 2, \
|
||||
f"CatELU expects feature dimension in the last axis, got shape {x.shape}"
|
||||
|
||||
y1 = self.elu(x)
|
||||
y2 = self.elu(-x)
|
||||
return torch.cat((y1, y2), dim=-1)
|
||||
|
||||
def get_activation(act_name):
|
||||
if act_name == "elu":
|
||||
return nn.ELU()
|
||||
elif act_name == "selu":
|
||||
return nn.SELU()
|
||||
elif act_name == "relu":
|
||||
return nn.ReLU()
|
||||
elif act_name == "crelu":
|
||||
return nn.ReLU()
|
||||
elif act_name == "lrelu":
|
||||
return nn.LeakyReLU()
|
||||
elif act_name == "tanh":
|
||||
return nn.Tanh()
|
||||
elif act_name == "sigmoid":
|
||||
return nn.Sigmoid()
|
||||
elif act_name == "cat_elu":
|
||||
return CatELU()
|
||||
else:
|
||||
print("invalid activation function!")
|
||||
return None
|
||||
|
||||
class L2Norm(nn.Module):
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
def forward(self, x):
|
||||
return F.normalize(x, p=2.0, dim=-1)
|
||||
|
||||
class SimNorm(nn.Module):
|
||||
"""
|
||||
Simplicial normalization.
|
||||
Adapted from https://arxiv.org/abs/2204.00616.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.dim = 8 # for latent dim 512
|
||||
|
||||
def forward(self, x):
|
||||
shp = x.shape
|
||||
x = x.view(*shp[:-1], -1, self.dim)
|
||||
x = F.softmax(x, dim=-1)
|
||||
return x.view(*shp)
|
||||
|
||||
def __repr__(self):
|
||||
return f"SimNorm(dim={self.dim})"
|
||||
|
||||
# MLP implementation for MoE
|
||||
class MLP(nn.Module):
|
||||
def __init__(self, input_dim, output_dim, hidden_dims, activation='elu', last_activation: str | None = None):
|
||||
super().__init__()
|
||||
|
||||
dims = [input_dim] + hidden_dims
|
||||
act_func = get_activation(activation)
|
||||
layers = []
|
||||
last_dim = dims[0]
|
||||
for h_dim in dims[1:]:
|
||||
layers.append(nn.Linear(last_dim, h_dim))
|
||||
layers.append(act_func)
|
||||
if activation == 'cat_elu':
|
||||
last_dim = h_dim * 2
|
||||
else:
|
||||
last_dim = h_dim
|
||||
|
||||
if isinstance(output_dim, int):
|
||||
layers.append(nn.Linear(last_dim, output_dim))
|
||||
elif isinstance(output_dim, tuple) or isinstance(output_dim, list):
|
||||
layers.append(nn.Linear(last_dim, np.prod(output_dim)))
|
||||
layers.append(nn.Unflatten(dim=-1, unflattened_size=output_dim))
|
||||
else:
|
||||
raise ValueError("output_dim must be int, tuple or list")
|
||||
|
||||
if last_activation is not None:
|
||||
last_act_func = get_activation(last_activation)
|
||||
layers.append(last_act_func)
|
||||
self.network = nn.Sequential(*layers)
|
||||
|
||||
def forward(self, x):
|
||||
return self.network(x)
|
||||
|
||||
class Experts(nn.Module):
|
||||
def __init__(self,
|
||||
expert_num,
|
||||
input_dim,
|
||||
backbone_hidden_dims,
|
||||
expert_hidden_dim,
|
||||
output_dim,
|
||||
activation='elu',
|
||||
):
|
||||
super().__init__()
|
||||
self.expert_num = expert_num
|
||||
self.output_dim = output_dim
|
||||
|
||||
self.backbone = MLP(input_dim, expert_num * expert_hidden_dim, backbone_hidden_dims, activation, last_activation=activation)
|
||||
self.experts = nn.Conv1d(
|
||||
in_channels=expert_num*expert_hidden_dim if activation != 'cat_elu' else expert_num*expert_hidden_dim*2,
|
||||
out_channels=expert_num*output_dim,
|
||||
kernel_size=1,
|
||||
groups=expert_num,
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
shared_features = self.backbone(x).unsqueeze(-1) # (B, expert_num * expert_hidden_dim, 1)
|
||||
expert_outs = self.experts(shared_features).squeeze(-1) # (B, expert_num * output_dim)
|
||||
expert_outs = expert_outs.reshape(-1, self.expert_num, self.output_dim)
|
||||
return expert_outs
|
||||
|
||||
class MoE(nn.Module):
|
||||
def __init__(self,
|
||||
expert_num,
|
||||
input_dim,
|
||||
hidden_dims,
|
||||
output_dim,
|
||||
activation='elu',
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
# Expert networks
|
||||
self.experts = Experts(
|
||||
expert_num=expert_num,
|
||||
input_dim=input_dim,
|
||||
backbone_hidden_dims=hidden_dims[:-1],
|
||||
expert_hidden_dim=hidden_dims[-1],
|
||||
output_dim=output_dim,
|
||||
activation=activation,
|
||||
)
|
||||
|
||||
# Gating network
|
||||
self.gating_network = nn.Sequential(
|
||||
MLP(input_dim, expert_num, hidden_dims, activation),
|
||||
nn.Softmax(dim=-1)
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
weights = self.gating_network(x) # (B, expert_num)
|
||||
expert_outs = self.experts(x) # (B, expert_num, output_dim)
|
||||
output = torch.sum(weights.unsqueeze(-1) * expert_outs, dim=1) # (B, output_dim)
|
||||
return output, weights
|
||||
|
||||
120
source/rsl_rl/rsl_rl/networks/normalization.py
Normal file
120
source/rsl_rl/rsl_rl/networks/normalization.py
Normal file
@@ -0,0 +1,120 @@
|
||||
# Copyright (c) 2021-2026, 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: int | tuple[int] | list[int], eps: float = 1e-2, until: int | None = None) -> None:
|
||||
"""Initialize EmpiricalNormalization module.
|
||||
|
||||
.. note:: The normalization parameters are computed over the whole batch, not for each environment separately.
|
||||
|
||||
Args:
|
||||
shape: Shape of input values except batch axis.
|
||||
eps: Small value for stability.
|
||||
until: If this arg is specified, the module learns input values until the sum of batch sizes exceeds it.
|
||||
"""
|
||||
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) -> torch.Tensor:
|
||||
return self._mean.squeeze(0).clone()
|
||||
|
||||
@property
|
||||
def std(self) -> torch.Tensor:
|
||||
return self._std.squeeze(0).clone()
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
"""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: torch.Tensor) -> None:
|
||||
"""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: torch.Tensor) -> torch.Tensor:
|
||||
"""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: int | tuple[int] | list[int], eps: float = 1e-2, gamma: float = 0.99, until: int | None = None
|
||||
) -> None:
|
||||
super().__init__()
|
||||
|
||||
self.emp_norm = EmpiricalNormalization(shape, eps, until)
|
||||
self.disc_avg = _DiscountedAverage(gamma)
|
||||
|
||||
def forward(self, rew: torch.Tensor) -> torch.Tensor:
|
||||
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
|
||||
|
||||
|
||||
class _DiscountedAverage:
|
||||
r"""Discounted average of rewards.
|
||||
|
||||
The discounted average is defined as:
|
||||
|
||||
.. math::
|
||||
|
||||
\bar{R}_t = \gamma \bar{R}_{t-1} + r_t
|
||||
"""
|
||||
|
||||
def __init__(self, gamma: float) -> None:
|
||||
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
|
||||
12
source/rsl_rl/rsl_rl/runners/__init__.py
Normal file
12
source/rsl_rl/rsl_rl/runners/__init__.py
Normal file
@@ -0,0 +1,12 @@
|
||||
# Copyright (c) 2021-2026, 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 # noqa: I001
|
||||
from .distillation_runner import DistillationRunner
|
||||
from .on_policy_runner_cts import OnPolicyRunnerCTS
|
||||
|
||||
__all__ = ["DistillationRunner", "OnPolicyRunner", "OnPolicyRunnerCTS"]
|
||||
57
source/rsl_rl/rsl_rl/runners/distillation_runner.py
Normal file
57
source/rsl_rl/rsl_rl/runners/distillation_runner.py
Normal file
@@ -0,0 +1,57 @@
|
||||
# Copyright (c) 2021-2026, ETH Zurich and NVIDIA CORPORATION
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from tensordict import TensorDict
|
||||
|
||||
from rsl_rl.algorithms import Distillation
|
||||
from rsl_rl.modules import StudentTeacher, StudentTeacherRecurrent
|
||||
from rsl_rl.runners import OnPolicyRunner
|
||||
from rsl_rl.storage import RolloutStorage
|
||||
from rsl_rl.utils import resolve_callable
|
||||
|
||||
|
||||
class DistillationRunner(OnPolicyRunner):
|
||||
"""Distillation runner for training and evaluation of teacher-student methods."""
|
||||
|
||||
def learn(self, num_learning_iterations: int, init_at_random_ep_len: bool = False) -> None:
|
||||
# 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.")
|
||||
|
||||
super().learn(num_learning_iterations, init_at_random_ep_len)
|
||||
|
||||
def _get_default_obs_sets(self) -> list[str]:
|
||||
"""Get the the default observation sets required for the algorithm.
|
||||
|
||||
.. note::
|
||||
See :func:`resolve_obs_groups` for more details on the handling of observation sets.
|
||||
"""
|
||||
return ["teacher"]
|
||||
|
||||
def _construct_algorithm(self, obs: TensorDict) -> Distillation:
|
||||
"""Construct the distillation algorithm."""
|
||||
# Initialize the policy
|
||||
student_teacher_class = resolve_callable(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 storage
|
||||
storage = RolloutStorage(
|
||||
"distillation", self.env.num_envs, self.cfg["num_steps_per_env"], obs, [self.env.num_actions], self.device
|
||||
)
|
||||
|
||||
# Initialize the algorithm
|
||||
alg_class = resolve_callable(self.alg_cfg.pop("class_name"))
|
||||
alg: Distillation = alg_class(
|
||||
student_teacher, storage, device=self.device, **self.alg_cfg, multi_gpu_cfg=self.multi_gpu_cfg
|
||||
)
|
||||
|
||||
# Set RND configuration to None as it does not apply to distillation
|
||||
self.cfg["algorithm"]["rnd_cfg"] = None
|
||||
|
||||
return alg
|
||||
286
source/rsl_rl/rsl_rl/runners/on_policy_runner.py
Normal file
286
source/rsl_rl/rsl_rl/runners/on_policy_runner.py
Normal file
@@ -0,0 +1,286 @@
|
||||
# Copyright (c) 2021-2026, ETH Zurich and NVIDIA CORPORATION
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import time
|
||||
import torch
|
||||
import warnings
|
||||
from tensordict import TensorDict
|
||||
|
||||
from rsl_rl.algorithms import PPO
|
||||
from rsl_rl.env import VecEnv
|
||||
from rsl_rl.modules import (
|
||||
ActorCritic,
|
||||
ActorCriticCNN,
|
||||
ActorCriticRecurrent,
|
||||
resolve_rnd_config,
|
||||
resolve_symmetry_config,
|
||||
)
|
||||
from rsl_rl.storage import RolloutStorage
|
||||
from rsl_rl.utils import resolve_callable, resolve_obs_groups
|
||||
from rsl_rl.utils.logger import Logger
|
||||
|
||||
|
||||
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: str = "cpu") -> None:
|
||||
self.cfg = train_cfg
|
||||
self.policy_cfg = train_cfg["policy"]
|
||||
self.alg_cfg = train_cfg["algorithm"]
|
||||
self.device = device
|
||||
self.env = env
|
||||
|
||||
# Setup multi-GPU training if enabled
|
||||
self._configure_multi_gpu()
|
||||
|
||||
# Query observations from environment for algorithm construction
|
||||
obs = self.env.get_observations()
|
||||
self.cfg["obs_groups"] = resolve_obs_groups(obs, self.cfg["obs_groups"], self._get_default_obs_sets())
|
||||
|
||||
# Create the algorithm
|
||||
self.alg = self._construct_algorithm(obs)
|
||||
|
||||
# Create the logger
|
||||
self.logger = Logger(
|
||||
log_dir=log_dir,
|
||||
cfg=self.cfg,
|
||||
env_cfg=self.env.cfg,
|
||||
num_envs=self.env.num_envs,
|
||||
is_distributed=self.is_distributed,
|
||||
gpu_world_size=self.gpu_world_size,
|
||||
gpu_global_rank=self.gpu_global_rank,
|
||||
device=self.device,
|
||||
)
|
||||
|
||||
self.current_learning_iteration = 0
|
||||
|
||||
def learn(self, num_learning_iterations: int, init_at_random_ep_len: bool = False) -> None:
|
||||
# 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)
|
||||
|
||||
# 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_it = self.current_learning_iteration
|
||||
total_it = start_it + num_learning_iterations
|
||||
for it in range(start_it, total_it):
|
||||
start = time.time()
|
||||
# Rollout
|
||||
with torch.inference_mode():
|
||||
for _ in range(self.cfg["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_cfg["rnd_cfg"] else None
|
||||
# Book keeping
|
||||
self.logger.process_env_step(rewards, dones, extras, intrinsic_rewards)
|
||||
|
||||
stop = time.time()
|
||||
collect_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 information
|
||||
self.logger.log(
|
||||
it=it,
|
||||
start_it=start_it,
|
||||
total_it=total_it,
|
||||
collect_time=collect_time,
|
||||
learn_time=learn_time,
|
||||
loss_dict=loss_dict,
|
||||
learning_rate=self.alg.learning_rate,
|
||||
action_std=self.alg.policy.action_std,
|
||||
rnd_weight=self.alg.rnd.weight if self.alg_cfg["rnd_cfg"] else None,
|
||||
)
|
||||
|
||||
# Save model
|
||||
if it % self.cfg["save_interval"] == 0:
|
||||
self.save(os.path.join(self.logger.log_dir, f"model_{it}.pt")) # type: ignore
|
||||
|
||||
# Save the final model after training
|
||||
if self.logger.log_dir is not None and not self.logger.disable_logs:
|
||||
self.save(os.path.join(self.logger.log_dir, f"model_{self.current_learning_iteration}.pt"))
|
||||
|
||||
def save(self, path: str, infos: dict | None = None) -> 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 self.alg_cfg["rnd_cfg"]:
|
||||
saved_dict["rnd_state_dict"] = self.alg.rnd.state_dict()
|
||||
if self.alg.rnd_optimizer:
|
||||
saved_dict["rnd_optimizer_state_dict"] = self.alg.rnd_optimizer.state_dict()
|
||||
torch.save(saved_dict, path)
|
||||
|
||||
# Upload model to external logging services
|
||||
self.logger.save_model(path, self.current_learning_iteration)
|
||||
|
||||
def load(self, path: str, load_optimizer: bool = True, map_location: str | None = None) -> dict:
|
||||
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 self.alg_cfg["rnd_cfg"]:
|
||||
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 self.alg_cfg["rnd_cfg"]:
|
||||
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: str | None = None) -> callable:
|
||||
self.eval_mode() # Switch to evaluation mode (e.g. for dropout)
|
||||
if device is not None:
|
||||
self.alg.policy.to(device)
|
||||
return self.alg.policy.act_inference
|
||||
|
||||
def train_mode(self) -> None:
|
||||
# PPO
|
||||
self.alg.policy.train()
|
||||
# RND
|
||||
if self.alg_cfg["rnd_cfg"]:
|
||||
self.alg.rnd.train()
|
||||
|
||||
def eval_mode(self) -> None:
|
||||
# PPO
|
||||
self.alg.policy.eval()
|
||||
# RND
|
||||
if self.alg_cfg["rnd_cfg"]:
|
||||
self.alg.rnd.eval()
|
||||
|
||||
def add_git_repo_to_log(self, repo_file_path: str) -> None:
|
||||
self.logger.git_status_repos.append(repo_file_path)
|
||||
|
||||
def _get_default_obs_sets(self) -> list[str]:
|
||||
"""Get the the default observation sets required for the algorithm.
|
||||
|
||||
.. note::
|
||||
See :func:`resolve_obs_groups` for more details on the handling of observation sets.
|
||||
"""
|
||||
default_sets = ["critic"]
|
||||
if "rnd_cfg" in self.alg_cfg and self.alg_cfg["rnd_cfg"] is not None:
|
||||
default_sets.append("rnd_state")
|
||||
return default_sets
|
||||
|
||||
def _configure_multi_gpu(self) -> None:
|
||||
"""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: TensorDict) -> PPO:
|
||||
"""Construct the actor-critic algorithm."""
|
||||
# Resolve RND config if used
|
||||
self.alg_cfg = resolve_rnd_config(self.alg_cfg, obs, self.cfg["obs_groups"], self.env)
|
||||
|
||||
# Resolve symmetry config if used
|
||||
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 policy
|
||||
actor_critic_class = resolve_callable(self.policy_cfg.pop("class_name"))
|
||||
actor_critic: ActorCritic | ActorCriticRecurrent | ActorCriticCNN = actor_critic_class(
|
||||
obs, self.cfg["obs_groups"], self.env.num_actions, **self.policy_cfg
|
||||
).to(self.device)
|
||||
|
||||
# Initialize the storage
|
||||
storage = RolloutStorage(
|
||||
"rl", self.env.num_envs, self.cfg["num_steps_per_env"], obs, [self.env.num_actions], self.device
|
||||
)
|
||||
|
||||
# Initialize the algorithm
|
||||
alg_class = resolve_callable(self.alg_cfg.pop("class_name"))
|
||||
alg: PPO = alg_class(
|
||||
actor_critic, storage, device=self.device, **self.alg_cfg, multi_gpu_cfg=self.multi_gpu_cfg
|
||||
)
|
||||
|
||||
return alg
|
||||
290
source/rsl_rl/rsl_rl/runners/on_policy_runner_cts.py
Normal file
290
source/rsl_rl/rsl_rl/runners/on_policy_runner_cts.py
Normal file
@@ -0,0 +1,290 @@
|
||||
# Copyright (c) 2021-2026, ETH Zurich and NVIDIA CORPORATION
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import time
|
||||
import torch
|
||||
import warnings
|
||||
from tensordict import TensorDict
|
||||
|
||||
from rsl_rl.algorithms import MoECTS
|
||||
from rsl_rl.env import VecEnv
|
||||
from rsl_rl.modules import (
|
||||
ActorCriticMoECTS,
|
||||
resolve_rnd_config,
|
||||
resolve_symmetry_config,
|
||||
)
|
||||
from rsl_rl.storage import RolloutStorageCTS
|
||||
from rsl_rl.utils import resolve_callable, resolve_obs_groups
|
||||
from rsl_rl.utils.logger_cts import LoggerCTS
|
||||
|
||||
|
||||
class OnPolicyRunnerCTS:
|
||||
"""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: str = "cpu") -> None:
|
||||
self.cfg = train_cfg
|
||||
self.policy_cfg = train_cfg["policy"]
|
||||
self.alg_cfg = train_cfg["algorithm"]
|
||||
self.device = device
|
||||
self.env = env
|
||||
|
||||
# Setup multi-GPU training if enabled
|
||||
self._configure_multi_gpu()
|
||||
|
||||
# Query observations from environment for algorithm construction
|
||||
obs = self.env.get_observations()
|
||||
self.cfg["obs_groups"] = resolve_obs_groups(obs, self.cfg["obs_groups"], self._get_default_obs_sets())
|
||||
|
||||
# Create the algorithm
|
||||
self.alg = self._construct_algorithm(obs)
|
||||
|
||||
# Create the logger
|
||||
self.logger = LoggerCTS(
|
||||
log_dir=log_dir,
|
||||
cfg=self.cfg,
|
||||
env_cfg=self.env.cfg,
|
||||
num_envs=self.env.num_envs,
|
||||
is_distributed=self.is_distributed,
|
||||
gpu_world_size=self.gpu_world_size,
|
||||
gpu_global_rank=self.gpu_global_rank,
|
||||
teacher_env_idxs=self.alg.teacher_env_idxs,
|
||||
device=self.device,
|
||||
)
|
||||
|
||||
self.current_learning_iteration = 0
|
||||
|
||||
def learn(self, num_learning_iterations: int, init_at_random_ep_len: bool = False) -> None:
|
||||
# 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)
|
||||
|
||||
# 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_it = self.current_learning_iteration
|
||||
total_it = start_it + num_learning_iterations
|
||||
for it in range(start_it, total_it):
|
||||
start = time.time()
|
||||
# Rollout
|
||||
with torch.inference_mode():
|
||||
for _ in range(self.cfg["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_cfg["rnd_cfg"] else None
|
||||
# Book keeping
|
||||
self.logger.process_env_step(rewards, dones, extras, intrinsic_rewards)
|
||||
|
||||
stop = time.time()
|
||||
collect_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 information
|
||||
self.logger.log(
|
||||
it=it,
|
||||
start_it=start_it,
|
||||
total_it=total_it,
|
||||
collect_time=collect_time,
|
||||
learn_time=learn_time,
|
||||
loss_dict=loss_dict,
|
||||
learning_rate=self.alg.learning_rate,
|
||||
action_std=self.alg.policy.action_std,
|
||||
rnd_weight=self.alg.rnd.weight if self.alg_cfg["rnd_cfg"] else None,
|
||||
)
|
||||
|
||||
# Save model
|
||||
if it % self.cfg["save_interval"] == 0:
|
||||
self.save(os.path.join(self.logger.log_dir, f"model_{it}.pt")) # type: ignore
|
||||
|
||||
# Save the final model after training
|
||||
if self.logger.log_dir is not None and not self.logger.disable_logs:
|
||||
self.save(os.path.join(self.logger.log_dir, f"model_{self.current_learning_iteration}.pt"))
|
||||
|
||||
def save(self, path: str, infos: dict | None = None) -> None:
|
||||
# Save model
|
||||
saved_dict = {
|
||||
"model_state_dict": self.alg.policy.state_dict(),
|
||||
"optimizer_state_dict": self.alg.optimizer.state_dict(),
|
||||
"optimizer_stu_enc_state_dict": self.alg.optimizer_stu_enc.state_dict(),
|
||||
"iter": self.current_learning_iteration,
|
||||
"infos": infos,
|
||||
}
|
||||
# Save RND model if used
|
||||
if self.alg_cfg["rnd_cfg"]:
|
||||
saved_dict["rnd_state_dict"] = self.alg.rnd.state_dict()
|
||||
if self.alg.rnd_optimizer:
|
||||
saved_dict["rnd_optimizer_state_dict"] = self.alg.rnd_optimizer.state_dict()
|
||||
torch.save(saved_dict, path)
|
||||
|
||||
# Upload model to external logging services
|
||||
self.logger.save_model(path, self.current_learning_iteration)
|
||||
|
||||
def load(self, path: str, load_optimizer: bool = True, map_location: str | None = None) -> dict:
|
||||
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 self.alg_cfg["rnd_cfg"]:
|
||||
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"])
|
||||
# Student encoder optimizer
|
||||
self.alg.optimizer_stu_enc.load_state_dict(loaded_dict["optimizer_stu_enc_state_dict"])
|
||||
# RND optimizer if used
|
||||
if self.alg_cfg["rnd_cfg"]:
|
||||
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: str | None = None) -> callable:
|
||||
self.eval_mode() # Switch to evaluation mode (e.g. for dropout)
|
||||
if device is not None:
|
||||
self.alg.policy.to(device)
|
||||
return self.alg.policy.act_inference
|
||||
|
||||
def train_mode(self) -> None:
|
||||
# PPO
|
||||
self.alg.policy.train()
|
||||
# RND
|
||||
if self.alg_cfg["rnd_cfg"]:
|
||||
self.alg.rnd.train()
|
||||
|
||||
def eval_mode(self) -> None:
|
||||
# PPO
|
||||
self.alg.policy.eval()
|
||||
# RND
|
||||
if self.alg_cfg["rnd_cfg"]:
|
||||
self.alg.rnd.eval()
|
||||
|
||||
def add_git_repo_to_log(self, repo_file_path: str) -> None:
|
||||
self.logger.git_status_repos.append(repo_file_path)
|
||||
|
||||
def _get_default_obs_sets(self) -> list[str]:
|
||||
"""Get the the default observation sets required for the algorithm.
|
||||
|
||||
.. note::
|
||||
See :func:`resolve_obs_groups` for more details on the handling of observation sets.
|
||||
"""
|
||||
default_sets = ["critic"]
|
||||
if "rnd_cfg" in self.alg_cfg and self.alg_cfg["rnd_cfg"] is not None:
|
||||
default_sets.append("rnd_state")
|
||||
return default_sets
|
||||
|
||||
def _configure_multi_gpu(self) -> None:
|
||||
"""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: TensorDict) -> MoECTS:
|
||||
"""Construct the actor-critic algorithm."""
|
||||
# Resolve RND config if used
|
||||
self.alg_cfg = resolve_rnd_config(self.alg_cfg, obs, self.cfg["obs_groups"], self.env)
|
||||
|
||||
# Resolve symmetry config if used
|
||||
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 policy
|
||||
# actor_critic_class = resolve_callable(self.policy_cfg.pop("class_name"))
|
||||
actor_critic_class = eval(self.policy_cfg.pop("class_name")) # temporally use eval to avoid import bugs
|
||||
actor_critic: ActorCriticMoECTS = actor_critic_class(
|
||||
obs, self.cfg["obs_groups"], self.env.num_actions, **self.policy_cfg
|
||||
).to(self.device)
|
||||
|
||||
# Initialize the storage
|
||||
storage = RolloutStorageCTS(
|
||||
"rl", self.env.num_envs, max(int(self.env.num_envs*self.alg_cfg["teacher_env_ratio"]), 1), self.cfg["num_steps_per_env"], obs, [self.env.num_actions], self.device
|
||||
)
|
||||
|
||||
# Initialize the algorithm
|
||||
# alg_class = resolve_callable(self.alg_cfg.pop("class_name"))
|
||||
alg_class = eval(self.alg_cfg.pop("class_name")) # temporally use eval to avoid import bugs
|
||||
alg: MoECTS = alg_class(
|
||||
actor_critic, storage, self.env.num_envs, device=self.device, **self.alg_cfg, multi_gpu_cfg=self.multi_gpu_cfg
|
||||
)
|
||||
|
||||
return alg
|
||||
11
source/rsl_rl/rsl_rl/storage/__init__.py
Normal file
11
source/rsl_rl/rsl_rl/storage/__init__.py
Normal file
@@ -0,0 +1,11 @@
|
||||
# Copyright (c) 2021-2026, 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 .rollout_storage_cts import RolloutStorageCTS
|
||||
|
||||
__all__ = ["RolloutStorage", "RolloutStorageCTS"]
|
||||
276
source/rsl_rl/rsl_rl/storage/rollout_storage.py
Normal file
276
source/rsl_rl/rsl_rl/storage/rollout_storage.py
Normal file
@@ -0,0 +1,276 @@
|
||||
# Copyright (c) 2021-2026, ETH Zurich and NVIDIA CORPORATION
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
from collections.abc import Generator
|
||||
from tensordict import TensorDict
|
||||
|
||||
from rsl_rl.networks import HiddenState
|
||||
from rsl_rl.utils import split_and_pad_trajectories
|
||||
|
||||
|
||||
class RolloutStorage:
|
||||
"""Storage for the data collected during a rollout.
|
||||
|
||||
The rollout storage is populated by adding transitions during the rollout phase. It then returns a generator for
|
||||
learning, depending on the algorithm and the policy architecture.
|
||||
"""
|
||||
|
||||
class Transition:
|
||||
"""Storage for a single state transition."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.observations: TensorDict | None = None
|
||||
self.actions: torch.Tensor | None = None
|
||||
self.privileged_actions: torch.Tensor | None = None
|
||||
self.rewards: torch.Tensor | None = None
|
||||
self.dones: torch.Tensor | None = None
|
||||
self.values: torch.Tensor | None = None
|
||||
self.actions_log_prob: torch.Tensor
|
||||
self.action_mean: torch.Tensor | None = None
|
||||
self.action_sigma: torch.Tensor | None = None
|
||||
self.hidden_states: tuple[HiddenState, HiddenState] = (None, None)
|
||||
|
||||
def clear(self) -> None:
|
||||
self.__init__()
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
training_type: str,
|
||||
num_envs: int,
|
||||
num_transitions_per_env: int,
|
||||
obs: TensorDict,
|
||||
actions_shape: tuple[int] | list[int],
|
||||
device: str = "cpu",
|
||||
) -> None:
|
||||
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_state_a = None
|
||||
self.saved_hidden_state_c = None
|
||||
|
||||
# Counter for the number of transitions stored
|
||||
self.step = 0
|
||||
|
||||
def add_transition(self, transition: Transition) -> None:
|
||||
# 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 clear(self) -> None:
|
||||
self.step = 0
|
||||
|
||||
# For distillation
|
||||
def generator(self) -> Generator:
|
||||
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: int, num_epochs: int = 8) -> Generator:
|
||||
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
|
||||
stop = (i + 1) * mini_batch_size
|
||||
batch_idx = indices[start:stop]
|
||||
|
||||
# Create the mini-batch
|
||||
obs_batch = observations[batch_idx]
|
||||
actions_batch = actions[batch_idx]
|
||||
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]
|
||||
|
||||
hidden_state_a_batch = None
|
||||
hidden_state_c_batch = None
|
||||
masks_batch = None
|
||||
|
||||
# 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,
|
||||
(
|
||||
hidden_state_a_batch,
|
||||
hidden_state_c_batch,
|
||||
),
|
||||
masks_batch,
|
||||
)
|
||||
|
||||
# For reinforcement learning with recurrent networks
|
||||
def recurrent_mini_batch_generator(self, num_mini_batches: int, num_epochs: int = 8) -> Generator:
|
||||
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])
|
||||
last_was_done = last_was_done.permute(1, 0)
|
||||
# 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]
|
||||
hidden_state_a_batch = [
|
||||
saved_hidden_state.permute(2, 0, 1, 3)[last_was_done][first_traj:last_traj]
|
||||
.transpose(1, 0)
|
||||
.contiguous()
|
||||
for saved_hidden_state in self.saved_hidden_state_a
|
||||
]
|
||||
hidden_state_c_batch = [
|
||||
saved_hidden_state.permute(2, 0, 1, 3)[last_was_done][first_traj:last_traj]
|
||||
.transpose(1, 0)
|
||||
.contiguous()
|
||||
for saved_hidden_state in self.saved_hidden_state_c
|
||||
]
|
||||
# Remove the tuple for GRU
|
||||
hidden_state_a_batch = (
|
||||
hidden_state_a_batch[0] if len(hidden_state_a_batch) == 1 else hidden_state_a_batch
|
||||
)
|
||||
hidden_state_c_batch = (
|
||||
hidden_state_c_batch[0] if len(hidden_state_c_batch) == 1 else hidden_state_c_batch
|
||||
)
|
||||
|
||||
# Yield the mini-batch
|
||||
yield (
|
||||
obs_batch,
|
||||
actions_batch,
|
||||
values_batch,
|
||||
advantages_batch,
|
||||
returns_batch,
|
||||
old_actions_log_prob_batch,
|
||||
old_mu_batch,
|
||||
old_sigma_batch,
|
||||
(
|
||||
hidden_state_a_batch,
|
||||
hidden_state_c_batch,
|
||||
),
|
||||
masks_batch,
|
||||
)
|
||||
|
||||
first_traj = last_traj
|
||||
|
||||
def _save_hidden_states(self, hidden_states: tuple[HiddenState, HiddenState]) -> None:
|
||||
if hidden_states == (None, None):
|
||||
return
|
||||
# Make a tuple out of GRU hidden states to match the LSTM format
|
||||
hidden_state_a = hidden_states[0] if isinstance(hidden_states[0], tuple) else (hidden_states[0],)
|
||||
hidden_state_c = hidden_states[1] if isinstance(hidden_states[1], tuple) else (hidden_states[1],)
|
||||
# Initialize hidden states if needed
|
||||
if self.saved_hidden_state_a is None:
|
||||
self.saved_hidden_state_a = [
|
||||
torch.zeros(self.observations.shape[0], *hidden_state_a[i].shape, device=self.device)
|
||||
for i in range(len(hidden_state_a))
|
||||
]
|
||||
self.saved_hidden_state_c = [
|
||||
torch.zeros(self.observations.shape[0], *hidden_state_c[i].shape, device=self.device)
|
||||
for i in range(len(hidden_state_c))
|
||||
]
|
||||
# Copy the states
|
||||
for i in range(len(hidden_state_a)):
|
||||
self.saved_hidden_state_a[i][self.step].copy_(hidden_state_a[i])
|
||||
self.saved_hidden_state_c[i][self.step].copy_(hidden_state_c[i])
|
||||
207
source/rsl_rl/rsl_rl/storage/rollout_storage_cts.py
Normal file
207
source/rsl_rl/rsl_rl/storage/rollout_storage_cts.py
Normal file
@@ -0,0 +1,207 @@
|
||||
# Copyright (c) 2021-2026, ETH Zurich and NVIDIA CORPORATION
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
from collections.abc import Generator
|
||||
from tensordict import TensorDict
|
||||
|
||||
from rsl_rl.networks import HiddenState
|
||||
from rsl_rl.utils import split_and_pad_trajectories
|
||||
from functools import partial
|
||||
|
||||
class RolloutStorageCTS:
|
||||
"""Storage for the data collected during a rollout.
|
||||
|
||||
The rollout storage is populated by adding transitions during the rollout phase. It then returns a generator for
|
||||
learning, depending on the algorithm and the policy architecture.
|
||||
"""
|
||||
|
||||
class Transition:
|
||||
"""Storage for a single state transition."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.observations: TensorDict | None = None
|
||||
self.actions: torch.Tensor | None = None
|
||||
self.privileged_actions: torch.Tensor | None = None
|
||||
self.rewards: torch.Tensor | None = None
|
||||
self.dones: torch.Tensor | None = None
|
||||
self.values: torch.Tensor | None = None
|
||||
self.actions_log_prob: torch.Tensor
|
||||
self.action_mean: torch.Tensor | None = None
|
||||
self.action_sigma: torch.Tensor | None = None
|
||||
self.hidden_states: tuple[HiddenState, HiddenState] = (None, None)
|
||||
|
||||
def clear(self) -> None:
|
||||
self.__init__()
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
training_type: str,
|
||||
num_envs: int,
|
||||
teacher_num_envs: int,
|
||||
num_transitions_per_env: int,
|
||||
obs: TensorDict,
|
||||
actions_shape: tuple[int] | list[int],
|
||||
device: str = "cpu",
|
||||
) -> None:
|
||||
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
|
||||
self.teacher_num_envs = teacher_num_envs
|
||||
self.student_num_envs = num_envs - teacher_num_envs
|
||||
|
||||
# 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_state_a = None
|
||||
self.saved_hidden_state_c = None
|
||||
|
||||
# Counter for the number of transitions stored
|
||||
self.step = 0
|
||||
|
||||
def add_transition(self, transition: Transition) -> None:
|
||||
# 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 clear(self) -> None:
|
||||
self.step = 0
|
||||
|
||||
# For distillation
|
||||
def generator(self) -> Generator:
|
||||
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: int, num_epochs: int = 8) -> Generator:
|
||||
if self.training_type != "rl":
|
||||
raise ValueError("This function is only available for reinforcement learning training.")
|
||||
|
||||
# Prepare indices
|
||||
teacher_samples_num = self.teacher_num_envs * self.num_transitions_per_env
|
||||
student_samples_num = self.student_num_envs * self.num_transitions_per_env
|
||||
teacher_mini_batch_size = teacher_samples_num // num_mini_batches
|
||||
student_mini_batch_size = student_samples_num // num_mini_batches
|
||||
teacher_indices = torch.randperm(teacher_samples_num, requires_grad=False, device=self.device)
|
||||
student_indices = teacher_samples_num + torch.randperm(student_samples_num, requires_grad=False, device=self.device)
|
||||
|
||||
# Core
|
||||
observations = self.observations.transpose(0, 1).flatten(0, 1)
|
||||
actions = self.actions.transpose(0, 1).flatten(0, 1)
|
||||
values = self.values.transpose(0, 1).flatten(0, 1)
|
||||
returns = self.returns.transpose(0, 1).flatten(0, 1)
|
||||
|
||||
# For PPO
|
||||
old_actions_log_prob = self.actions_log_prob.transpose(0, 1).flatten(0, 1)
|
||||
advantages = self.advantages.transpose(0, 1).flatten(0, 1)
|
||||
old_mu = self.mu.transpose(0, 1).flatten(0, 1)
|
||||
old_sigma = self.sigma.transpose(0, 1).flatten(0, 1)
|
||||
|
||||
def _get_teacher_student_samples(data, slice):
|
||||
(i1, i2), (j1, j2) = slice
|
||||
return torch.cat([data[teacher_indices[i1:i2]], data[student_indices[j1:j2]]], 0).detach()
|
||||
|
||||
for epoch in range(num_epochs):
|
||||
for i in range(num_mini_batches):
|
||||
# Select the indices for the mini-batch
|
||||
slice = (
|
||||
(i * teacher_mini_batch_size, (i+1) * teacher_mini_batch_size),
|
||||
(i * student_mini_batch_size, (i+1) * student_mini_batch_size),
|
||||
)
|
||||
|
||||
# Create the mini-batch
|
||||
get_batch = partial(_get_teacher_student_samples, slice=slice)
|
||||
obs_batch, actions_batch, target_values_batch, returns_batch, \
|
||||
old_actions_log_prob_batch, advantages_batch, old_mu_batch, \
|
||||
old_sigma_batch = map(get_batch, [
|
||||
observations,
|
||||
actions,
|
||||
values,
|
||||
returns,
|
||||
old_actions_log_prob,
|
||||
advantages,
|
||||
old_mu,
|
||||
old_sigma
|
||||
])
|
||||
|
||||
hidden_state_a_batch = None
|
||||
hidden_state_c_batch = None
|
||||
masks_batch = None
|
||||
|
||||
# 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,
|
||||
(
|
||||
hidden_state_a_batch,
|
||||
hidden_state_c_batch,
|
||||
),
|
||||
masks_batch,
|
||||
)
|
||||
|
||||
# For reinforcement learning with recurrent networks
|
||||
def recurrent_mini_batch_generator(self, num_mini_batches: int, num_epochs: int = 8) -> Generator:
|
||||
return NotImplementedError("CTS rollout storage does not support RNNs yet.")
|
||||
|
||||
def _save_hidden_states(self, hidden_states: tuple[HiddenState, HiddenState]) -> None:
|
||||
return NotImplementedError("CTS rollout storage does not support RNNs yet.")
|
||||
26
source/rsl_rl/rsl_rl/utils/__init__.py
Normal file
26
source/rsl_rl/rsl_rl/utils/__init__.py
Normal file
@@ -0,0 +1,26 @@
|
||||
# Copyright (c) 2021-2026, ETH Zurich and NVIDIA CORPORATION
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
|
||||
"""Helper functions."""
|
||||
|
||||
from .utils import (
|
||||
get_param,
|
||||
resolve_callable,
|
||||
resolve_nn_activation,
|
||||
resolve_obs_groups,
|
||||
resolve_optimizer,
|
||||
split_and_pad_trajectories,
|
||||
unpad_trajectories,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"get_param",
|
||||
"resolve_callable",
|
||||
"resolve_nn_activation",
|
||||
"resolve_obs_groups",
|
||||
"resolve_optimizer",
|
||||
"split_and_pad_trajectories",
|
||||
"unpad_trajectories",
|
||||
]
|
||||
290
source/rsl_rl/rsl_rl/utils/logger.py
Normal file
290
source/rsl_rl/rsl_rl/utils/logger.py
Normal file
@@ -0,0 +1,290 @@
|
||||
# Copyright (c) 2021-2026, ETH Zurich and NVIDIA CORPORATION
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import git
|
||||
import os
|
||||
import pathlib
|
||||
import statistics
|
||||
import time
|
||||
import torch
|
||||
from collections import deque
|
||||
|
||||
import rsl_rl
|
||||
|
||||
|
||||
class Logger:
|
||||
"""Logger to save the learning metrics to different logging services."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
log_dir: str | None,
|
||||
cfg: dict,
|
||||
env_cfg: dict | object,
|
||||
num_envs: int,
|
||||
is_distributed: bool,
|
||||
gpu_world_size: int,
|
||||
gpu_global_rank: int,
|
||||
device: str,
|
||||
) -> None:
|
||||
self.log_dir = log_dir
|
||||
self.cfg = cfg
|
||||
self.num_envs = num_envs
|
||||
self.gpu_world_size = gpu_world_size
|
||||
self.device = device
|
||||
self.git_status_repos = [rsl_rl.__file__]
|
||||
self.tot_timesteps = 0
|
||||
self.tot_time = 0
|
||||
|
||||
# Create buffers
|
||||
self.ep_extras = []
|
||||
self.rewbuffer = deque(maxlen=100)
|
||||
self.lenbuffer = deque(maxlen=100)
|
||||
self.cur_reward_sum = torch.zeros(self.num_envs, dtype=torch.float, device=self.device)
|
||||
self.cur_episode_length = torch.zeros(self.num_envs, dtype=torch.float, device=self.device)
|
||||
|
||||
# Create RND buffers
|
||||
if self.cfg["algorithm"]["rnd_cfg"]:
|
||||
self.erewbuffer = deque(maxlen=100)
|
||||
self.irewbuffer = deque(maxlen=100)
|
||||
self.cur_ereward_sum = torch.zeros(self.num_envs, dtype=torch.float, device=self.device)
|
||||
self.cur_ireward_sum = torch.zeros(self.num_envs, dtype=torch.float, device=self.device)
|
||||
|
||||
# Decide whether to disable logging
|
||||
# Note: We only log from the process with rank 0 (main process)
|
||||
self.disable_logs = is_distributed and gpu_global_rank != 0
|
||||
|
||||
# Initialize the writer
|
||||
self._prepare_logging_writer()
|
||||
|
||||
# Log code state
|
||||
self._store_code_state()
|
||||
|
||||
# Log configuration
|
||||
if self.writer and not self.disable_logs and self.logger_type in ["wandb", "neptune"]:
|
||||
self.writer.store_config(env_cfg, self.cfg)
|
||||
|
||||
def process_env_step(
|
||||
self,
|
||||
rewards: torch.Tensor,
|
||||
dones: torch.Tensor,
|
||||
extras: dict,
|
||||
intrinsic_rewards: torch.Tensor | None = None,
|
||||
) -> None:
|
||||
"""Add metrics from the environment step to the buffers."""
|
||||
if self.log_dir is not None:
|
||||
if "episode" in extras:
|
||||
self.ep_extras.append(extras["episode"])
|
||||
elif "log" in extras:
|
||||
self.ep_extras.append(extras["log"])
|
||||
|
||||
# Update rewards and episode length
|
||||
if intrinsic_rewards is not None:
|
||||
self.cur_ereward_sum += rewards
|
||||
self.cur_ireward_sum += intrinsic_rewards
|
||||
self.cur_reward_sum += rewards + intrinsic_rewards
|
||||
else:
|
||||
self.cur_reward_sum += rewards
|
||||
self.cur_episode_length += 1
|
||||
|
||||
# Clear data for completed episodes
|
||||
new_ids = (dones > 0).nonzero(as_tuple=False)
|
||||
self.rewbuffer.extend(self.cur_reward_sum[new_ids][:, 0].cpu().numpy().tolist())
|
||||
self.lenbuffer.extend(self.cur_episode_length[new_ids][:, 0].cpu().numpy().tolist())
|
||||
self.cur_reward_sum[new_ids] = 0
|
||||
self.cur_episode_length[new_ids] = 0
|
||||
if intrinsic_rewards is not None:
|
||||
self.erewbuffer.extend(self.cur_ereward_sum[new_ids][:, 0].cpu().numpy().tolist())
|
||||
self.irewbuffer.extend(self.cur_ireward_sum[new_ids][:, 0].cpu().numpy().tolist())
|
||||
self.cur_ereward_sum[new_ids] = 0
|
||||
self.cur_ireward_sum[new_ids] = 0
|
||||
|
||||
def log(
|
||||
self,
|
||||
it: int,
|
||||
start_it: int,
|
||||
total_it: int,
|
||||
collect_time: float,
|
||||
learn_time: float,
|
||||
loss_dict: dict,
|
||||
learning_rate: float,
|
||||
action_std: torch.Tensor,
|
||||
rnd_weight: float | None,
|
||||
print_minimal: bool = False,
|
||||
width: int = 80,
|
||||
pad: int = 40,
|
||||
) -> None:
|
||||
"""Log the training metrics to the logging service and print them to the console."""
|
||||
if self.log_dir is not None and not self.disable_logs:
|
||||
collection_size = self.cfg["num_steps_per_env"] * self.num_envs * self.gpu_world_size
|
||||
iteration_time = collect_time + learn_time
|
||||
self.tot_timesteps += collection_size
|
||||
self.tot_time += iteration_time
|
||||
|
||||
# Log episode extras
|
||||
extras_string = ""
|
||||
if self.ep_extras:
|
||||
# Iterate over all keys in the episode info dictionary
|
||||
for key in self.ep_extras[0]:
|
||||
infotensor = torch.tensor([], device=self.device)
|
||||
# Iterate over all steps
|
||||
for ep_info in self.ep_extras:
|
||||
# Handle missing, scalar, and zero dimensional tensors
|
||||
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)
|
||||
if "/" in key:
|
||||
self.writer.add_scalar(key, value, it)
|
||||
extras_string += f"""{f"{key}:":>{pad}} {value:.4f}\n"""
|
||||
else:
|
||||
self.writer.add_scalar("Episode/" + key, value, it)
|
||||
extras_string += f"""{f"Mean episode {key}:":>{pad}} {value:.4f}\n"""
|
||||
|
||||
# Log losses
|
||||
for key, value in loss_dict.items():
|
||||
self.writer.add_scalar(f"Loss/{key}", value, it)
|
||||
self.writer.add_scalar("Loss/learning_rate", learning_rate, it)
|
||||
|
||||
# Log noise std
|
||||
self.writer.add_scalar("Policy/mean_noise_std", action_std.mean().item(), it)
|
||||
|
||||
# Log performance
|
||||
fps = int(collection_size / (collect_time + learn_time))
|
||||
self.writer.add_scalar("Perf/total_fps", fps, it)
|
||||
self.writer.add_scalar("Perf/collection_time", collect_time, it)
|
||||
self.writer.add_scalar("Perf/learning_time", learn_time, it)
|
||||
|
||||
# Log rewards and episode length
|
||||
if len(self.rewbuffer) > 0:
|
||||
if self.cfg["algorithm"]["rnd_cfg"]:
|
||||
self.writer.add_scalar("Rnd/mean_extrinsic_reward", statistics.mean(self.erewbuffer), it)
|
||||
self.writer.add_scalar("Rnd/mean_intrinsic_reward", statistics.mean(self.irewbuffer), it)
|
||||
self.writer.add_scalar("Rnd/weight", rnd_weight, it)
|
||||
self.writer.add_scalar("Train/mean_reward", statistics.mean(self.rewbuffer), it)
|
||||
self.writer.add_scalar("Train/mean_episode_length", statistics.mean(self.lenbuffer), it)
|
||||
if self.logger_type != "wandb":
|
||||
self.writer.add_scalar(
|
||||
"Train/mean_reward/time", statistics.mean(self.rewbuffer), int(self.tot_time)
|
||||
)
|
||||
self.writer.add_scalar(
|
||||
"Train/mean_episode_length/time", statistics.mean(self.lenbuffer), int(self.tot_time)
|
||||
)
|
||||
|
||||
# Print to console
|
||||
log_string = f"""{"#" * width}\n"""
|
||||
log_string += f"""\033[1m{f" Learning iteration {it}/{total_it} ".center(width)}\033[0m \n\n"""
|
||||
|
||||
# Print run name if provided
|
||||
run_name = self.cfg.get("run_name")
|
||||
log_string += f"""{"Run name:":>{pad}} {run_name}\n""" if run_name else ""
|
||||
|
||||
# Print performance
|
||||
log_string += (
|
||||
f"""{"Total steps:":>{pad}} {self.tot_timesteps} \n"""
|
||||
f"""{"Steps per second:":>{pad}} {fps:.0f} \n"""
|
||||
f"""{"Collection time:":>{pad}} {collect_time:.3f}s \n"""
|
||||
f"""{"Learning time:":>{pad}} {learn_time:.3f}s \n"""
|
||||
)
|
||||
|
||||
# Print losses
|
||||
for key, value in loss_dict.items():
|
||||
log_string += f"""{f"Mean {key} loss:":>{pad}} {value:.4f}\n"""
|
||||
|
||||
# Print rewards and episode length
|
||||
if len(self.rewbuffer) > 0:
|
||||
if self.cfg["algorithm"]["rnd_cfg"]:
|
||||
log_string += f"""{"Mean extrinsic reward:":>{pad}} {statistics.mean(self.erewbuffer):.2f}\n"""
|
||||
log_string += f"""{"Mean intrinsic reward:":>{pad}} {statistics.mean(self.irewbuffer):.2f}\n"""
|
||||
log_string += f"""{"Mean reward:":>{pad}} {statistics.mean(self.rewbuffer):.2f}\n"""
|
||||
log_string += f"""{"Mean episode length:":>{pad}} {statistics.mean(self.lenbuffer):.2f}\n"""
|
||||
|
||||
# Print noise std
|
||||
log_string += f"""{"Mean action noise std:":>{pad}} {action_std.mean().item():.2f}\n"""
|
||||
|
||||
# Print episode extras
|
||||
if not print_minimal:
|
||||
log_string += extras_string
|
||||
|
||||
# Print footer
|
||||
done_it = it + 1 - start_it
|
||||
remaining_it = total_it - start_it - done_it
|
||||
eta = self.tot_time / done_it * remaining_it
|
||||
log_string += (
|
||||
f"""{"-" * width}\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(eta))}\n"""
|
||||
)
|
||||
print(log_string)
|
||||
|
||||
# Clear extras buffer
|
||||
self.ep_extras.clear()
|
||||
|
||||
def save_model(self, path: str, it: int) -> None:
|
||||
"""Save the model to external logging services if specified."""
|
||||
if self.writer and not self.disable_logs and self.logger_type in ["neptune", "wandb"]:
|
||||
self.writer.save_model(path, it)
|
||||
|
||||
def _prepare_logging_writer(self) -> None:
|
||||
"""Prepare the logging writer, which can be either Tensorboard, W&B or Neptune."""
|
||||
if self.log_dir is not None and not self.disable_logs:
|
||||
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)
|
||||
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)
|
||||
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 'wandb', 'neptune', or 'tensorboard'.")
|
||||
else:
|
||||
self.writer = None
|
||||
|
||||
def _store_code_state(self) -> None:
|
||||
"""Store the current git diff of the code repositories involved in the experiment."""
|
||||
if self.log_dir is not None and not self.disable_logs:
|
||||
git_log_dir = os.path.join(self.log_dir, "git")
|
||||
os.makedirs(git_log_dir, exist_ok=True)
|
||||
file_paths = []
|
||||
# Iterate over all repositories to log
|
||||
for repository_file_path in self.git_status_repos:
|
||||
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.")
|
||||
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)
|
||||
|
||||
# Upload diff files to external logging services
|
||||
if self.writer and self.logger_type in ["wandb", "neptune"] and file_paths:
|
||||
for path in file_paths:
|
||||
self.writer.save_file(path)
|
||||
311
source/rsl_rl/rsl_rl/utils/logger_cts.py
Normal file
311
source/rsl_rl/rsl_rl/utils/logger_cts.py
Normal file
@@ -0,0 +1,311 @@
|
||||
# Copyright (c) 2021-2026, ETH Zurich and NVIDIA CORPORATION
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import git
|
||||
import os
|
||||
import pathlib
|
||||
import statistics
|
||||
import time
|
||||
import torch
|
||||
from collections import deque
|
||||
|
||||
import rsl_rl
|
||||
|
||||
|
||||
class LoggerCTS:
|
||||
"""Logger to save the learning metrics to different logging services."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
log_dir: str | None,
|
||||
cfg: dict,
|
||||
env_cfg: dict | object,
|
||||
num_envs: int,
|
||||
is_distributed: bool,
|
||||
gpu_world_size: int,
|
||||
gpu_global_rank: int,
|
||||
teacher_env_idxs: torch.Tensor,
|
||||
device: str,
|
||||
) -> None:
|
||||
self.log_dir = log_dir
|
||||
self.cfg = cfg
|
||||
self.num_envs = num_envs
|
||||
self.gpu_world_size = gpu_world_size
|
||||
self.device = device
|
||||
self.git_status_repos = [rsl_rl.__file__]
|
||||
self.tot_timesteps = 0
|
||||
self.tot_time = 0
|
||||
self.teacher_env_idxs = teacher_env_idxs # Indices of teacher envs in the total envs
|
||||
|
||||
# Create buffers
|
||||
self.ep_extras = []
|
||||
self.teacher_rewbuffer = deque(maxlen=100)
|
||||
self.teacher_lenbuffer = deque(maxlen=100)
|
||||
self.student_rewbuffer = deque(maxlen=100)
|
||||
self.student_lenbuffer = deque(maxlen=100)
|
||||
self.cur_reward_sum = torch.zeros(self.num_envs, dtype=torch.float, device=self.device)
|
||||
self.cur_episode_length = torch.zeros(self.num_envs, dtype=torch.float, device=self.device)
|
||||
|
||||
# Create RND buffers
|
||||
if self.cfg["algorithm"]["rnd_cfg"]:
|
||||
self.erewbuffer = deque(maxlen=100)
|
||||
self.irewbuffer = deque(maxlen=100)
|
||||
self.cur_ereward_sum = torch.zeros(self.num_envs, dtype=torch.float, device=self.device)
|
||||
self.cur_ireward_sum = torch.zeros(self.num_envs, dtype=torch.float, device=self.device)
|
||||
|
||||
# Decide whether to disable logging
|
||||
# Note: We only log from the process with rank 0 (main process)
|
||||
self.disable_logs = is_distributed and gpu_global_rank != 0
|
||||
|
||||
# Initialize the writer
|
||||
self._prepare_logging_writer()
|
||||
|
||||
# Log code state
|
||||
self._store_code_state()
|
||||
|
||||
# Log configuration
|
||||
if self.writer and not self.disable_logs and self.logger_type in ["wandb", "neptune"]:
|
||||
self.writer.store_config(env_cfg, self.cfg)
|
||||
|
||||
def process_env_step(
|
||||
self,
|
||||
rewards: torch.Tensor,
|
||||
dones: torch.Tensor,
|
||||
extras: dict,
|
||||
intrinsic_rewards: torch.Tensor | None = None,
|
||||
) -> None:
|
||||
"""Add metrics from the environment step to the buffers."""
|
||||
if self.log_dir is not None:
|
||||
if "episode" in extras:
|
||||
self.ep_extras.append(extras["episode"])
|
||||
elif "log" in extras:
|
||||
self.ep_extras.append(extras["log"])
|
||||
|
||||
# Update rewards and episode length
|
||||
if intrinsic_rewards is not None:
|
||||
self.cur_ereward_sum += rewards
|
||||
self.cur_ireward_sum += intrinsic_rewards
|
||||
self.cur_reward_sum += rewards + intrinsic_rewards
|
||||
else:
|
||||
self.cur_reward_sum += rewards
|
||||
self.cur_episode_length += 1
|
||||
|
||||
# Clear data for completed episodes
|
||||
new_ids = (dones > 0).nonzero(as_tuple=False)
|
||||
if new_ids.shape[0]:
|
||||
teacher_ids = new_ids[torch.isin(new_ids, self.teacher_env_idxs)]
|
||||
student_ids = new_ids[~torch.isin(new_ids, self.teacher_env_idxs)]
|
||||
self.teacher_rewbuffer.extend(self.cur_reward_sum[teacher_ids].cpu().numpy().tolist())
|
||||
self.teacher_lenbuffer.extend(self.cur_episode_length[teacher_ids].cpu().numpy().tolist())
|
||||
self.student_rewbuffer.extend(self.cur_reward_sum[student_ids].cpu().numpy().tolist())
|
||||
self.student_lenbuffer.extend(self.cur_episode_length[student_ids].cpu().numpy().tolist())
|
||||
self.cur_reward_sum[new_ids] = 0
|
||||
self.cur_episode_length[new_ids] = 0
|
||||
|
||||
if intrinsic_rewards is not None:
|
||||
self.erewbuffer.extend(self.cur_ereward_sum[new_ids][:, 0].cpu().numpy().tolist())
|
||||
self.irewbuffer.extend(self.cur_ireward_sum[new_ids][:, 0].cpu().numpy().tolist())
|
||||
self.cur_ereward_sum[new_ids] = 0
|
||||
self.cur_ireward_sum[new_ids] = 0
|
||||
|
||||
def log(
|
||||
self,
|
||||
it: int,
|
||||
start_it: int,
|
||||
total_it: int,
|
||||
collect_time: float,
|
||||
learn_time: float,
|
||||
loss_dict: dict,
|
||||
learning_rate: float,
|
||||
action_std: torch.Tensor,
|
||||
rnd_weight: float | None,
|
||||
print_minimal: bool = False,
|
||||
width: int = 80,
|
||||
pad: int = 40,
|
||||
) -> None:
|
||||
"""Log the training metrics to the logging service and print them to the console."""
|
||||
if self.log_dir is not None and not self.disable_logs:
|
||||
collection_size = self.cfg["num_steps_per_env"] * self.num_envs * self.gpu_world_size
|
||||
iteration_time = collect_time + learn_time
|
||||
self.tot_timesteps += collection_size
|
||||
self.tot_time += iteration_time
|
||||
|
||||
# Log episode extras
|
||||
extras_string = ""
|
||||
if self.ep_extras:
|
||||
# Iterate over all keys in the episode info dictionary
|
||||
for key in self.ep_extras[0]:
|
||||
infotensor = torch.tensor([], device=self.device)
|
||||
# Iterate over all steps
|
||||
for ep_info in self.ep_extras:
|
||||
# Handle missing, scalar, and zero dimensional tensors
|
||||
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)
|
||||
if "/" in key:
|
||||
self.writer.add_scalar(key, value, it)
|
||||
extras_string += f"""{f"{key}:":>{pad}} {value:.4f}\n"""
|
||||
else:
|
||||
self.writer.add_scalar("Episode/" + key, value, it)
|
||||
extras_string += f"""{f"Mean episode {key}:":>{pad}} {value:.4f}\n"""
|
||||
|
||||
# Log losses
|
||||
for key, value in loss_dict.items():
|
||||
self.writer.add_scalar(f"Loss/{key}", value, it)
|
||||
self.writer.add_scalar("Loss/learning_rate", learning_rate, it)
|
||||
|
||||
# Log noise std
|
||||
self.writer.add_scalar("Policy/mean_noise_std", action_std.mean().item(), it)
|
||||
|
||||
# Log performance
|
||||
fps = int(collection_size / (collect_time + learn_time))
|
||||
self.writer.add_scalar("Perf/total_fps", fps, it)
|
||||
self.writer.add_scalar("Perf/collection_time", collect_time, it)
|
||||
self.writer.add_scalar("Perf/learning_time", learn_time, it)
|
||||
|
||||
# Log rewards and episode length
|
||||
if self.cfg["algorithm"]["rnd_cfg"] and len(self.erewbuffer) > 0:
|
||||
self.writer.add_scalar("Rnd/mean_extrinsic_reward", statistics.mean(self.erewbuffer), it)
|
||||
self.writer.add_scalar("Rnd/mean_intrinsic_reward", statistics.mean(self.irewbuffer), it)
|
||||
self.writer.add_scalar("Rnd/weight", rnd_weight, it)
|
||||
if len(self.teacher_rewbuffer) > 0:
|
||||
self.writer.add_scalar("Train/mean_teacher_reward", statistics.mean(self.teacher_rewbuffer), it)
|
||||
self.writer.add_scalar("Train/mean_teacher_episode_length", statistics.mean(self.teacher_lenbuffer), it)
|
||||
if self.logger_type != "wandb":
|
||||
self.writer.add_scalar(
|
||||
"Train/mean_teacher_reward/time", statistics.mean(self.teacher_rewbuffer), int(self.tot_time)
|
||||
)
|
||||
self.writer.add_scalar(
|
||||
"Train/mean_teacher_episode_length/time", statistics.mean(self.teacher_lenbuffer), int(self.tot_time)
|
||||
)
|
||||
if len(self.student_rewbuffer) > 0:
|
||||
self.writer.add_scalar("Train/mean_student_reward", statistics.mean(self.student_rewbuffer), it)
|
||||
self.writer.add_scalar("Train/mean_student_episode_length", statistics.mean(self.student_lenbuffer), it)
|
||||
if self.logger_type != "wandb":
|
||||
self.writer.add_scalar(
|
||||
"Train/mean_student_reward/time", statistics.mean(self.student_rewbuffer), int(self.tot_time)
|
||||
)
|
||||
self.writer.add_scalar(
|
||||
"Train/mean_student_episode_length/time", statistics.mean(self.student_lenbuffer), int(self.tot_time)
|
||||
)
|
||||
# Print to console
|
||||
log_string = f"""{"#" * width}\n"""
|
||||
log_string += f"""\033[1m{f" Learning iteration {it}/{total_it} ".center(width)}\033[0m \n\n"""
|
||||
|
||||
# Print run name if provided
|
||||
run_name = self.cfg.get("run_name")
|
||||
log_string += f"""{"Run name:":>{pad}} {run_name}\n""" if run_name else ""
|
||||
|
||||
# Print performance
|
||||
log_string += (
|
||||
f"""{"Total steps:":>{pad}} {self.tot_timesteps} \n"""
|
||||
f"""{"Steps per second:":>{pad}} {fps:.0f} \n"""
|
||||
f"""{"Collection time:":>{pad}} {collect_time:.3f}s \n"""
|
||||
f"""{"Learning time:":>{pad}} {learn_time:.3f}s \n"""
|
||||
)
|
||||
|
||||
# Print losses
|
||||
for key, value in loss_dict.items():
|
||||
log_string += f"""{f"Mean {key} loss:":>{pad}} {value:.4f}\n"""
|
||||
|
||||
# Print rewards and episode length
|
||||
if self.cfg["algorithm"]["rnd_cfg"] and len(self.erewbuffer) > 0:
|
||||
log_string += f"""{"Mean extrinsic reward:":>{pad}} {statistics.mean(self.erewbuffer):.2f}\n"""
|
||||
log_string += f"""{"Mean intrinsic reward:":>{pad}} {statistics.mean(self.irewbuffer):.2f}\n"""
|
||||
if len(self.teacher_rewbuffer) > 0:
|
||||
log_string += f"""{"Mean teacher reward:":>{pad}} {statistics.mean(self.teacher_rewbuffer):.2f}\n"""
|
||||
log_string += f"""{"Mean teacher episode length:":>{pad}} {statistics.mean(self.teacher_lenbuffer):.2f}\n"""
|
||||
if len(self.student_rewbuffer) > 0:
|
||||
log_string += f"""{"Mean student reward:":>{pad}} {statistics.mean(self.student_rewbuffer):.2f}\n"""
|
||||
log_string += f"""{"Mean student episode length:":>{pad}} {statistics.mean(self.student_lenbuffer):.2f}\n"""
|
||||
# Print noise std
|
||||
log_string += f"""{"Mean action noise std:":>{pad}} {action_std.mean().item():.2f}\n"""
|
||||
|
||||
# Print episode extras
|
||||
if not print_minimal:
|
||||
log_string += extras_string
|
||||
|
||||
# Print footer
|
||||
done_it = it + 1 - start_it
|
||||
remaining_it = total_it - start_it - done_it
|
||||
eta = self.tot_time / done_it * remaining_it
|
||||
log_string += (
|
||||
f"""{"-" * width}\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(eta))}\n"""
|
||||
)
|
||||
print(log_string)
|
||||
|
||||
# Clear extras buffer
|
||||
self.ep_extras.clear()
|
||||
|
||||
def save_model(self, path: str, it: int) -> None:
|
||||
"""Save the model to external logging services if specified."""
|
||||
if self.writer and not self.disable_logs and self.logger_type in ["neptune", "wandb"]:
|
||||
self.writer.save_model(path, it)
|
||||
|
||||
def _prepare_logging_writer(self) -> None:
|
||||
"""Prepare the logging writer, which can be either Tensorboard, W&B or Neptune."""
|
||||
if self.log_dir is not None and not self.disable_logs:
|
||||
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)
|
||||
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)
|
||||
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 'wandb', 'neptune', or 'tensorboard'.")
|
||||
else:
|
||||
self.writer = None
|
||||
|
||||
def _store_code_state(self) -> None:
|
||||
"""Store the current git diff of the code repositories involved in the experiment."""
|
||||
if self.log_dir is not None and not self.disable_logs:
|
||||
git_log_dir = os.path.join(self.log_dir, "git")
|
||||
os.makedirs(git_log_dir, exist_ok=True)
|
||||
file_paths = []
|
||||
# Iterate over all repositories to log
|
||||
for repository_file_path in self.git_status_repos:
|
||||
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.")
|
||||
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)
|
||||
|
||||
# Upload diff files to external logging services
|
||||
if self.writer and self.logger_type in ["wandb", "neptune"] and file_paths:
|
||||
for path in file_paths:
|
||||
self.writer.save_file(path)
|
||||
96
source/rsl_rl/rsl_rl/utils/neptune_utils.py
Normal file
96
source/rsl_rl/rsl_rl/utils/neptune_utils.py
Normal file
@@ -0,0 +1,96 @@
|
||||
# Copyright (c) 2021-2026, 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.") from None
|
||||
|
||||
|
||||
class NeptuneSummaryWriter(SummaryWriter):
|
||||
"""Summary writer for Neptune."""
|
||||
|
||||
def __init__(self, log_dir: str, flush_secs: int, cfg: dict) -> None:
|
||||
super().__init__(log_dir, flush_secs)
|
||||
|
||||
# Get the run name
|
||||
run_name = os.path.split(log_dir)[-1]
|
||||
|
||||
# Get neptune project and entity
|
||||
try:
|
||||
project = cfg["neptune_project"]
|
||||
except KeyError:
|
||||
raise KeyError("Please specify neptune_project in the runner config, e.g. legged_gym.") from None
|
||||
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"
|
||||
) from None
|
||||
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"
|
||||
) from None
|
||||
|
||||
# Initialize neptune
|
||||
neptune_project = entity + "/" + project
|
||||
self.run = neptune.init_run(project=neptune_project, api_token=token)
|
||||
self.run["log_dir"].log(run_name)
|
||||
|
||||
# Name mapping for incompatible characters
|
||||
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: dict | object, train_cfg: dict) -> None:
|
||||
self.run["runner_cfg"] = train_cfg
|
||||
self.run["policy_cfg"] = train_cfg["policy"]
|
||||
self.run["alg_cfg"] = train_cfg["algorithm"]
|
||||
try:
|
||||
self.run["env_cfg"] = env_cfg.to_dict()
|
||||
except Exception:
|
||||
self.run["env_cfg"] = asdict(env_cfg)
|
||||
|
||||
def add_scalar(
|
||||
self,
|
||||
tag: str,
|
||||
scalar_value: float,
|
||||
global_step: int | None = None,
|
||||
walltime: float | None = None,
|
||||
new_style: bool = False,
|
||||
) -> None:
|
||||
super().add_scalar(
|
||||
tag,
|
||||
scalar_value,
|
||||
global_step=global_step,
|
||||
walltime=walltime,
|
||||
new_style=new_style,
|
||||
)
|
||||
self.run[self._map_path(tag)].log(scalar_value, step=global_step)
|
||||
|
||||
def stop(self) -> None:
|
||||
self.run.stop()
|
||||
|
||||
def save_model(self, model_path: str, it: int) -> None:
|
||||
self.run["model/saved_model_" + str(it)].upload(model_path)
|
||||
|
||||
def save_file(self, path: str) -> None:
|
||||
name = path.rsplit("/", 1)[-1].split(".")[0]
|
||||
self.run["git_diff/" + name].upload(path)
|
||||
|
||||
def _map_path(self, path: str) -> str:
|
||||
if path in self.name_map:
|
||||
return self.name_map[path]
|
||||
else:
|
||||
return path
|
||||
340
source/rsl_rl/rsl_rl/utils/utils.py
Normal file
340
source/rsl_rl/rsl_rl/utils/utils.py
Normal file
@@ -0,0 +1,340 @@
|
||||
# Copyright (c) 2021-2026, ETH Zurich and NVIDIA CORPORATION
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib
|
||||
import pkgutil
|
||||
import torch
|
||||
import warnings
|
||||
from tensordict import TensorDict
|
||||
from typing import Any, Callable
|
||||
|
||||
import rsl_rl
|
||||
|
||||
|
||||
def get_param(param: Any, idx: int) -> Any:
|
||||
"""Get a parameter for the given index.
|
||||
|
||||
Args:
|
||||
param: Parameter or list/tuple of parameters.
|
||||
idx: Index to get the parameter for.
|
||||
"""
|
||||
if isinstance(param, (tuple, list)):
|
||||
return param[idx]
|
||||
else:
|
||||
return param
|
||||
|
||||
|
||||
def resolve_nn_activation(act_name: str) -> torch.nn.Module:
|
||||
"""Resolve the activation function from the name.
|
||||
|
||||
Args:
|
||||
act_name: 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:
|
||||
"""Resolve the optimizer from the name.
|
||||
|
||||
Args:
|
||||
optimizer_name: 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]:
|
||||
"""Split trajectories at done indices.
|
||||
|
||||
Split trajectories, concatenate them and pad with zeros up to the length of the longest trajectory. Return masks
|
||||
corresponding to valid parts of the trajectories.
|
||||
|
||||
Example (transposed for readability):
|
||||
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 the 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) # type: ignore
|
||||
# Remove the added trajectory
|
||||
padded_trajectories[k] = padded_trajectories[k][:, :-1]
|
||||
padded_trajectories = TensorDict(
|
||||
padded_trajectories, batch_size=[tensor.batch_size[0], len(trajectory_lengths_list)], device=tensor.device
|
||||
)
|
||||
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) # type: ignore
|
||||
# Remove the added trajectory
|
||||
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: torch.Tensor | TensorDict, masks: torch.Tensor) -> torch.Tensor | TensorDict:
|
||||
"""Do 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 resolve_callable(callable_or_name: type | Callable | str) -> Callable:
|
||||
"""Resolve a callable from a string, type, or return callable directly.
|
||||
|
||||
This function enables passing custom classes or functions directly or as strings. The following formats are
|
||||
supported:
|
||||
- Direct callable: Pass a type or function directly (e.g., MyClass, my_func)
|
||||
- Qualified name with colon: "module.path:Attr.Nested" (explicit, recommended)
|
||||
- Qualified name with dot: "module.path.ClassName" (implicit)
|
||||
- Simple name: e.g. "PPO", "ActorCritic", ... (looks for callable in rsl_rl)
|
||||
|
||||
Args:
|
||||
callable_or_name: A callable (type/function) or string name.
|
||||
|
||||
Returns:
|
||||
The resolved callable.
|
||||
|
||||
Raises:
|
||||
TypeError: If input is neither a callable nor a string.
|
||||
ImportError: If the module cannot be imported.
|
||||
AttributeError: If the attribute cannot be found in the module.
|
||||
ValueError: If a simple name cannot be found in rsl_rl packages.
|
||||
"""
|
||||
# Already a callable - return directly
|
||||
if callable(callable_or_name):
|
||||
return callable_or_name
|
||||
|
||||
# Must be a string at this point
|
||||
if not isinstance(callable_or_name, str):
|
||||
raise TypeError(f"Expected callable or string, got {type(callable_or_name)}")
|
||||
|
||||
# Handle qualified name with colon separator (e.g., "module.path:Attr.Nested")
|
||||
if ":" in callable_or_name:
|
||||
module_path, attr_path = callable_or_name.rsplit(":", 1)
|
||||
# Try to import the module
|
||||
module = importlib.import_module(module_path)
|
||||
# Try to get the attribute
|
||||
obj = module
|
||||
for attr in attr_path.split("."):
|
||||
obj = getattr(obj, attr)
|
||||
return obj # type: ignore
|
||||
|
||||
# Handle qualified name with dot separator (e.g., "module.path.ClassName")
|
||||
if "." in callable_or_name:
|
||||
parts = callable_or_name.split(".")
|
||||
module_found = False
|
||||
for i in range(len(parts) - 1, 0, -1):
|
||||
# Try to import the module with the first i parts
|
||||
module_path = ".".join(parts[:i])
|
||||
attr_parts = parts[i:]
|
||||
try:
|
||||
module = importlib.import_module(module_path)
|
||||
except ModuleNotFoundError:
|
||||
continue
|
||||
module_found = True
|
||||
# Once a module is found, try to get the attribute
|
||||
obj = module
|
||||
try:
|
||||
for attr in attr_parts:
|
||||
obj = getattr(obj, attr)
|
||||
return obj # type: ignore
|
||||
except AttributeError:
|
||||
continue
|
||||
if module_found:
|
||||
raise AttributeError(f"Could not resolve '{callable_or_name}': attribute not found in module")
|
||||
else:
|
||||
raise ImportError(f"Could not resolve '{callable_or_name}': no valid module.attr split found")
|
||||
|
||||
# Simple name - look for it in rsl_rl
|
||||
for _, module_name, _ in pkgutil.iter_modules(rsl_rl.__path__, "rsl_rl."):
|
||||
module = importlib.import_module(module_name)
|
||||
if hasattr(module, callable_or_name):
|
||||
return getattr(module, callable_or_name)
|
||||
|
||||
# Raise error if no approach worked
|
||||
raise ValueError(
|
||||
f"Could not resolve '{callable_or_name}'. Use qualified name like 'module.path:ClassName' "
|
||||
f"or pass the class directly."
|
||||
)
|
||||
|
||||
|
||||
def resolve_obs_groups(
|
||||
obs: TensorDict, obs_groups: dict[str, list[str]], default_sets: list[str]
|
||||
) -> dict[str, list[str]]:
|
||||
"""Validate the observation configuration and resolve 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:
|
||||
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:
|
||||
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
|
||||
74
source/rsl_rl/rsl_rl/utils/wandb_utils.py
Normal file
74
source/rsl_rl/rsl_rl/utils/wandb_utils.py
Normal file
@@ -0,0 +1,74 @@
|
||||
# Copyright (c) 2021-2026, 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 package is required to log to Weights and Biases.") from None
|
||||
|
||||
|
||||
class WandbSummaryWriter(SummaryWriter):
|
||||
"""Summary writer for Weights and Biases."""
|
||||
|
||||
def __init__(self, log_dir: str, flush_secs: int, cfg: dict) -> None:
|
||||
super().__init__(log_dir, flush_secs)
|
||||
|
||||
# Get the run name
|
||||
run_name = os.path.split(log_dir)[-1]
|
||||
|
||||
# Get wandb project and entity
|
||||
try:
|
||||
project = cfg["wandb_project"]
|
||||
except KeyError:
|
||||
raise KeyError("Please specify wandb_project in the runner config, e.g. legged_gym.") from None
|
||||
try:
|
||||
entity = os.environ["WANDB_USERNAME"]
|
||||
except KeyError:
|
||||
entity = None
|
||||
|
||||
# Initialize wandb
|
||||
wandb.init(project=project, entity=entity, name=run_name)
|
||||
wandb.config.update({"log_dir": log_dir})
|
||||
|
||||
def store_config(self, env_cfg: dict | object, train_cfg: dict) -> None:
|
||||
wandb.config.update({"runner_cfg": train_cfg})
|
||||
wandb.config.update({"policy_cfg": train_cfg["policy"]})
|
||||
wandb.config.update({"alg_cfg": train_cfg["algorithm"]})
|
||||
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: str,
|
||||
scalar_value: float,
|
||||
global_step: int | None = None,
|
||||
walltime: float | None = None,
|
||||
new_style: bool = False,
|
||||
) -> None:
|
||||
super().add_scalar(
|
||||
tag,
|
||||
scalar_value,
|
||||
global_step=global_step,
|
||||
walltime=walltime,
|
||||
new_style=new_style,
|
||||
)
|
||||
wandb.log({tag: scalar_value}, step=global_step)
|
||||
|
||||
def stop(self) -> None:
|
||||
wandb.finish()
|
||||
|
||||
def save_model(self, model_path: str, it: int) -> None:
|
||||
wandb.save(model_path, base_path=os.path.dirname(model_path))
|
||||
|
||||
def save_file(self, path: str) -> None:
|
||||
wandb.save(path, base_path=os.path.dirname(path))
|
||||
71
source/rsl_rl/ruff.toml
Normal file
71
source/rsl_rl/ruff.toml
Normal file
@@ -0,0 +1,71 @@
|
||||
line-length = 120
|
||||
target-version = "py39"
|
||||
preview = true
|
||||
|
||||
[lint]
|
||||
select = [
|
||||
# pycodestyle
|
||||
"E", "W",
|
||||
# pydocstyle
|
||||
"D",
|
||||
# pylint for later
|
||||
# "PL",
|
||||
# pyflakes
|
||||
"F",
|
||||
# pyupgrade
|
||||
"UP",
|
||||
# pep8-naming
|
||||
"N",
|
||||
# flake8-bugbear
|
||||
"B",
|
||||
# flake8-simplify
|
||||
"SIM",
|
||||
# flake8-tidy-imports
|
||||
"TID",
|
||||
# flake8-annotations
|
||||
"ANN",
|
||||
# isort
|
||||
"I",
|
||||
# perflint
|
||||
"PERF",
|
||||
# ruff
|
||||
"RUF",
|
||||
]
|
||||
ignore = ["B006",
|
||||
"B007",
|
||||
"B028",
|
||||
"ANN401",
|
||||
"D100",
|
||||
"D101",
|
||||
"D102",
|
||||
"D103",
|
||||
"D104",
|
||||
"D105",
|
||||
"D106",
|
||||
"D107",
|
||||
"D203",
|
||||
"D213",
|
||||
"D413",
|
||||
]
|
||||
per-file-ignores = {"*/__init__.py" = ["F401"]}
|
||||
|
||||
[lint.isort]
|
||||
# Order of imports
|
||||
section-order = [
|
||||
"future",
|
||||
"standard-library",
|
||||
"third-party",
|
||||
"first-party",
|
||||
"local-folder",
|
||||
]
|
||||
# Extra standard libraries considered as part of python (permissive licenses)
|
||||
extra-standard-library = [
|
||||
"numpy",
|
||||
"torch",
|
||||
"tensordict",
|
||||
"warp",
|
||||
"typing_extensions",
|
||||
"git",
|
||||
]
|
||||
# Imports from this repository
|
||||
known-first-party = ["rsl_rl"]
|
||||
8
source/rsl_rl/setup.py
Normal file
8
source/rsl_rl/setup.py
Normal file
@@ -0,0 +1,8 @@
|
||||
# Copyright (c) 2021-2026, ETH Zurich and NVIDIA CORPORATION
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
|
||||
from setuptools import setup
|
||||
|
||||
setup()
|
||||
6
source/rsl_rl/tests/__init__.py
Normal file
6
source/rsl_rl/tests/__init__.py
Normal file
@@ -0,0 +1,6 @@
|
||||
# Copyright (c) 2021-2026, ETH Zurich and NVIDIA CORPORATION
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
|
||||
"""Tests for rsl_rl."""
|
||||
6
source/rsl_rl/tests/utils/__init__.py
Normal file
6
source/rsl_rl/tests/utils/__init__.py
Normal file
@@ -0,0 +1,6 @@
|
||||
# Copyright (c) 2021-2026, ETH Zurich and NVIDIA CORPORATION
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
|
||||
"""Tests for the utils module of rsl_rl."""
|
||||
148
source/rsl_rl/tests/utils/test_resolve_callable.py
Normal file
148
source/rsl_rl/tests/utils/test_resolve_callable.py
Normal file
@@ -0,0 +1,148 @@
|
||||
# Copyright (c) 2021-2026, ETH Zurich and NVIDIA CORPORATION
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
|
||||
"""Tests for resolve_callable utility function."""
|
||||
|
||||
import pytest
|
||||
|
||||
from rsl_rl.utils import resolve_callable
|
||||
|
||||
|
||||
# Test fixtures - nested class for testing nested attribute resolution
|
||||
class OuterClass:
|
||||
"""Outer class for testing nested attribute resolution."""
|
||||
|
||||
class InnerClass:
|
||||
"""Inner nested class."""
|
||||
|
||||
pass
|
||||
|
||||
@staticmethod
|
||||
def static_method() -> str:
|
||||
return "static"
|
||||
|
||||
|
||||
def sample_function() -> str:
|
||||
"""Sample function for testing."""
|
||||
return "sample"
|
||||
|
||||
|
||||
class TestResolveCallableDirect:
|
||||
"""Tests for direct callable passing."""
|
||||
|
||||
def test_direct_class(self) -> None:
|
||||
"""Passing a class directly should return it unchanged."""
|
||||
from rsl_rl.algorithms import PPO
|
||||
|
||||
result = resolve_callable(PPO)
|
||||
assert result is PPO
|
||||
|
||||
def test_direct_function(self) -> None:
|
||||
"""Passing a function directly should return it unchanged."""
|
||||
result = resolve_callable(sample_function)
|
||||
assert result is sample_function
|
||||
|
||||
def test_direct_builtin(self) -> None:
|
||||
"""Passing a builtin should return it unchanged."""
|
||||
result = resolve_callable(len)
|
||||
assert result is len
|
||||
|
||||
|
||||
class TestResolveCallableColonFormat:
|
||||
"""Tests for colon-separated format 'module:attr'."""
|
||||
|
||||
def test_colon_format_class(self) -> None:
|
||||
"""Should resolve 'module:Class' format."""
|
||||
result = resolve_callable("rsl_rl.algorithms:PPO")
|
||||
from rsl_rl.algorithms import PPO
|
||||
|
||||
assert result is PPO
|
||||
|
||||
def test_colon_format_nested(self) -> None:
|
||||
"""Should resolve 'module:Outer.Inner' nested format."""
|
||||
result = resolve_callable("tests.utils.test_resolve_callable:OuterClass.InnerClass")
|
||||
assert result is OuterClass.InnerClass
|
||||
|
||||
def test_colon_format_static_method(self) -> None:
|
||||
"""Should resolve nested static methods."""
|
||||
result = resolve_callable("tests.utils.test_resolve_callable:OuterClass.static_method")
|
||||
assert result is OuterClass.static_method
|
||||
|
||||
def test_colon_format_invalid_module(self) -> None:
|
||||
"""Should raise ImportError for invalid module."""
|
||||
with pytest.raises(ImportError):
|
||||
resolve_callable("nonexistent_module:SomeClass")
|
||||
|
||||
def test_colon_format_invalid_attr(self) -> None:
|
||||
"""Should raise AttributeError for invalid attribute."""
|
||||
with pytest.raises(AttributeError):
|
||||
resolve_callable("rsl_rl.algorithms:NonexistentClass")
|
||||
|
||||
|
||||
class TestResolveCallableDotFormat:
|
||||
"""Tests for dot-separated format 'module.attr'."""
|
||||
|
||||
def test_dot_format_class(self) -> None:
|
||||
"""Should resolve 'module.Class' format."""
|
||||
result = resolve_callable("rsl_rl.algorithms.PPO")
|
||||
from rsl_rl.algorithms import PPO
|
||||
|
||||
assert result is PPO
|
||||
|
||||
def test_dot_format_nested(self) -> None:
|
||||
"""Should resolve 'module.Outer.Inner' nested format."""
|
||||
# This tests the progressive module path splitting
|
||||
result = resolve_callable("tests.utils.test_resolve_callable.OuterClass.InnerClass")
|
||||
assert result is OuterClass.InnerClass
|
||||
|
||||
def test_dot_format_static_method(self) -> None:
|
||||
"""Should resolve nested static methods."""
|
||||
result = resolve_callable("tests.utils.test_resolve_callable.OuterClass.static_method")
|
||||
assert result is OuterClass.static_method
|
||||
|
||||
def test_dot_format_invalid_module(self) -> None:
|
||||
"""Should raise ImportError for invalid module."""
|
||||
with pytest.raises(ImportError):
|
||||
resolve_callable("nonexistent_module.SomeClass")
|
||||
|
||||
def test_dot_format_invalid_attr(self) -> None:
|
||||
"""Should raise AttributeError for invalid attribute."""
|
||||
with pytest.raises(AttributeError):
|
||||
resolve_callable("rsl_rl.algorithms.NonexistentClass")
|
||||
|
||||
|
||||
class TestResolveCallableSimpleName:
|
||||
"""Tests for simple name resolution via rsl_rl packages."""
|
||||
|
||||
def test_simple_name(self) -> None:
|
||||
"""Should resolve 'PPO' from rsl_rl.algorithms."""
|
||||
result = resolve_callable("PPO")
|
||||
from rsl_rl.algorithms import PPO
|
||||
|
||||
assert result is PPO
|
||||
|
||||
def test_simple_name_unknown(self) -> None:
|
||||
"""Should raise ValueError for unknown simple names."""
|
||||
with pytest.raises(ValueError, match="Could not resolve"):
|
||||
resolve_callable("NonexistentClassName")
|
||||
|
||||
|
||||
class TestResolveCallableErrors:
|
||||
"""Tests for error handling."""
|
||||
|
||||
def test_type_error_none(self) -> None:
|
||||
"""Should raise TypeError for None input."""
|
||||
with pytest.raises(TypeError, match="Expected callable or string"):
|
||||
resolve_callable(None)
|
||||
|
||||
def test_type_error_int(self) -> None:
|
||||
"""Should raise TypeError for int input."""
|
||||
with pytest.raises(TypeError, match="Expected callable or string"):
|
||||
resolve_callable(42)
|
||||
|
||||
def test_type_error_list(self) -> None:
|
||||
"""Should raise TypeError for list input."""
|
||||
with pytest.raises(TypeError, match="Expected callable or string"):
|
||||
resolve_callable(["PPO"])
|
||||
Reference in New Issue
Block a user