cyy/go2_rl_robotlab
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

init commit.

Browse Source
This commit is contained in:
wertyuilife
2026-03-23 03:17:59 +08:00
commit fce9166e87
99 changed files with 13778 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
source/rsl_rl/.gitignore vendored Normal file
View 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
View 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
View 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
View 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
View 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
View 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
View 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
View 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
View 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
View 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
View 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.

19
source/rsl_rl/licenses/dependencies/pre-commit-hooks-license.txt Normal file
View 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.

19
source/rsl_rl/licenses/dependencies/pre-commit-license.txt Normal file
View 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
View 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
View 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
View 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
View 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
View 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
View 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
View 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
View 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
View 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
View 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
View 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
View 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
View 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
View 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
View 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
View 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
View 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
View 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
View 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
View 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
View 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
View 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
View 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
View 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
View 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
View 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
View 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
View 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
View 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
View 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
View 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
View 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
View 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
View 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
View 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
View 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
View 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
View 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
View 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
View 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
View 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
View 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
View 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
View 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
View 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"])