commit 1aeeb34653b506e86b2c1cfd74933841dc733011 Author: huangweifeng <719566231@qq.com> Date: Tue Oct 13 21:25:03 2026 +0800 Add standalone Huoyi training project diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 0000000..d917487 --- /dev/null +++ b/.gitattributes @@ -0,0 +1,18 @@ +# copied and modified from https://github.com/NVIDIA/warp/blob/main/.gitattributes +* text=auto +*.sh text eol=LF + +# # copied from https://github.com/isaac-sim/IsaacLab/blob/main/.gitattributes +# *.usd filter=lfs diff=lfs merge=lfs -text +# *.dae filter=lfs diff=lfs merge=lfs -text +# *.mtl filter=lfs diff=lfs merge=lfs -text +# *.obj filter=lfs diff=lfs merge=lfs -text +# *.gif filter=lfs diff=lfs merge=lfs -text +# *.png filter=lfs diff=lfs merge=lfs -text +# *.jpg filter=lfs diff=lfs merge=lfs -text +# *.psd filter=lfs diff=lfs merge=lfs -text +# *.mp4 filter=lfs diff=lfs merge=lfs -text +# *.usda filter=lfs diff=lfs merge=lfs -text +# *.hdr filter=lfs diff=lfs merge=lfs -text +# *.pt filter=lfs diff=lfs merge=lfs -text +# *.jit filter=lfs diff=lfs merge=lfs -text diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..389dd89 --- /dev/null +++ b/.gitignore @@ -0,0 +1,39 @@ +# Omniverse +**/*.dmp +**/.thumbs + +# Python +.DS_Store +**/*.egg-info/ +**/__pycache__/ +**/.pytest_cache/ +**/*.pyc +**/*.pb + +# IDE +**/.idea/ +**/.vscode/ +# Don't ignore the top-level .vscode directory as it is +# used to configure VS Code settings +!.vscode + +# Outputs +**/runs/* +**/logs/* +**/recordings/* +**/output/* +**/outputs/* +**/videos/* +**/wandb/* +**/.neptune/* +docker/artifacts/ +*.tmp + +# Isaac-Sim packman +_isaac_sim* +_repo +_build +.lastformat + + +**/usd/* diff --git a/.vscode/.gitignore b/.vscode/.gitignore new file mode 100644 index 0000000..10b0af3 --- /dev/null +++ b/.vscode/.gitignore @@ -0,0 +1,10 @@ +# Note: These files are kept for development purposes only. +!tools/launch.template.json +!tools/settings.template.json +!tools/setup_vscode.py +!extensions.json +!tasks.json + +# Ignore all other files +.python.env +*.json diff --git a/.vscode/tools/setup_vscode.py b/.vscode/tools/setup_vscode.py new file mode 100644 index 0000000..07bf6f4 --- /dev/null +++ b/.vscode/tools/setup_vscode.py @@ -0,0 +1,220 @@ +# Copyright (c) 2024-2025, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""This script sets up the vs-code settings for the Isaac Lab project. + +This script merges the python.analysis.extraPaths from the "{ISAACSIM_DIR}/.vscode/settings.json" file into +the ".vscode/settings.json" file. + +This is necessary because Isaac Sim 2022.2.1 onwards does not add the necessary python packages to the python path +when the "setup_python_env.sh" is run as part of the vs-code launch configuration. +""" + +import argparse +import os +import pathlib +import platform +import re +import sys + +PROJECT_DIR = pathlib.Path(__file__).parents[2] +"""Path to the the project directory.""" + +try: + import isaacsim # noqa: F401 + + isaacsim_dir = os.environ.get("ISAAC_PATH", "") +except ModuleNotFoundError or ImportError: + # Create a parser to get the isaac-sim path + parser = argparse.ArgumentParser(description="Setup the VSCode settings for the project.") + parser.add_argument("--isaac_path", type=str, help="The absolute path to the Isaac Sim installation.") + args = parser.parse_args() + + # parse the isaac-sim directory + isaacsim_dir = args.isaac_path + # check if the isaac-sim directory is provided + if not os.path.exists(isaacsim_dir): + raise FileNotFoundError( + f"Could not find the isaac-sim directory: {isaacsim_dir}. Please provide the correct path to the Isaac Sim" + " installation." + ) +except EOFError: + print("Unable to trigger EULA acceptance. This is likely due to the script being run in a non-interactive shell.") + print("Please run the script in an interactive shell to accept the EULA.") + print("Skipping the setup of the VSCode settings...") + sys.exit(0) + +# check if the isaac-sim directory exists +if not os.path.exists(isaacsim_dir): + raise FileNotFoundError( + f"Could not find the isaac-sim directory: {isaacsim_dir}. There are two possible reasons for this:" + "\n\t1. The Isaac Sim directory does not exist as provided CLI path." + "\n\t2. The script could import the 'isaacsim' package. This could be due to the 'isaacsim' package not being " + "installed in the Python environment.\n" + "\nPlease make sure that the Isaac Sim directory exists or that the 'isaacsim' package is installed." + ) + +ISAACSIM_DIR = isaacsim_dir +"""Path to the isaac-sim directory.""" + + +def overwrite_python_analysis_extra_paths(isaaclab_settings: str) -> str: + """Overwrite the python.analysis.extraPaths in the Isaac Lab settings file. + + The extraPaths are replaced with the path names from the isaac-sim settings file that exists in the + "{ISAACSIM_DIR}/.vscode/settings.json" file. + + If the isaac-sim settings file does not exist, the extraPaths are not overwritten. + + Args: + isaaclab_settings: The settings string to use as template. + + Returns: + The settings string with overwritten python analysis extra paths. + """ + # isaac-sim settings + isaacsim_vscode_filename = os.path.join(ISAACSIM_DIR, ".vscode", "settings.json") + + # we use the isaac-sim settings file to get the python.analysis.extraPaths for kit extensions + # if this file does not exist, we will not add any extra paths + if os.path.exists(isaacsim_vscode_filename): + # read the path names from the isaac-sim settings file + with open(isaacsim_vscode_filename) as f: + vscode_settings = f.read() + # extract the path names + # search for the python.analysis.extraPaths section and extract the contents + settings = re.search( + r"\"python.analysis.extraPaths\": \[.*?\]", vscode_settings, flags=re.MULTILINE | re.DOTALL + ) + settings = settings.group(0) + settings = settings.split('"python.analysis.extraPaths": [')[-1] + settings = settings.split("]")[0] + + # read the path names from the isaac-sim settings file + path_names = settings.split(",") + path_names = [path_name.strip().strip('"') for path_name in path_names] + path_names = [path_name for path_name in path_names if len(path_name) > 0] + + # change the path names to be relative to the Isaac Lab directory + rel_path = os.path.relpath(ISAACSIM_DIR, PROJECT_DIR) + path_names = ['"${workspaceFolder}/' + rel_path + "/" + path_name + '"' for path_name in path_names] + else: + path_names = [] + print( + f"[WARN] Could not find Isaac Sim VSCode settings: {isaacsim_vscode_filename}." + "\n\tThis will result in missing 'python.analysis.extraPaths' in the VSCode" + "\n\tsettings, which limits the functionality of the Python language server." + "\n\tHowever, it does not affect the functionality of the Isaac Lab project." + "\n\tWe are working on a fix for this issue with the Isaac Sim team." + ) + + # add the path names that are in the Isaac Lab extensions directory + isaaclab_extensions = os.listdir(os.path.join(PROJECT_DIR, "source")) + path_names.extend(['"${workspaceFolder}/source/' + ext + '"' for ext in isaaclab_extensions]) + + # combine them into a single string + path_names = ",\n\t\t".expandtabs(4).join(path_names) + # deal with the path separator being different on Windows and Unix + path_names = path_names.replace("\\", "/") + + # replace the path names in the Isaac Lab settings file with the path names parsed + isaaclab_settings = re.sub( + r"\"python.analysis.extraPaths\": \[.*?\]", + '"python.analysis.extraPaths": [\n\t\t'.expandtabs(4) + path_names + "\n\t]".expandtabs(4), + isaaclab_settings, + flags=re.DOTALL, + ) + # return the Isaac Lab settings string + return isaaclab_settings + + +def overwrite_default_python_interpreter(isaaclab_settings: str) -> str: + """Overwrite the default python interpreter in the Isaac Lab settings file. + + The default python interpreter is replaced with the path to the python interpreter used by the + isaac-sim project. This is necessary because the default python interpreter is the one shipped with + isaac-sim. + + Args: + isaaclab_settings: The settings string to use as template. + + Returns: + The settings string with overwritten default python interpreter. + """ + # read executable name + python_exe = os.path.normpath(sys.executable) + + # replace with Isaac Sim's python.sh or python.bat scripts to make sure python with correct + # source paths is set as default + if f"kit{os.sep}python{os.sep}bin{os.sep}python" in python_exe: + # Check if the OS is Windows or Linux to use appropriate shell file + if platform.system() == "Windows": + python_exe = python_exe.replace(f"kit{os.sep}python{os.sep}bin{os.sep}python3", "python.bat") + else: + python_exe = python_exe.replace(f"kit{os.sep}python{os.sep}bin{os.sep}python3", "python.sh") + + # replace the default python interpreter in the Isaac Lab settings file with the path to the + # python interpreter in the Isaac Lab directory + isaaclab_settings = re.sub( + r"\"python.defaultInterpreterPath\": \".*?\"", + f'"python.defaultInterpreterPath": "{python_exe}"', + isaaclab_settings, + flags=re.DOTALL, + ) + # return the Isaac Lab settings file + return isaaclab_settings + + +def main(): + # Isaac Lab template settings + isaaclab_vscode_template_filename = os.path.join(PROJECT_DIR, ".vscode", "tools", "settings.template.json") + # make sure the Isaac Lab template settings file exists + if not os.path.exists(isaaclab_vscode_template_filename): + raise FileNotFoundError( + f"Could not find the Isaac Lab template settings file: {isaaclab_vscode_template_filename}" + ) + # read the Isaac Lab template settings file + with open(isaaclab_vscode_template_filename) as f: + isaaclab_template_settings = f.read() + + # overwrite the python.analysis.extraPaths in the Isaac Lab settings file with the path names + isaaclab_settings = overwrite_python_analysis_extra_paths(isaaclab_template_settings) + # overwrite the default python interpreter in the Isaac Lab settings file with the path to the + # python interpreter used to call this script + isaaclab_settings = overwrite_default_python_interpreter(isaaclab_settings) + + # add template notice to the top of the file + header_message = ( + "// This file is a template and is automatically generated by the setup_vscode.py script.\n" + "// Do not edit this file directly.\n" + "// \n" + f"// Generated from: {isaaclab_vscode_template_filename}\n" + ) + isaaclab_settings = header_message + isaaclab_settings + + # write the Isaac Lab settings file + isaaclab_vscode_filename = os.path.join(PROJECT_DIR, ".vscode", "settings.json") + with open(isaaclab_vscode_filename, "w") as f: + f.write(isaaclab_settings) + + # copy the launch.json file if it doesn't exist + isaaclab_vscode_launch_filename = os.path.join(PROJECT_DIR, ".vscode", "launch.json") + isaaclab_vscode_template_launch_filename = os.path.join(PROJECT_DIR, ".vscode", "tools", "launch.template.json") + if not os.path.exists(isaaclab_vscode_launch_filename): + # read template launch settings + with open(isaaclab_vscode_template_launch_filename) as f: + isaaclab_template_launch_settings = f.read() + # add header + header_message = header_message.replace( + isaaclab_vscode_template_filename, isaaclab_vscode_template_launch_filename + ) + isaaclab_launch_settings = header_message + isaaclab_template_launch_settings + # write the Isaac Lab launch settings file + with open(isaaclab_vscode_launch_filename, "w") as f: + f.write(isaaclab_launch_settings) + + +if __name__ == "__main__": + main() diff --git a/CONTRIBUTORS.md b/CONTRIBUTORS.md new file mode 100644 index 0000000..829eafb --- /dev/null +++ b/CONTRIBUTORS.md @@ -0,0 +1,23 @@ +# RL Training Developers and Contributors + +This is the official list of RL Training Project developers and contributors. + +To see the full list of contributors, please check the revision history in the source control. + +Guidelines for modifications: + +* Please keep the lists sorted alphabetically. +* Names should be added to this file as: *individual names* or *organizations*. + +## Developers + +* [Bo Peng](https://www.linkedin.com/in/bo-peng-647068272/) + +## Contributors + +* [Ziqi Fan](https://github.com/fan-ziqi) +* Zhongliang Pan + +## Acknowledgements + + diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..6c1a48b --- /dev/null +++ b/LICENSE @@ -0,0 +1,28 @@ +BSD 3-Clause License + +Copyright (c) 2025, DEEPRoboticsLab + +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. diff --git a/LICENSE-robot_lab b/LICENSE-robot_lab new file mode 100644 index 0000000..e08a51a --- /dev/null +++ b/LICENSE-robot_lab @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright 2024-2025 Ziqi Fan + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/README.md b/README.md new file mode 100644 index 0000000..4c5228e --- /dev/null +++ b/README.md @@ -0,0 +1,180 @@ +rl_training — 火蚁独立训练版本 + +> [!IMPORTANT] +> 本仓库是火蚁（Solenopsis）轮足平台的独立发布副本。火蚁专用训练、Stage 2、URDF与资产说明请先阅读 [README_HUOYI.md](README_HUOYI.md)。模型资产已直接包含在仓库中，不需要初始化 Git 子模块。 + +[![IsaacSim](https://img.shields.io/badge/IsaacSim-5.1.0-silver.svg)](https://docs.omniverse.nvidia.com/isaacsim/latest/overview.html) +[![Isaac Lab](https://img.shields.io/badge/IsaacLab-2.3.0-silver)](https://isaac-sim.github.io/IsaacLab) +[![Discord](https://img.shields.io/badge/-Discord-5865F2?style=flat&logo=Discord&logoColor=white)](https://discord.gg/gdM9mQutC8) +[![Python](https://img.shields.io/badge/python-3.11-blue.svg)](https://docs.python.org/3/whatsnew/3.11.html) +[![Linux platform](https://img.shields.io/badge/platform-linux--64-orange.svg)](https://releases.ubuntu.com/22.04/) +[![License](https://img.shields.io/badge/license-BSD%203--Clause-blue.svg)](https://opensource.org/license/bsd-3-clause) + +## Tutorial Videos +We've released the following tutorials for training and deploying a reinforcement learning policy. Please check it out on [Bilibili](https://b23.tv/UoIqsFn) or [YouTube](https://youtube.com/playlist?list=PLy9YHJvMnjO0X4tx_NTWugTUMJXUrOgFH&si=pjUGF5PbFf3tGLFz)! + +## Overview + +**rl_training** is a RL training library for deeprobotics robots, based on IsaacLab. The table below lists all available environments: + +| Robot Model | Environment Name (ID) | Screenshot | +|---------------------|------------------------------------------------------------|------------| +| [Deeprobotics Lite3](https://www.deeprobotics.cn/robot/index/product1.html) | Rough-Deeprobotics-Lite3-v0 | Lite3 +| [Deeprobotics M20](https://www.deeprobotics.cn/robot/index/lynx.html) | Rough-Deeprobotics-M20-v0 | deeprobotics_m20 + +> [!NOTE] +> If you want to deploy policies in mujoco or real robots, please use the corresponding deploy repo in [Deep Robotics Github Center](https://github.com/DeepRoboticsLab). +## Contribution + +Everyone is welcome to contribute to this repo. If you discover a bug or optimize our training config, just submit a pull request and we will look into it. +## Installation + +- Install Isaac Lab by following the [installation guide](https://isaac-sim.github.io/IsaacLab/main/source/setup/installation/index.html). We recommend using the conda installation as it simplifies calling Python scripts from the terminal. + +- Clone this repository separately from the Isaac Lab installation (i.e. outside the `IsaacLab` directory): + + ```bash + git clone https://github.com/weifenghuang812/huoyi_training.git + cd huoyi_training + ``` + +- Using a python interpreter that has Isaac Lab installed, install the library + + ```bash + python -m pip install -e source/rl_training + ``` + +- Verify that the extension is correctly installed by running the following command to print all the available environments in the extension: + + ```bash + python scripts/tools/list_envs.py + ``` + +
+ +Setup as Omniverse Extension (Optional, click to expand) + +We provide an example UI extension that will load upon enabling your extension defined in `source/rl_training/rl_training/ui_extension_example.py`. + +To enable your extension, follow these steps: + +1. **Add the search path of your repository** to the extension manager: + - Navigate to the extension manager using `Window` -> `Extensions`. + - Click on the **Hamburger Icon** (☰), then go to `Settings`. + - In the `Extension Search Paths`, enter the absolute path to `rl_trainingb/source` + - If not already present, in the `Extension Search Paths`, enter the path that leads to Isaac Lab's extension directory directory (`IsaacLab/source`) + - Click on the **Hamburger Icon** (☰), then click `Refresh`. + +2. **Search and enable your extension**: + - Find your extension under the `Third Party` category. + - Toggle it to enable your extension. + +
+ +## Try examples + +Deeprobotics Lite3: + +```bash +# Train +python scripts/reinforcement_learning/rsl_rl/train.py --task=Rough-Deeprobotics-Lite3-v0 --headless + +# Play +python scripts/reinforcement_learning/rsl_rl/play.py --task=Rough-Deeprobotics-Lite3-v0 --num_envs=10 +``` + +Deeprobotics M20: + +```bash +# Train +python scripts/reinforcement_learning/rsl_rl/train.py --task=Rough-Deeprobotics-M20-v0 --headless + +# Play +python scripts/reinforcement_learning/rsl_rl/play.py --task=Rough-Deeprobotics-M20-v0 --num_envs=10 +``` + +> [!NOTE] +> If you want to control a **SINGLE ROBOT** with the keyboard during playback, add `--keyboard` at the end of the play script. +> +> ``` +> Key bindings: +> ====================== ========================= ======================== +> Command Key (+ve axis) Key (-ve axis) +> ====================== ========================= ======================== +> Move along x-axis Numpad 8 / Arrow Up Numpad 2 / Arrow Down +> Move along y-axis Numpad 4 / Arrow Right Numpad 6 / Arrow Left +> Rotate along z-axis Numpad 7 / Z Numpad 9 / X +> ====================== ========================= ======================== +> ``` +* Record video of a trained agent (requires installing `ffmpeg`), add `--video --video_length 200` +* Play/Train with 32 environments, add `--num_envs 32` +* Play on specific folder or checkpoint, add `--load_run run_folder_name --checkpoint model.pt` +* Resume training from folder or checkpoint, add `--resume --load_run run_folder_name --checkpoint model.pt` + +## Trained Results +You can download our trained examples for reference of the training process. This should be expected if you run our training process properly. [M20](https://drive.google.com/file/d/1PAGIHOI6Lsge_JA9kU3w7RU-GyqXybob/view?usp=drive_link) and [Lite3](https://drive.google.com/file/d/1EgCVQA552x3YiaRhPRv675TSd01yqF0c/view?usp=drive_link) + + + +## Multi-gpu acceleration +* To train with multiple GPUs, use the following command, where --nproc_per_node represents the number of available GPUs: + ```bash + python -m torch.distributed.run --nnodes=1 --nproc_per_node=2 scripts/reinforcement_learning/rsl_rl/train.py --task= --headless + python -m torch.distributed.run --nnodes=1 --nproc_per_node=2 scripts/reinforcement_learning/rsl_rl/train.py --task=Rough-Deeprobotics-Lite3-v0 --headless --distributed --num_envs=2048 + ``` +* Note: each gpu will have the same number of envs specified in the config, to use the previous total number of envs, devide it by the number of gpus. +* To scale up training beyond multiple GPUs on a single machine, it is also possible to train across multiple nodes. To train across multiple nodes/machines, it is required to launch an individual process on each node. + + For the master node, use the following command, where --nproc_per_node represents the number of available GPUs, and --nnodes represents the number of nodes: + ```bash + python -m torch.distributed.run --nproc_per_node=2 --nnodes=2 --node_rank=0 --rdzv_id=123 --rdzv_backend=c10d --rdzv_endpoint=localhost:5555 scripts/reinforcement_learning/rsl_rl/train.py --task= --headless --distributed + ``` + Note that the port (`5555`) can be replaced with any other available port. + For non-master nodes, use the following command, replacing `--node_rank` with the index of each machine: + ```bash + python -m torch.distributed.run --nproc_per_node=2 --nnodes=2 --node_rank=1 --rdzv_id=123 --rdzv_backend=c10d --rdzv_endpoint=ip_of_master_machine:5555 scripts/reinforcement_learning/rsl_rl/train.py --task= --headless --distributed + ``` + +## Tensorboard + +To view tensorboard, run: + +```bash +tensorboard --logdir=logs +``` + +## Troubleshooting + +### Pylance Missing Indexing of Extensions + +In some VsCode versions, the indexing of part of the extensions is missing. In this case, add the path to your extension in `.vscode/settings.json` under the key `"python.analysis.extraPaths"`. + +**Note: Replace `` with your own IsaacLab path.** + +```json +{ + "python.languageServer": "Pylance", + "python.analysis.extraPaths": [ + "${workspaceFolder}/source/rl_training", + "//source/isaaclab", + "//source/isaaclab_assets", + "//source/isaaclab_mimic", + "//source/isaaclab_rl", + "//source/isaaclab_tasks", + ] +} +``` + +### Clean USD Caches + +Temporary USD files are generated in `/tmp/IsaacLab/usd_{date}_{time}_{random}` during simulation runs. These files can consume significant disk space and can be cleaned by: + +```bash +rm -rf /tmp/IsaacLab/usd_* +``` + +## Acknowledgements + +The project uses some code from the following open-source code repositories: + +- [fan-ziqi/robot_lab](https://github.com/fan-ziqi/robot_lab) diff --git a/README_HUOYI.md b/README_HUOYI.md new file mode 100644 index 0000000..6985a9b --- /dev/null +++ b/README_HUOYI.md @@ -0,0 +1,76 @@ +# 火蚁轮足机器人强化学习项目 + +这是火蚁（Solenopsis）轮足平台当前可运行训练版本的独立发布副本。 + +## 包含内容 + +- Isaac Lab / RSL-RL 训练与播放脚本 +- 火蚁 Stage 1 与 Stage 2 多地形任务配置 +- 火蚁质量、惯量、电机和接触参数 +- `Solenopsis.urdf` +- 当前 URDF 使用的轻量可视化 STL（碰撞体仍为简化几何体） +- 评估脚本和模型对比文档 + +为了保持仓库轻量，本仓库不包含训练日志、TensorBoard 数据、历史 checkpoint、原始高精度 CAD STL、旧 USD 备份或本地 Codex 快照。 + +## 克隆 + +```bash +git clone https://github.com/weifenghuang812/huoyi_training.git +cd huoyi_training +``` + +本仓库已经将火蚁模型直接包含在 `deep_robotics_model/` 中，不需要执行 `git submodule update`。 + +## 安装 + +先按项目原始 `README.md` 安装匹配版本的 Isaac Lab，然后在仓库根目录执行： + +```bash +python -m pip install -e source/rl_training +``` + +## 从零训练 Stage 1 + +```bash +python scripts/reinforcement_learning/rsl_rl/train.py \ + --task Rough-Deeprobotics-M20-v0 \ + --headless +``` + +## 训练 Stage 2 + +Stage 2 应从 Stage 1 checkpoint 恢复： + +```bash +python scripts/reinforcement_learning/rsl_rl/train.py \ + --task Rough-Deeprobotics-M20-Support-v0 \ + --headless \ + --resume \ + --load_run /绝对路径/到/stage1运行目录 \ + --load_checkpoint model_xxx.pt +``` + +## 播放 + +```bash +python scripts/reinforcement_learning/rsl_rl/play.py \ + --task Rough-Deeprobotics-M20-v0 \ + --checkpoint /绝对路径/到/model_xxx.pt +``` + +键盘控制时添加： + +```bash +--keyboard +``` + +## 资产说明 + +活动模型路径为： + +```text +deep_robotics_model/M20/M20_urdf/urdf/Solenopsis.urdf +``` + +训练加载 URDF；碰撞体使用方块和圆柱简化，可视化使用轻量 STL。 diff --git a/VERSION b/VERSION new file mode 100644 index 0000000..afaf360 --- /dev/null +++ b/VERSION @@ -0,0 +1 @@ +1.0.0 \ No newline at end of file diff --git a/deep_robotics_model/LICENSE.txt b/deep_robotics_model/LICENSE.txt new file mode 100644 index 0000000..cdc6c06 --- /dev/null +++ b/deep_robotics_model/LICENSE.txt @@ -0,0 +1,28 @@ +BSD 3-Clause License + +Copyright (c) 2024, DeepRoboticsLab + +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. diff --git a/deep_robotics_model/M20/M20_urdf/meshes_solenopsis_visual/base_link.STL b/deep_robotics_model/M20/M20_urdf/meshes_solenopsis_visual/base_link.STL new file mode 100644 index 0000000..657d94c Binary files /dev/null and b/deep_robotics_model/M20/M20_urdf/meshes_solenopsis_visual/base_link.STL differ diff --git a/deep_robotics_model/M20/M20_urdf/meshes_solenopsis_visual/fl_knee.STL b/deep_robotics_model/M20/M20_urdf/meshes_solenopsis_visual/fl_knee.STL new file mode 100644 index 0000000..9866003 Binary files /dev/null and b/deep_robotics_model/M20/M20_urdf/meshes_solenopsis_visual/fl_knee.STL differ diff --git a/deep_robotics_model/M20/M20_urdf/meshes_solenopsis_visual/fl_wheel.STL b/deep_robotics_model/M20/M20_urdf/meshes_solenopsis_visual/fl_wheel.STL new file mode 100644 index 0000000..f3e1ef6 Binary files /dev/null and b/deep_robotics_model/M20/M20_urdf/meshes_solenopsis_visual/fl_wheel.STL differ diff --git a/deep_robotics_model/M20/M20_urdf/meshes_solenopsis_visual/fr_knee.STL b/deep_robotics_model/M20/M20_urdf/meshes_solenopsis_visual/fr_knee.STL new file mode 100644 index 0000000..21b7289 Binary files /dev/null and b/deep_robotics_model/M20/M20_urdf/meshes_solenopsis_visual/fr_knee.STL differ diff --git a/deep_robotics_model/M20/M20_urdf/meshes_solenopsis_visual/fr_wheel.STL b/deep_robotics_model/M20/M20_urdf/meshes_solenopsis_visual/fr_wheel.STL new file mode 100644 index 0000000..8aa87bb Binary files /dev/null and b/deep_robotics_model/M20/M20_urdf/meshes_solenopsis_visual/fr_wheel.STL differ diff --git a/deep_robotics_model/M20/M20_urdf/meshes_solenopsis_visual/hl_knee.STL b/deep_robotics_model/M20/M20_urdf/meshes_solenopsis_visual/hl_knee.STL new file mode 100644 index 0000000..53126d4 Binary files /dev/null and b/deep_robotics_model/M20/M20_urdf/meshes_solenopsis_visual/hl_knee.STL differ diff --git a/deep_robotics_model/M20/M20_urdf/meshes_solenopsis_visual/hl_wheel.STL b/deep_robotics_model/M20/M20_urdf/meshes_solenopsis_visual/hl_wheel.STL new file mode 100644 index 0000000..e88f13a Binary files /dev/null and b/deep_robotics_model/M20/M20_urdf/meshes_solenopsis_visual/hl_wheel.STL differ diff --git a/deep_robotics_model/M20/M20_urdf/meshes_solenopsis_visual/hr_knee.STL b/deep_robotics_model/M20/M20_urdf/meshes_solenopsis_visual/hr_knee.STL new file mode 100644 index 0000000..1d49836 Binary files /dev/null and b/deep_robotics_model/M20/M20_urdf/meshes_solenopsis_visual/hr_knee.STL differ diff --git a/deep_robotics_model/M20/M20_urdf/meshes_solenopsis_visual/hr_wheel.STL b/deep_robotics_model/M20/M20_urdf/meshes_solenopsis_visual/hr_wheel.STL new file mode 100644 index 0000000..4086a0c Binary files /dev/null and b/deep_robotics_model/M20/M20_urdf/meshes_solenopsis_visual/hr_wheel.STL differ diff --git a/deep_robotics_model/M20/M20_urdf/meshes_solenopsis_visual/optimization_manifest.json b/deep_robotics_model/M20/M20_urdf/meshes_solenopsis_visual/optimization_manifest.json new file mode 100644 index 0000000..7ffc024 --- /dev/null +++ b/deep_robotics_model/M20/M20_urdf/meshes_solenopsis_visual/optimization_manifest.json @@ -0,0 +1,92 @@ +{ + "source_dir": "/home/ycpy/huangweifeng/hwf_workspace/rl_training_huoyi/rl_training/deep_robotics_model/M20/M20_urdf/meshes_solenopsis", + "grid_mm": 2.0, + "maximum_vertex_displacement_mm": 1.7320508075688772, + "source_faces": 2146372, + "optimized_faces": 234820, + "source_bytes": 107319356, + "optimized_bytes": 11741756, + "meshes": [ + { + "file": "base_link.STL", + "source_vertices": 568586, + "source_faces": 1140292, + "optimized_vertices": 35724, + "optimized_faces": 86170, + "source_bytes": 57014684, + "optimized_bytes": 4308584 + }, + { + "file": "fl_knee.STL", + "source_vertices": 54172, + "source_faces": 108311, + "optimized_vertices": 6368, + "optimized_faces": 14405, + "source_bytes": 5415634, + "optimized_bytes": 720334 + }, + { + "file": "fl_wheel.STL", + "source_vertices": 71661, + "source_faces": 143210, + "optimized_vertices": 11121, + "optimized_faces": 22793, + "source_bytes": 7160584, + "optimized_bytes": 1139734 + }, + { + "file": "fr_knee.STL", + "source_vertices": 54172, + "source_faces": 108311, + "optimized_vertices": 6363, + "optimized_faces": 14398, + "source_bytes": 5415634, + "optimized_bytes": 719984 + }, + { + "file": "fr_wheel.STL", + "source_vertices": 71668, + "source_faces": 143224, + "optimized_vertices": 11122, + "optimized_faces": 22796, + "source_bytes": 7161284, + "optimized_bytes": 1139884 + }, + { + "file": "hl_knee.STL", + "source_vertices": 54172, + "source_faces": 108311, + "optimized_vertices": 6340, + "optimized_faces": 14346, + "source_bytes": 5415634, + "optimized_bytes": 717384 + }, + { + "file": "hl_wheel.STL", + "source_vertices": 71660, + "source_faces": 143208, + "optimized_vertices": 11120, + "optimized_faces": 22790, + "source_bytes": 7160484, + "optimized_bytes": 1139584 + }, + { + "file": "hr_knee.STL", + "source_vertices": 54172, + "source_faces": 108311, + "optimized_vertices": 6331, + "optimized_faces": 14331, + "source_bytes": 5415634, + "optimized_bytes": 716634 + }, + { + "file": "hr_wheel.STL", + "source_vertices": 71653, + "source_faces": 143194, + "optimized_vertices": 11121, + "optimized_faces": 22791, + "source_bytes": 7159784, + "optimized_bytes": 1139634 + } + ] +} diff --git a/deep_robotics_model/M20/M20_urdf/urdf/Solenopsis.urdf b/deep_robotics_model/M20/M20_urdf/urdf/Solenopsis.urdf new file mode 100644 index 0000000..f09c2d8 --- /dev/null +++ b/deep_robotics_model/M20/M20_urdf/urdf/Solenopsis.urdf @@ -0,0 +1,292 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/deep_robotics_model/README.md b/deep_robotics_model/README.md new file mode 100644 index 0000000..d54bd42 --- /dev/null +++ b/deep_robotics_model/README.md @@ -0,0 +1,18 @@ +# Robot Models Repository + +This repository contains 3D models for three Deep Robotics robots: **Lite3**, **M20**, and **X30**. Each model is provided in MJCF, URDF, and USD formats. To visualize the model and do some measuring, I recommend using this [tool](https://viewer.robotsfan.com/). Just drag your folder to the web page and you can see everything. + +**Note**: This repository only contains low resolution files. To download the high resolution files, please use [this link](https://drive.google.com/drive/folders/1EOELXUYSBPEJeD0rUIkJxnlLMDr6IHBV?usp=sharing) + + +## Model Overview +**Note**: High resolution models are good for new simulators like isaaclab/isaacsim but old simulators like pybullet cannot open files this big. +| Robot Model | High Resolution Image | Low Resolution Image | +|:-----------:|:---------------------:|:--------------------:| +| **Lite3** |
Lite3/Lite3_urdf/urdf/Lite3_high_res.urdf |
Lite3/Lite3_urdf/urdf/Lite3.urdf | +| **X30** |
X30/X30_urdf/urdf/X30_high_res.urdf |
X30/X30_urdf/urdf/X30.urdf | +| **M20** |
M20/M20_urdf/urdf/M20_high_res.urdf |
M20/M20_urdf/urdf/M20.urdf | + +## Contributors +See the [Contributors](Contributors.md) page for a list of contributors. + diff --git a/docs/fire_ant_prototype_specs.md b/docs/fire_ant_prototype_specs.md new file mode 100644 index 0000000..de445f4 --- /dev/null +++ b/docs/fire_ant_prototype_specs.md @@ -0,0 +1,178 @@ +# 火蚁样机参数简表 + +本文档依据当前火蚁仿真模型与训练配置整理： + +- URDF: `deep_robotics_model/M20/M20_urdf/urdf/M20.urdf` +- 执行器配置: `source/rl_training/rl_training/assets/deeprobotics.py` +- 训练配置: `source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/config/wheeled/deeprobotics_m20/rough_env_cfg.py` + +## 1. 结构形式 + +火蚁当前为 8 自由度轮腿机器人 +## 2. 质量参数 +整机仿真总质量约：23.690 kg + +| 部件 | 单个质量 | 数量 | 小计 | +| --- | ---: | ---: | ---: | +| 机身 base_link | 15.882 kg | 1 | 15.882 kg | +| 膝部连杆（按 P65MAX-S 质量建模） | 0.612 kg | 4 | 2.448 kg | +| 轮毂电机与轮胎（W190-S） | 1.340 kg | 4 | 5.360 kg | + +## 3. 尺寸参数 + +机身主体： + + +长 x 宽 x 高 = 0.60 m x 0.32 m x 0.12 m + + +轮子： + + +半径 = 0.096 m +直径 = 0.192 m +轴向总尺寸 = 0.083 m（手册工程图标注 82.95 mm） +``` + +膝部连杆简化尺寸： + + +主连杆 box = 0.05 m x 0.024 m x 0.20 m +末端圆柱半径 = 0.04 m +末端圆柱长度 = 0.09 m + + +零关节角几何尺寸： + +```text +轴距 = 0.60 m +轮距 = 0.50 m +轮心相对机身中心高度 = -0.21 m +``` + +训练初始姿态下，前膝为 `+0.5 rad`，后膝为 `-0.5 rad`： + +```text +有效轴距约 = 0.801 m +轮距 = 0.50 m +机身中心接地高度约 = 0.280 m +机身底部离地约 = 0.220 m +``` + +## 4. 电机与执行器参数 + +### 膝关节电机 + +仿真限制： + +```text +型号 = P65MAX-S +额定扭矩 = 21.3 Nm +峰值扭矩 = 75.0 Nm +速度上限 = 17.8024 rad/s ≈ 170 rpm +``` + +IsaacLab 控制参数： + +```text +位置刚度 stiffness = 80.0 +阻尼 damping = 2.0 +控制延迟 = 0 ~ 5 step +``` + +当前 URDF 与 Isaac Lab actuator 均使用相同的峰值扭矩和速度上限。 + +### 轮端电机 + +仿真实际执行器限制： + +```text +型号 = W190-S +额定扭矩 = 8.7 Nm +峰值扭矩 = 28.68 Nm +额定点转速 = 94.457 rad/s ≈ 902 rpm +速度上限 = 109.956 rad/s ≈ 1050 rpm +理论空载轮缘线速度 = 10.56 m/s +``` + +URDF 中轮关节写有： + +```text +effort = 28.68 Nm +velocity = 109.956 rad/s +``` + +URDF 与 Isaac Lab actuator 的限制现已保持一致。 + +## 5. 训练需求 + +Stage 1 基础运动： + +```text +前后速度 lin_vel_x = -2.0 ~ 2.0 m/s +侧向速度 lin_vel_y = 0 +转向 yaw_rate = -1.5 ~ 1.5 rad/s +``` + +Stage 2 超支撑越障： + +```text +前进速度 lin_vel_x = 0.25 ~ 1.2 m/s +侧向速度 lin_vel_y = 0 +转向 yaw_rate = -0.6 ~ 0.6 rad/s +``` + +Stage 2 主要地形难度： + +```text +boxes 高度 = 0.03 ~ 0.22 m +上台阶高度 = 0.04 ~ 0.24 m +下台阶高度 = 0.04 ~ 0.20 m +随机粗糙高度 = 0.01 ~ 0.14 m +``` + +## 6. 样机设计重点 + +优先级建议： + +1. 膝关节峰值扭矩优先于轮端极速。 +2. 轮端要保证低速大扭矩，避免顶障碍时堵转。 +3. 机身底部、膝部、轮架必须能承受超支撑接触冲击。 +4. 轮胎材料要保证摩擦和一定缓冲，不能只看刚性轮模型。 +5. 电机驱动器需要考虑峰值电流、连续散热和限流策略。 + +## 7. 当前仿真限制 + +接触材质配置： + +| 接触对象 | 静摩擦随机范围 | 动摩擦随机范围 | 恢复系数随机范围 | +| --- | ---: | ---: | ---: | +| W190-S 轮子 | 0.9 ~ 1.4 | 0.6 ~ 0.9 | 0.0 ~ 0.05 | +| base_link 与 knee links | 0.4 ~ 0.8 | 0.25 ~ 0.4 | 0.0 ~ 0.05 | + +地形静、动摩擦均为 1.0，摩擦使用乘法组合，因此轮地有效摩擦范围与轮子材质随机范围一致。火蚁任务将地形恢复系数覆盖为 0.0，因此轮地碰撞的有效恢复系数为 0。质量与惯量随机化保持启用。 + +当前模型不能直接等价为最终硬件设计： + +- 几何为简化 box/cylinder，不是完整 CAD。 +- 没有真实电机扭矩-转速曲线。 +- 没有电池、电流限制、热衰减。 +- 没有减速器效率、背隙、柔性结构。 +- 没有轮胎变形模型。 +- 接触与摩擦比真实世界更理想。 +- 轮关节 URDF effort 与 IsaacLab actuator effort 不一致，训练实际以 actuator 为准。 + +## 8. 初版样机目标 + +建议以如下指标作为第一版样机基准： + +```text +整机仿真质量：约 23.69 kg +轮径：0.192 m +轮距：0.50 m +有效轴距：约 0.80 m +膝关节峰值扭矩：75 Nm +轮端峰值扭矩：28.68 Nm +机身底部离地：约 0.22 m +目标越障高度：0.20 ~ 0.24 m +``` diff --git a/docs/imgs/deeprobotics_lite3.png b/docs/imgs/deeprobotics_lite3.png new file mode 100644 index 0000000..8b09a83 Binary files /dev/null and b/docs/imgs/deeprobotics_lite3.png differ diff --git a/docs/imgs/deeprobotics_m20.png b/docs/imgs/deeprobotics_m20.png new file mode 100644 index 0000000..5570644 Binary files /dev/null and b/docs/imgs/deeprobotics_m20.png differ diff --git a/docs/motor_model_update_comparison.md b/docs/motor_model_update_comparison.md new file mode 100644 index 0000000..ba41986 --- /dev/null +++ b/docs/motor_model_update_comparison.md @@ -0,0 +1,51 @@ +# P65MAX-S / W190-S 模型更新对比 + +数据来源：智身科技《关节选型手册 2026》。本次将简化 M20 URDF 的关节电机替换为 P65MAX-S，将轮毂电机替换为 W190-S。 + +## 关节与执行器 + +| 参数 | 修改前 | 修改后 | +| --- | ---: | ---: | +| 膝关节峰值扭矩 | 76.4 Nm | 75.0 Nm | +| 膝关节速度上限 | 22.4 rad/s（约 214 rpm） | 17.8024 rad/s（170 rpm） | +| 单个 knee link 质量 | 1.000 kg | 0.612 kg | +| 轮关节 URDF 扭矩上限 | 40.0 Nm | 28.68 Nm | +| 轮关节训练执行器扭矩上限 | 21.6 Nm | 28.68 Nm | +| 轮关节速度上限 | 79.3 rad/s（约 757 rpm） | 109.956 rad/s（1050 rpm） | + +## 轮毂几何与质量 + +| 参数 | 修改前 | 修改后 | +| --- | ---: | ---: | +| 轮半径 | 0.140 m | 0.096 m | +| 轮直径 | 0.280 m | 0.192 m | +| 轴向尺寸 | 0.090 m | 0.083 m | +| 单轮质量 | 1.400 kg | 1.340 kg | +| 绕轮轴转动惯量 | 0.013700 kg·m² | 0.006174720 kg·m² | +| 径向转动惯量 | 0.007500 kg·m² | 0.003856632 kg·m² | +| 空载理论轮缘速度 | 11.10 m/s | 10.56 m/s | + +W190-S 的新转动惯量按质量 1.34 kg、半径 0.096 m、轴向尺寸 0.083 m 的均匀实心圆柱重新计算。 + +## 整机影响 + +| 参数 | 修改前 | 修改后 | 变化 | +| --- | ---: | ---: | ---: | +| URDF 总质量 | 25.482 kg | 23.690 kg | -1.792 kg | +| 初始生成高度 | 0.350 m | 0.306 m | -0.044 m | +| 初始姿态稳定后机身中心高度 | 约 0.324 m | 约 0.280 m | -0.044 m | +| 初始姿态机身底部离地 | 约 0.264 m | 约 0.220 m | -0.044 m | +| 轮距 | 0.500 m | 0.500 m | 不变 | +| 有效轴距 | 约 0.801 m | 约 0.801 m | 不变 | + +## 同步修改位置 + +- `deep_robotics_model/M20/M20_urdf/urdf/M20.urdf`：质量、惯量、轮子几何、关节限制。 +- `source/rl_training/rl_training/assets/deeprobotics.py`：训练执行器限制和初始生成高度。 +- `source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/config/wheeled/deeprobotics_m20/rough_env_cfg.py`：差速转向奖励和轮滑观测所用轮半径。 + +## 建模边界 + +- 简化 URDF 没有独立的电机刚体，因此将 knee link 质量直接设为 P65MAX-S 的 0.612 kg，并按质量比例缩放原 knee link 惯量。 +- W190-S 手册给出的是约 83 mm 的轴向总尺寸，简化圆柱模型将其作为轮宽。 +- 当前执行器仍使用固定扭矩与固定速度上限，没有拟合手册中的完整扭矩—转速曲线、热衰减或电流限制。 diff --git a/docs/solenopsis_urdf_comparison.md b/docs/solenopsis_urdf_comparison.md new file mode 100644 index 0000000..6b90b7b --- /dev/null +++ b/docs/solenopsis_urdf_comparison.md @@ -0,0 +1,114 @@ +# Solenopsis CAD URDF 与旧 M20 简化 URDF 对比 + +对比对象： + +- 旧模型：`deep_robotics_model/M20/M20_urdf/urdf/M20.urdf` +- 新模型：`deep_robotics_model/M20/M20_urdf/urdf/Solenopsis.urdf` +- CAD 来源：`Solenopsis-00 火蚁-四单臂轮足平台总成.STEP` + +## 质量闭合 + +用户提供的刚体分组质量为：机身组 17.274 kg、单个摆臂组 0.955 kg、单个轮组 1.340 kg；整机质量为 26.581 kg。 + +```text +17.274 + 4 × 0.955 + 4 × 1.340 = 26.454 kg +26.581 - 26.454 = 0.127 kg +``` + +按确认结果，将差额 0.127 kg（即 127 g，不是 0.127 g）合并到 `base_link`，因此新 `base_link` 质量为 17.401 kg。该差额等效放在已报告的机身组质心处，所以不改变机身组质心坐标和惯量张量；没有假定未知零部件的空间分布。 + +| 刚体 | 旧模型单体质量 | 新模型单体质量 | 单体变化 | +| --- | ---: | ---: | ---: | +| `base_link` | 15.882 kg | 17.401 kg | +1.519 kg（+9.5643%） | +| 每个 `*_knee` | 0.612 kg | 0.955 kg | +0.343 kg（+56.0458%） | +| 每个 `*_wheel` | 1.340 kg | 1.340 kg | 0 | +| 整机 | 23.690 kg | 26.581 kg | +2.891 kg（+12.2035%） | + +## 坐标换算 + +SolidWorks 总装坐标为 Y 轴向上；现有训练模型采用 ROS/URDF 的 Z 轴向上。新模型使用以下右手坐标变换： + +```text +(x, y, z)_URDF = (z, x, y)_SolidWorks +``` + +摆臂和轮子的质量属性原始坐标系是运动模组坐标系，X 轴为电机旋转轴，Y 轴沿轮心指向关节中心。为了保持现有控制器的关节符号约定，新 URDF 统一定义： + +- 四个关节轴均为 `(0, -1, 0)`； +- `q = 0` 时四个摆臂均竖直向下； +- 关节中心到轮轴中心的距离为 CAD 中的精确值 0.300 m； +- 左右侧采用不同的正规旋转矩阵变换质量属性，没有使用镜像矩阵变换惯量。 + +## 质心与惯量 + +所有新惯量均为“关于各自刚体质心、并与对应 URDF link 坐标轴对齐”的惯量，单位为 kg·m²。SolidWorks 的 kg·mm² 已严格乘以 `10^-6`，并按 `I_URDF = R I_SW R^T` 旋转。 + +### base_link + +| 参数 | 旧模型 | 新模型 | +| --- | ---: | ---: | +| 质心 x | 0.003551059 m | -0.000072 m | +| 质心 y | 0.001650758 m | 0.000171 m | +| 质心 z | -0.003464621 m | 0.076892 m | +| Ixx | 0.088078700 | 0.133964838 | +| Iyy | 0.533816000 | 0.363696493 | +| Izz | 0.563915000 | 0.450717599 | +| Ixy | 0.001396000 | -0.000732466 | +| Ixz | 0.000111000 | 0.000157104 | +| Iyz | 0.000519000 | 0.000165399 | + +### 单个 knee + +新左侧 knee 质心为 `(-0.000040, -0.001088, -0.158988) m`，新右侧 knee 质心为 `(0.000040, 0.001088, -0.158988) m`。 + +| 参数 | 旧模型代表值 | 新模型左侧 | 新模型右侧 | +| --- | ---: | ---: | ---: | +| Ixx | 0.005202000 | 0.014124479 | 0.014124479 | +| Iyy | 0.005202000 | 0.014470094 | 0.014470094 | +| Izz | 0.000734400 | 0.000478899 | 0.000478899 | +| Ixy | ±0.000024480 | 0.000002919 | 0.000002919 | +| Ixz | ±0.000091800 | 0.000004851 | -0.000004851 | +| Iyz | ±0.000673200 | -0.000261898 | 0.000261898 | + +### 单个 wheel + +新左轮质心为 `(-0.000002, 0.059326, -0.000009) m`，新右轮质心为 `(0.000002, -0.059326, -0.000009) m`。这保留了 SolidWorks 报告的轮组轴向质心偏置，不再把轮组质心强制放在关节原点。 + +| 参数 | 旧均匀圆柱模型 | 新模型左侧 | 新模型右侧 | +| --- | ---: | ---: | ---: | +| Ixx | 0.003856632 | 0.003296552 | 0.003296552 | +| Iyy（轮轴） | 0.006174720 | 0.005826562 | 0.005826562 | +| Izz | 0.003856632 | 0.003296995 | 0.003296995 | +| Ixy | 0 | 0.000000144 | 0.000000144 | +| Ixz | 0 | -0.000000077 | 0.000000077 | +| Iyz | 0 | 0.000000424 | -0.000000424 | + +## 关节几何 + +| 参数 | 旧简化模型 | 新 CAD 模型 | +| --- | ---: | ---: | +| 前/后关节中心 x | ±0.300 m | ±0.117577164466 m | +| 左/右关节中心 y | ±0.140 m | ±0.167 m | +| 关节中心 z | 0 m | 0.070 m | +| 膝关节至轮轴向量 | `(0, ±0.110, -0.210) m` | `(0, 0, -0.300) m` | +| 膝关节至轮轴距离 | 0.237065 m | 0.300000 m | +| 关节轴 | `(0, -1, 0)` | `(0, -1, 0)` | +| 自由度 | 8 | 8 | + +关节名称、顺序、扭矩上限和速度上限均保持不变，因此软件接口仍是四个 knee 加四个 wheel。不过，质量、质心、惯量、关节位置和杆长都有实质变化，旧策略不能被视为动力学等价策略，应重新训练或至少重新做系统性验证。 + +## 显示与碰撞几何 + +新 URDF 使用从 STEP 总成拆分得到的 9 个 STL 作为 visual：一个机身、四个摆臂、四个轮组。STL 保留毫米坐标，URDF 使用 `0.001` 比例换算到米。当前训练引用 `meshes_solenopsis_visual` 中的可视化网格；高精度原始导出保留在 `meshes_solenopsis` 中。 + +collision 不使用 STL，只使用方块和圆柱，尺寸来自 CAD 外廓： + +- `base_link`：一个 `0.706 × 0.310 × 0.1714 m` 方块； +- 每个 `*_knee`：一个 `0.088 × 0.031 × 0.3765 m` 摆臂方块，加一个 P65 转动端圆柱； +- 每个 `*_wheel`：半径 `0.098804 m`、轴向宽度 `0.082959 m` 的圆柱。 + +训练可视化网格约 23.5 万个三角形、11.7 MB，顶点相对高精度 STL 的最大量化位移不超过 1.733 mm。非闭合边、内部电气件和紧固件不会参与碰撞求解。这样既保留 CAD 外观，又避免 214 万面高精度 STL 拖慢 Isaac Sim 导入，同时让并行训练继续使用稳定的基础几何碰撞体。 + +## 当前切换状态 + +旧 `M20.urdf` 没有被覆盖，IsaacLab训练入口已经切换到 `Solenopsis.urdf`。导入时，visual加载STEP拆分得到的STL用于正常显示，collision只加载方块和圆柱参与物理计算。 diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..02815af --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,85 @@ +[build-system] +requires = ["setuptools", "toml"] +build-backend = "setuptools.build_meta" + +[tool.isort] + +atomic = true +profile = "black" +line_length = 120 +py_version = 310 +skip_glob = ["docs/*", "logs/*"] +group_by_package = true + +sections = [ + "FUTURE", + "STDLIB", + "THIRDPARTY", + "ISAACLABPARTY", + "FIRSTPARTY", + "LOCALFOLDER", +] +extra_standard_library = [ + "numpy", + "h5py", + "open3d", + "torch", + "tensordict", + "bpy", + "matplotlib", + "gymnasium", + "gym", + "scipy", + "hid", + "yaml", + "prettytable", + "toml", + "trimesh", + "tqdm", + "psutil", +] +known_thirdparty = [ + "isaacsim.core.api", + "omni.replicator.core", + "pxr", + "omni.kit.*", + "warp", + "carb", + "Semantics", +] +known_isaaclabparty = [ + "isaaclab", + "isaaclab_tasks", + "isaaclab_assets", + "isaaclab_mimic", + "isaaclab_rl" +] + +# Modify the following to include the package names of your first-party code +known_firstparty = "rl_training" +known_local_folder = "config" + +[tool.pyright] + +exclude = [ + "**/__pycache__", + "**/docs", + "**/logs", + ".git", + ".vscode", +] + +typeCheckingMode = "basic" +pythonVersion = "3.11" +pythonPlatform = "Linux" +enableTypeIgnoreComments = true + +# This is required as the CI pre-commit does not download the module (i.e. numpy, torch, prettytable) +# Therefore, we have to ignore missing imports +reportMissingImports = "none" +# This is required to ignore for type checks of modules with stubs missing. +reportMissingModuleSource = "none" # -> most common: prettytable in mdp managers + +reportGeneralTypeIssues = "none" # -> raises 218 errors (usage of literal MISSING in dataclasses) +reportOptionalMemberAccess = "warning" # -> raises 8 errors +reportPrivateUsage = "warning" diff --git a/scripts/reinforcement_learning/rl_utils.py b/scripts/reinforcement_learning/rl_utils.py new file mode 100644 index 0000000..d344b3c --- /dev/null +++ b/scripts/reinforcement_learning/rl_utils.py @@ -0,0 +1,44 @@ +# Copyright (c) 2025 Deep Robotics +# SPDX-License-Identifier: BSD 3-Clause + +# Copyright (c) 2024-2025 Ziqi Fan +# SPDX-License-Identifier: Apache-2.0 + +import torch + +import isaaclab.utils.math as math_utils + + +def camera_follow(env): + if not hasattr(camera_follow, "smooth_camera_positions"): + camera_follow.smooth_camera_positions = [] + robot_pos = env.unwrapped.scene["robot"].data.root_pos_w[0] + robot_quat = env.unwrapped.scene["robot"].data.root_quat_w[0] + camera_offset = torch.tensor([-3.0, 0.0, 0.5], dtype=torch.float32, device=env.device) + camera_pos = math_utils.transform_points( + camera_offset.unsqueeze(0), pos=robot_pos.unsqueeze(0), quat=robot_quat.unsqueeze(0) + ).squeeze(0) + # camera_pos[2] = torch.clamp(camera_pos[2], min=0.1) + window_size = 50 + camera_follow.smooth_camera_positions.append(camera_pos) + if len(camera_follow.smooth_camera_positions) > window_size: + camera_follow.smooth_camera_positions.pop(0) + smooth_camera_pos = torch.mean(torch.stack(camera_follow.smooth_camera_positions), dim=0) + env.unwrapped.viewport_camera_controller.set_view_env_index(env_index=0) + env.unwrapped.viewport_camera_controller.update_view_location( + eye=smooth_camera_pos.cpu().numpy(), lookat=robot_pos.cpu().numpy() + ) + + + + +import os +import pickle +from typing import Any + +def reindex_feet(vec): + return vec[:, [1, 0, 3, 2]] + +def reindex(vec): + return vec[:, [3, 4, 5, 0, 1, 2, 9, 10, 11, 6, 7, 8]] + diff --git a/scripts/reinforcement_learning/rsl_rl/cli_args.py b/scripts/reinforcement_learning/rsl_rl/cli_args.py new file mode 100644 index 0000000..717037f --- /dev/null +++ b/scripts/reinforcement_learning/rsl_rl/cli_args.py @@ -0,0 +1,97 @@ +# Copyright (c) 2025 Deep Robotics +# SPDX-License-Identifier: BSD 3-Clause + +# Copyright (c) 2024-2025 Ziqi Fan +# SPDX-License-Identifier: Apache-2.0 + +# Copyright (c) 2024-2025, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import argparse +import random +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from isaaclab_rl.rsl_rl import RslRlOnPolicyRunnerCfg + + +def add_rsl_rl_args(parser: argparse.ArgumentParser): + """Add RSL-RL arguments to the parser. + + Args: + parser: The parser to add the arguments to. + """ + # create a new argument group + arg_group = parser.add_argument_group("rsl_rl", description="Arguments for RSL-RL agent.") + # -- experiment arguments + arg_group.add_argument( + "--experiment_name", type=str, default=None, help="Name of the experiment folder where logs will be stored." + ) + arg_group.add_argument("--run_name", type=str, default=None, help="Run name suffix to the log directory.") + # -- load arguments + arg_group.add_argument("--resume", action="store_true", default=False, help="Whether to resume from a checkpoint.") + arg_group.add_argument("--load_run", type=str, default=None, help="Name of the run folder to resume from.") + arg_group.add_argument("--checkpoint", type=str, default=None, help="Checkpoint file to resume from.") + # -- logger arguments + arg_group.add_argument( + "--logger", type=str, default=None, choices={"wandb", "tensorboard", "neptune"}, help="Logger module to use." + ) + arg_group.add_argument( + "--log_project_name", type=str, default=None, help="Name of the logging project when using wandb or neptune." + ) + + +def parse_rsl_rl_cfg(task_name: str, args_cli: argparse.Namespace) -> RslRlOnPolicyRunnerCfg: + """Parse configuration for RSL-RL agent based on inputs. + + Args: + task_name: The name of the environment. + args_cli: The command line arguments. + + Returns: + The parsed configuration for RSL-RL agent based on inputs. + """ + from isaaclab_tasks.utils.parse_cfg import load_cfg_from_registry + + # load the default configuration + rslrl_cfg: RslRlOnPolicyRunnerCfg = load_cfg_from_registry(task_name, "rsl_rl_cfg_entry_point") + rslrl_cfg = update_rsl_rl_cfg(rslrl_cfg, args_cli) + return rslrl_cfg + + +def update_rsl_rl_cfg(agent_cfg: RslRlOnPolicyRunnerCfg, args_cli: argparse.Namespace): + """Update configuration for RSL-RL agent based on inputs. + + Args: + agent_cfg: The configuration for RSL-RL agent. + args_cli: The command line arguments. + + Returns: + The updated configuration for RSL-RL agent based on inputs. + """ + # override the default configuration with CLI arguments + if hasattr(args_cli, "seed") and args_cli.seed is not None: + # randomly sample a seed if seed = -1 + if args_cli.seed == -1: + args_cli.seed = random.randint(0, 10000) + agent_cfg.seed = args_cli.seed + if args_cli.resume is not None: + agent_cfg.resume = args_cli.resume + if args_cli.load_run is not None: + agent_cfg.load_run = args_cli.load_run + if args_cli.checkpoint is not None: + agent_cfg.load_checkpoint = args_cli.checkpoint + if args_cli.run_name is not None: + agent_cfg.run_name = args_cli.run_name + if args_cli.logger is not None: + agent_cfg.logger = args_cli.logger + # set the project name for wandb and neptune + if agent_cfg.logger in {"wandb", "neptune"} and args_cli.log_project_name: + agent_cfg.wandb_project = args_cli.log_project_name + agent_cfg.neptune_project = args_cli.log_project_name + + return agent_cfg diff --git a/scripts/reinforcement_learning/rsl_rl/eval.py b/scripts/reinforcement_learning/rsl_rl/eval.py new file mode 100644 index 0000000..b45a2df --- /dev/null +++ b/scripts/reinforcement_learning/rsl_rl/eval.py @@ -0,0 +1,716 @@ +# Copyright (c) 2025 Deep Robotics +# SPDX-License-Identifier: BSD 3-Clause + +"""Evaluate an RSL-RL checkpoint with criterion metrics migrated from legged_gym.""" + +from __future__ import annotations + +import argparse +import copy +import json +import os +import sys +from datetime import datetime + +import numpy as np + +PROJECT_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "..")) +PROJECT_SOURCE_PATH = os.path.join(PROJECT_ROOT, "source", "rl_training") +if PROJECT_SOURCE_PATH not in sys.path: + sys.path.insert(0, PROJECT_SOURCE_PATH) + +from isaaclab.app import AppLauncher + +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) +import cli_args + +parser = argparse.ArgumentParser(description="Evaluate an RSL-RL checkpoint with criterion metrics.") +parser.add_argument("--num_envs", type=int, default=250, help="Number of parallel environments.") +parser.add_argument("--task", type=str, default="Rough-Deeprobotics-M20-v0", help="Name of the task.") +parser.add_argument( + "--agent", type=str, default="rsl_rl_cfg_entry_point", help="Name of the RL agent configuration entry point." +) +parser.add_argument("--seed", type=int, default=None, help="Seed used for the environment.") +parser.add_argument( + "--eval_type", + type=str, + default="all", + choices=[ + "all", + "upper_limits", + "velocity_error", + "access_rate", + "time_consumption", + "energy_consumption", + "motion_stability", + "stability_margin", + "robustness", + ], +) +parser.add_argument("--episode_length_s", type=float, default=10.0) +parser.add_argument("--command_x", type=float, default=1.0) +parser.add_argument("--command_y", type=float, default=0.0) +parser.add_argument("--command_yaw", type=float, default=0.0) +parser.add_argument("--terrain_size", type=float, default=2.0, help="Evaluation terrain tile size in meters.") +parser.add_argument("--success_distance", type=float, default=2.0) +parser.add_argument("--start_x", type=float, default=-0.9, help="Evaluation reset x position on the terrain tile.") +parser.add_argument("--start_y", type=float, default=0.0, help="Evaluation reset y position on the terrain tile.") +parser.add_argument("--num_levels", type=int, default=8, help="Number of difficulty levels for upper_limits.") +parser.add_argument("--envs_per_cell", type=int, default=8, help="Number of robots per terrain/difficulty cell.") +parser.add_argument( + "--terrain_type", + type=str, + default="all", + choices=["all", "stairs_up", "stairs_down", "smooth_slope", "rough_slope", "boxes"], + help="Terrain type used by upper_limits.", +) +parser.add_argument("--tilt_threshold", type=float, default=0.70, help="Tilt threshold in radians for success/stability.") +parser.add_argument("--acc_threshold", type=float, default=2.0, help="Base acceleration threshold for instability.") +parser.add_argument("--robot_mass", type=float, default=25.482, help="Robot mass used for COT.") +parser.add_argument("--output", type=str, default=None, help="Optional output JSON path.") +parser.add_argument("--print_json", action="store_true", default=False, help="Print the saved JSON payload to stdout.") +parser.add_argument("--verbose", action="store_true", default=False, help="Print per-terrain per-level debug lines.") +parser.add_argument( + "--include_summary", + action="store_true", + default=False, + help="When eval_type=all, also run a second mixed-terrain rollout for summary metrics.", +) +cli_args.add_rsl_rl_args(parser) +AppLauncher.add_app_launcher_args(parser) +args_cli, hydra_args = parser.parse_known_args() + +sys.argv = [sys.argv[0]] + hydra_args + +app_launcher = AppLauncher(args_cli) +simulation_app = app_launcher.app + +import gymnasium as gym +import torch +from rsl_rl.runners import OnPolicyRunner + +from isaaclab.envs import ( + DirectMARLEnv, + DirectMARLEnvCfg, + DirectRLEnvCfg, + ManagerBasedRLEnvCfg, + multi_agent_to_single_agent, +) +from isaaclab.terrains.config.rough import ROUGH_TERRAINS_CFG +from isaaclab_rl.rsl_rl import RslRlOnPolicyRunnerCfg, RslRlVecEnvWrapper +from isaaclab_tasks.utils import get_checkpoint_path +from isaaclab_tasks.utils.hydra import hydra_task_config + +import rl_training +import rl_training.tasks # noqa: F401 + +EVAL_TYPES = [ + "upper_limits", + "velocity_error", + "access_rate", + "time_consumption", + "energy_consumption", + "motion_stability", + "stability_margin", + "robustness", +] + +TERRAIN_LEVELS = { + "stairs_up": [0.04, 0.06, 0.08, 0.10, 0.12, 0.14, 0.16, 0.18], + "stairs_down": [0.04, 0.06, 0.08, 0.10, 0.12, 0.14, 0.16, 0.18], + "smooth_slope": [0.05, 0.10, 0.15, 0.20, 0.25, 0.30, 0.35, 0.40], + "rough_slope": [0.05, 0.10, 0.15, 0.20, 0.25, 0.30, 0.35, 0.40], + "boxes": [0.04, 0.06, 0.08, 0.10, 0.12, 0.14, 0.16, 0.18], +} + +TERRAIN_TO_SUBTERRAIN = { + "stairs_up": "pyramid_stairs", + "stairs_down": "pyramid_stairs_inv", + "smooth_slope": "hf_pyramid_slope", + "rough_slope": "random_rough", + "boxes": "boxes", +} + + +def apply_eval_terrain_baseline(env_cfg) -> None: + """Use the same evaluation terrain generator for every task.""" + if env_cfg.scene.terrain.terrain_type == "plane": + env_cfg.scene.terrain.terrain_type = "generator" + env_cfg.scene.terrain.terrain_generator = copy.deepcopy(ROUGH_TERRAINS_CFG) + env_cfg.scene.terrain.max_init_terrain_level = None + env_cfg.scene.terrain.terrain_generator.size = (args_cli.terrain_size, args_cli.terrain_size) + env_cfg.scene.terrain.terrain_generator.border_width = 2.0 + env_cfg.scene.terrain.terrain_generator.curriculum = False + env_cfg.scene.terrain.terrain_generator.num_rows = 5 + env_cfg.scene.terrain.terrain_generator.num_cols = 5 + env_cfg.scene.terrain.terrain_generator.use_cache = False + for sub_cfg in env_cfg.scene.terrain.terrain_generator.sub_terrains.values(): + if hasattr(sub_cfg, "platform_width"): + sub_cfg.platform_width = min(0.4, args_cli.terrain_size * 0.2) + if hasattr(sub_cfg, "border_width"): + sub_cfg.border_width = min(0.1, args_cli.terrain_size * 0.05) + if hasattr(sub_cfg, "step_width"): + sub_cfg.step_width = min(0.25, args_cli.terrain_size * 0.125) + if hasattr(sub_cfg, "grid_width"): + sub_cfg.grid_width = min(0.3, args_cli.terrain_size * 0.15) + + +def resolve_checkpoint(log_root_path: str, load_run: str | None, checkpoint: str | None) -> str: + if checkpoint is not None: + checkpoint_path = os.path.expanduser(checkpoint) + if os.path.isabs(checkpoint_path) or os.path.dirname(checkpoint_path): + checkpoint_path = os.path.abspath(checkpoint_path) + if not os.path.isfile(checkpoint_path): + raise FileNotFoundError(f"Unable to find checkpoint: {checkpoint_path}") + return checkpoint_path + + if load_run is not None: + run_path = os.path.expanduser(load_run) + if os.path.isabs(run_path) or os.path.dirname(run_path): + run_path = os.path.abspath(run_path) + if not os.path.isdir(run_path): + raise FileNotFoundError(f"Unable to find run directory: {run_path}") + return get_checkpoint_path(os.path.dirname(run_path), os.path.basename(run_path), checkpoint) + + return get_checkpoint_path(log_root_path, load_run, checkpoint) + + +def set_fixed_command(env_cfg, command_x: float, command_y: float, command_yaw: float, episode_length_s: float) -> None: + command_cfg = env_cfg.commands.base_velocity + command_cfg.heading_command = False + command_cfg.resampling_time_range = (episode_length_s + 0.01, episode_length_s + 0.01) + command_cfg.ranges.lin_vel_x = (command_x, command_x) + command_cfg.ranges.lin_vel_y = (command_y, command_y) + command_cfg.ranges.ang_vel_z = (command_yaw, command_yaw) + + +def configure_eval_env(env_cfg, eval_type: str) -> None: + env_cfg.scene.num_envs = args_cli.num_envs + env_cfg.episode_length_s = args_cli.episode_length_s + env_cfg.seed = args_cli.seed if args_cli.seed is not None else 0 + env_cfg.sim.device = args_cli.device if args_cli.device is not None else env_cfg.sim.device + env_cfg.observations.policy.enable_corruption = False + set_fixed_command(env_cfg, args_cli.command_x, args_cli.command_y, args_cli.command_yaw, args_cli.episode_length_s) + + apply_eval_terrain_baseline(env_cfg) + + env_cfg.curriculum.terrain_levels = None + env_cfg.curriculum.command_levels = None + + for name in ( + "randomize_rigid_body_material", + "randomize_wheel_material", + "randomize_support_material", + "randomize_rigid_body_mass", + "randomize_rigid_body_mass_base", + "randomize_rigid_body_inertia", + "randomize_com_positions", + "randomize_apply_external_force_torque", + "randomize_actuator_gains", + ): + if hasattr(env_cfg.events, name): + setattr(env_cfg.events, name, None) + + if hasattr(env_cfg.events, "randomize_reset_base") and env_cfg.events.randomize_reset_base is not None: + env_cfg.events.randomize_reset_base.params = { + "pose_range": { + "x": (args_cli.start_x, args_cli.start_x), + "y": (args_cli.start_y, args_cli.start_y), + "z": (0.0, 0.0), + "roll": (0.0, 0.0), + "pitch": (0.0, 0.0), + "yaw": (0.0, 0.0), + }, + "velocity_range": { + "x": (0.0, 0.0), + "y": (0.0, 0.0), + "z": (0.0, 0.0), + "roll": (0.0, 0.0), + "pitch": (0.0, 0.0), + "yaw": (0.0, 0.0), + }, + } + + if hasattr(env_cfg.events, "randomize_push_robot"): + env_cfg.events.randomize_push_robot = None + if eval_type == "robustness" and hasattr(env_cfg.events, "randomize_push_robot"): + # Keep robustness deterministic enough to compare runs, but enable interval pushes. + from isaaclab.managers import EventTermCfg as EventTerm + import rl_training.tasks.manager_based.locomotion.velocity.mdp as mdp + + env_cfg.events.randomize_push_robot = EventTerm( + func=mdp.push_by_setting_velocity, + mode="interval", + interval_range_s=(2.0, 2.01), + params={"velocity_range": {"x": (-0.5, 0.5), "y": (-0.5, 0.5)}}, + ) + + +def configure_upper_limit_env(env_cfg, terrain_types: list[str]) -> None: + """Configure one fixed terrain grid: rows are difficulty levels, columns are terrain types.""" + configure_eval_env(env_cfg, "all") + terrain_generator = env_cfg.scene.terrain.terrain_generator + terrain_generator.curriculum = True + terrain_generator.num_rows = args_cli.num_levels + terrain_generator.num_cols = len(terrain_types) + terrain_generator.difficulty_range = (0.0, 1.0) + if hasattr(terrain_generator, "seed"): + terrain_generator.seed = args_cli.seed if args_cli.seed is not None else 0 + env_cfg.scene.terrain.max_init_terrain_level = args_cli.num_levels - 1 + env_cfg.scene.num_envs = args_cli.num_levels * len(terrain_types) * args_cli.envs_per_cell + + selected_subterrains = [TERRAIN_TO_SUBTERRAIN[terrain_type] for terrain_type in terrain_types] + original_subterrains = terrain_generator.sub_terrains + terrain_generator.sub_terrains = { + subterrain_name: copy.deepcopy(original_subterrains[subterrain_name]) for subterrain_name in selected_subterrains + } + for name, sub_cfg in terrain_generator.sub_terrains.items(): + sub_cfg.proportion = selected_subterrains.count(name) / len(selected_subterrains) + + for terrain_type in terrain_types: + sub_cfg = terrain_generator.sub_terrains[TERRAIN_TO_SUBTERRAIN[terrain_type]] + levels = TERRAIN_LEVELS[terrain_type][: args_cli.num_levels] + if terrain_type in ("stairs_up", "stairs_down"): + sub_cfg.step_height_range = (levels[0], levels[-1]) + elif terrain_type == "smooth_slope": + sub_cfg.slope_range = (levels[0], levels[-1]) + elif terrain_type == "rough_slope": + sub_cfg.noise_range = (levels[0], levels[-1]) + sub_cfg.noise_step = 0.01 + elif terrain_type == "boxes": + sub_cfg.grid_height_range = (levels[0], levels[-1]) + + +def assign_upper_limit_origins(env, terrain_types: list[str]) -> tuple[torch.Tensor, torch.Tensor]: + """Assign an equal number of robots to every terrain/difficulty cell.""" + base_env = env.unwrapped + terrain = base_env.scene.terrain + device = base_env.device + num_levels = args_cli.num_levels + num_types = len(terrain_types) + levels = [] + types = [] + for terrain_type_id in range(num_types): + for level_id in range(num_levels): + levels.extend([level_id] * args_cli.envs_per_cell) + types.extend([terrain_type_id] * args_cli.envs_per_cell) + levels_tensor = torch.tensor(levels, dtype=torch.long, device=device) + types_tensor = torch.tensor(types, dtype=torch.long, device=device) + terrain.terrain_levels = levels_tensor + terrain.terrain_types = types_tensor + terrain.env_origins = terrain.terrain_origins[levels_tensor, types_tensor] + env.reset() + return levels_tensor, types_tensor + + +def tilt_from_projected_gravity(projected_gravity_b: torch.Tensor) -> torch.Tensor: + z_axis = torch.clamp(-projected_gravity_b[:, 2], min=-1.0, max=1.0) + return torch.acos(z_axis) + + +def support_margin_xy(root_xy: torch.Tensor, foot_xy: torch.Tensor) -> torch.Tensor: + min_xy = torch.min(foot_xy, dim=1).values + max_xy = torch.max(foot_xy, dim=1).values + inside_x = (root_xy[:, 0] >= min_xy[:, 0]) & (root_xy[:, 0] <= max_xy[:, 0]) + inside_y = (root_xy[:, 1] >= min_xy[:, 1]) & (root_xy[:, 1] <= max_xy[:, 1]) + inside = inside_x & inside_y + inside_dist = torch.minimum( + torch.minimum(root_xy[:, 0] - min_xy[:, 0], max_xy[:, 0] - root_xy[:, 0]), + torch.minimum(root_xy[:, 1] - min_xy[:, 1], max_xy[:, 1] - root_xy[:, 1]), + ) + outside_dx = torch.clamp(torch.maximum(min_xy[:, 0] - root_xy[:, 0], root_xy[:, 0] - max_xy[:, 0]), min=0.0) + outside_dy = torch.clamp(torch.maximum(min_xy[:, 1] - root_xy[:, 1], root_xy[:, 1] - max_xy[:, 1]), min=0.0) + outside_dist = torch.sqrt(outside_dx**2 + outside_dy**2) + return torch.where(inside, inside_dist, -outside_dist) + + +def run_rollout(env, policy, include_per_env: bool = False) -> dict[str, object]: + base_env = env.unwrapped + robot = base_env.scene["robot"] + dt = base_env.step_dt + device = base_env.device + max_steps = int(base_env.max_episode_length) + command = torch.tensor([args_cli.command_x, args_cli.command_y, args_cli.command_yaw], device=device) + command_xy = command[:2].unsqueeze(0) + command_yaw = command[2] + command_dir = torch.tensor([args_cli.command_x, args_cli.command_y], device=device) + if torch.linalg.norm(command_dir) < 1e-6: + command_dir = torch.tensor([1.0, 0.0], device=device) + command_dir = command_dir / torch.linalg.norm(command_dir) + + obs = env.get_observations() + start_pos = robot.data.root_pos_w[:, :3].clone() + last_lin_vel = robot.data.root_lin_vel_w[:, :3].clone() + complete_time = torch.full((base_env.num_envs,), float("nan"), device=device) + energy = torch.zeros(base_env.num_envs, device=device) + route = torch.zeros(base_env.num_envs, device=device) + lin_error_sum = torch.zeros(base_env.num_envs, device=device) + yaw_error_sum = torch.zeros(base_env.num_envs, device=device) + instability_count = torch.zeros(base_env.num_envs, device=device) + max_tilt = torch.zeros(base_env.num_envs, device=device) + stable_time = torch.zeros(base_env.num_envs, device=device) + margin_sum = torch.zeros(base_env.num_envs, device=device) + margin_count = 0 + max_progress = torch.zeros(base_env.num_envs, device=device) + reached = torch.zeros(base_env.num_envs, dtype=torch.bool, device=device) + ever_fallen = torch.zeros(base_env.num_envs, dtype=torch.bool, device=device) + support_steps = torch.zeros(base_env.num_envs, device=device) + support_events = torch.zeros(base_env.num_envs, device=device) + max_support_time = torch.zeros(base_env.num_envs, device=device) + previous_support_active = torch.zeros(base_env.num_envs, dtype=torch.bool, device=device) + + wheel_body_ids, _ = robot.find_bodies(".*_wheel") + contact_sensor = base_env.scene.sensors["contact_forces"] + support_body_ids, _ = contact_sensor.find_bodies(["base_link", ".*_knee"]) + + for step in range(max_steps): + pre_progress = torch.sum((robot.data.root_pos_w[:, :2] - start_pos[:, :2]) * command_dir.unsqueeze(0), dim=1) + max_progress = torch.maximum(max_progress, pre_progress) + reached |= pre_progress >= args_cli.success_distance + + with torch.inference_mode(): + actions = policy(obs) + if isinstance(actions, tuple): + actions = actions[0] + obs, _, dones, _ = env.step(actions) + + lin_vel_b = robot.data.root_lin_vel_b[:, :2] + yaw_vel_b = robot.data.root_ang_vel_b[:, 2] + lin_error_sum += torch.linalg.norm(lin_vel_b - command_xy, dim=1) + yaw_error_sum += torch.abs(yaw_vel_b - command_yaw) + + joint_power = torch.abs(robot.data.applied_torque * robot.data.joint_vel) + energy += torch.sum(joint_power, dim=1) * dt + route += torch.linalg.norm(robot.data.root_lin_vel_w[:, :2], dim=1) * dt + + acc = torch.linalg.norm(robot.data.root_lin_vel_w[:, :3] - last_lin_vel, dim=1) / dt + instability_count += (acc > args_cli.acc_threshold).float() + last_lin_vel = robot.data.root_lin_vel_w[:, :3].clone() + + tilt = tilt_from_projected_gravity(robot.data.projected_gravity_b) + max_tilt = torch.maximum(max_tilt, tilt) + stable_time += (tilt < args_cli.tilt_threshold).float() * dt + + progress = torch.sum((robot.data.root_pos_w[:, :2] - start_pos[:, :2]) * command_dir.unsqueeze(0), dim=1) + max_progress = torch.maximum(max_progress, progress) + reached |= progress >= args_cli.success_distance + newly_success = torch.isnan(complete_time) & (progress >= args_cli.success_distance) + complete_time[newly_success] = (step + 1) * dt + ever_fallen |= tilt > args_cli.tilt_threshold + + support_force = torch.linalg.norm( + contact_sensor.data.net_forces_w[:, support_body_ids, :], dim=-1 + ) + support_time = contact_sensor.data.current_contact_time[:, support_body_ids] + support_active = torch.any(support_force > 15.0, dim=1) + support_steps += support_active.float() + support_events += (support_active & ~previous_support_active).float() + max_support_time = torch.maximum(max_support_time, torch.amax(support_time, dim=1)) + previous_support_active = support_active + + if len(wheel_body_ids) >= 3: + foot_xy = robot.data.body_pos_w[:, wheel_body_ids, :2] + margin_sum += support_margin_xy(robot.data.root_pos_w[:, :2], foot_xy) + margin_count += 1 + + tilt = tilt_from_projected_gravity(robot.data.projected_gravity_b) + success = reached & ~ever_fallen + success_count = torch.sum(success).item() + success_den = max(base_env.num_envs, 1) + successful_time = complete_time[success & torch.isfinite(complete_time)] + avg_time = torch.mean(successful_time).item() if successful_time.numel() > 0 else float("nan") + + cot = energy / ((route + 1e-6) * args_cli.robot_mass * 9.81) + metrics = { + "num_envs": float(base_env.num_envs), + "episode_length_s": float(args_cli.episode_length_s), + "terrain_size_m": float(args_cli.terrain_size), + "success_distance_m": float(args_cli.success_distance), + "success_count": float(success_count), + "success_rate": float(success_count / success_den), + "avg_final_progress_m": float(torch.mean(max_progress).item()), + "avg_time_consumption_s": float(avg_time), + "time_efficiency": float(avg_time / (args_cli.success_distance / max(abs(args_cli.command_x), 1e-6))) + if successful_time.numel() > 0 + else float("nan"), + "mean_lin_vel_error": float(torch.mean(lin_error_sum / max_steps).item()), + "mean_yaw_vel_error": float(torch.mean(yaw_error_sum / max_steps).item()), + "avg_energy_j": float(torch.mean(energy).item()), + "avg_cost_of_transport": float(torch.mean(cot).item()), + "unstability_rate": float(torch.mean(instability_count / max_steps).item()), + "avg_stability_margin_m": float(torch.mean(margin_sum / max(margin_count, 1)).item()), + "avg_max_tilt_rad": float(torch.mean(max_tilt).item()), + "avg_stable_time_s": float(torch.mean(stable_time).item()), + "avg_support_contact_ratio": float(torch.mean(support_steps / max_steps).item()), + "avg_support_events": float(torch.mean(support_events).item()), + "avg_max_support_contact_time_s": float(torch.mean(max_support_time).item()), + "successful_runs_with_support_ratio": float( + torch.mean((support_events[success] > 0.0).float()).item() + ) + if success_count > 0 + else float("nan"), + } + if include_per_env: + metrics["_per_env"] = { + "success": success.detach().cpu().tolist(), + "max_progress": max_progress.detach().cpu().tolist(), + "max_tilt": max_tilt.detach().cpu().tolist(), + } + return metrics + + +def make_env_and_policy(env_cfg, agent_cfg, resume_path: str): + env = gym.make(args_cli.task, cfg=env_cfg) + if isinstance(env.unwrapped, DirectMARLEnv): + env = multi_agent_to_single_agent(env) + env = RslRlVecEnvWrapper(env, clip_actions=agent_cfg.clip_actions) + + runner = OnPolicyRunner(env, agent_cfg.to_dict(), log_dir=None, device=agent_cfg.device) + runner.load(resume_path) + policy = runner.get_inference_policy(device=env.unwrapped.device) + return env, policy + + +def print_upper_limit_matrix(title: str, terrain_types: list[str], matrix: list[list[float | None]]) -> None: + """Print an upper-limit success matrix with levels as rows and terrain types as columns.""" + col_width = 14 + print(f"\n{title}") + print("level".ljust(8) + "".join(terrain_type[: col_width - 1].rjust(col_width) for terrain_type in terrain_types)) + for level_index, row in enumerate(matrix, start=1): + cells = [] + for value in row: + cells.append(" -".rjust(col_width) if value is None else f"{value:.2f}".rjust(col_width)) + print(f"{level_index:<8}" + "".join(cells)) + + +def matrix_to_numpy(matrix: list[list[float | None]]) -> np.ndarray: + return np.array([[np.nan if value is None else value for value in row] for row in matrix], dtype=np.float32) + + +def print_numpy_matrix(title: str, matrix: list[list[float | None]]) -> None: + print(f"{title}:") + print(np.array2string(matrix_to_numpy(matrix).round(2), precision=2, suppress_small=True)) + + +def print_upper_limit_summary(result: dict[str, object], terrain_types: list[str]) -> None: + matrix = result["_matrix"] + print_numpy_matrix("success rate", matrix["success_rate"]) + print_numpy_matrix("success rate accumulated", matrix["success_rate_accumulated"]) + score = np.array([result[terrain_type]["score"] for terrain_type in terrain_types], dtype=np.float32) + print(f"score: {np.array2string(score.round(2), precision=2, suppress_small=True)} of {len(terrain_types)}") + + +def print_rollout_summary(metrics: dict[str, object]) -> None: + success_rate = float(metrics["success_rate"]) + success_count = int(round(float(metrics["success_count"]))) + avg_time = float(metrics["avg_time_consumption_s"]) + avg_time_text = "nan" if np.isnan(avg_time) else f"{avg_time:.2f}" + print(f"success rate: {success_rate:.2f}") + print(f"success count: {success_count} time consumption: {avg_time_text}") + print(f"average power consumption: {float(metrics['avg_energy_j']):.2f}") + print(f"average cost of transport: {float(metrics['avg_cost_of_transport']):.2f}") + print(f"unstability rate: {float(metrics['unstability_rate']) * 100:.1f}%") + print(f"average stability margin: {float(metrics['avg_stability_margin_m']):.2f}") + print( + f"average max tilt: {float(metrics['avg_max_tilt_rad']):.2f} " + f"average max time: {float(metrics['avg_stable_time_s']):.2f}" + ) + support_success_ratio = float(metrics["successful_runs_with_support_ratio"]) + support_success_text = "nan" if np.isnan(support_success_ratio) else f"{support_success_ratio * 100:.1f}%" + print( + f"support contact ratio: {float(metrics['avg_support_contact_ratio']) * 100:.1f}% " + f"events/episode: {float(metrics['avg_support_events']):.2f} " + f"max duration: {float(metrics['avg_max_support_contact_time_s']):.2f}s " + f"successful runs using support: {support_success_text}" + ) + + +def print_metric_result(eval_type: str, metrics: dict[str, object]) -> None: + if eval_type == "velocity_error": + print(f"mean lin vel error: {float(metrics['mean_lin_vel_error']):.2f}") + print(f"mean yaw vel error: {float(metrics['mean_yaw_vel_error']):.2f}") + elif eval_type == "access_rate": + print(f"success rate: {float(metrics['success_rate']):.2f}") + elif eval_type == "time_consumption": + success_count = int(round(float(metrics["success_count"]))) + avg_time = float(metrics["avg_time_consumption_s"]) + avg_time_text = "nan" if np.isnan(avg_time) else f"{avg_time:.2f}" + print(f"success count: {success_count} time consumption: {avg_time_text}") + elif eval_type == "energy_consumption": + print(f"average power consumption: {float(metrics['avg_energy_j']):.2f}") + print(f"average cost of transport: {float(metrics['avg_cost_of_transport']):.2f}") + elif eval_type == "motion_stability": + print(f"unstability rate: {float(metrics['unstability_rate']) * 100:.1f}%") + elif eval_type == "stability_margin": + print(f"average stability margin: {float(metrics['avg_stability_margin_m']):.2f}") + elif eval_type == "robustness": + print( + f"average max tilt: {float(metrics['avg_max_tilt_rad']):.2f} " + f"average max time: {float(metrics['avg_stable_time_s']):.2f}" + ) + + +def run_upper_limits(env_cfg, agent_cfg, resume_path: str) -> dict[str, object]: + terrain_types = list(TERRAIN_LEVELS) if args_cli.terrain_type == "all" else [args_cli.terrain_type] + eval_cfg = copy.deepcopy(env_cfg) + configure_upper_limit_env(eval_cfg, terrain_types) + env, policy = make_env_and_policy(eval_cfg, agent_cfg, resume_path) + level_ids, type_ids = assign_upper_limit_origins(env, terrain_types) + metrics = run_rollout(env, policy, include_per_env=True) + env.close() + + per_env = metrics.pop("_per_env") + success = torch.tensor(per_env["success"], dtype=torch.float32) + max_progress = torch.tensor(per_env["max_progress"], dtype=torch.float32) + max_tilt = torch.tensor(per_env["max_tilt"], dtype=torch.float32) + level_ids_cpu = level_ids.cpu() + type_ids_cpu = type_ids.cpu() + + result: dict[str, object] = {} + for terrain_type_id, terrain_type in enumerate(terrain_types): + level_values = TERRAIN_LEVELS[terrain_type][: args_cli.num_levels] + rates = [] + level_metrics = [] + for level_index, level_value in enumerate(level_values): + mask = (type_ids_cpu == terrain_type_id) & (level_ids_cpu == level_index) + success_count = float(torch.sum(success[mask]).item()) + total_count = int(torch.sum(mask).item()) + success_rate = success_count / max(total_count, 1) + avg_progress = float(torch.mean(max_progress[mask]).item()) if total_count > 0 else float("nan") + avg_tilt = float(torch.mean(max_tilt[mask]).item()) if total_count > 0 else float("nan") + rates.append(success_rate) + level_metrics.append( + { + "level": level_index + 1, + "difficulty": level_value, + "success_rate": success_rate, + "success_count": success_count, + "num_envs": total_count, + "avg_progress_m": avg_progress, + "avg_max_tilt_rad": avg_tilt, + } + ) + if args_cli.verbose: + print( + f"[EVAL] {terrain_type} level={level_index + 1} difficulty={level_value:.3f} " + f"success_rate={success_rate:.3f}" + ) + accumulated = [] + running = 1.0 + for rate in rates: + running *= rate + accumulated.append(running) + score = 0.0 + for idx, acc_rate in enumerate(accumulated): + next_fail = 1.0 - rates[idx + 1] if idx + 1 < len(rates) else 1.0 + score += acc_rate * (idx + 1) * next_fail + result[terrain_type] = { + "levels": level_metrics, + "success_rate": rates, + "success_rate_accumulated": accumulated, + "score": score, + "max_level": len(level_values), + "benchmark": "fixed_rough_terrain_v1", + "terrain_size_m": args_cli.terrain_size, + "success_distance_m": args_cli.success_distance, + "episode_length_s": args_cli.episode_length_s, + } + max_levels = max((len(result[terrain_type]["success_rate"]) for terrain_type in terrain_types), default=0) + success_matrix: list[list[float | None]] = [] + accumulated_matrix: list[list[float | None]] = [] + difficulty_matrix: list[list[float | None]] = [] + for level_index in range(max_levels): + success_row = [] + accumulated_row = [] + difficulty_row = [] + for terrain_type in terrain_types: + rates = result[terrain_type]["success_rate"] + accumulated = result[terrain_type]["success_rate_accumulated"] + levels = result[terrain_type]["levels"] + success_row.append(rates[level_index] if level_index < len(rates) else None) + accumulated_row.append(accumulated[level_index] if level_index < len(accumulated) else None) + difficulty_row.append(levels[level_index]["difficulty"] if level_index < len(levels) else None) + success_matrix.append(success_row) + accumulated_matrix.append(accumulated_row) + difficulty_matrix.append(difficulty_row) + + result["_matrix"] = { + "terrain_types": terrain_types, + "success_rate": success_matrix, + "success_rate_accumulated": accumulated_matrix, + "difficulty": difficulty_matrix, + "terrain_size_m": args_cli.terrain_size, + "success_distance_m": args_cli.success_distance, + "episode_length_s": args_cli.episode_length_s, + } + result["_summary"] = metrics + print_upper_limit_summary(result, terrain_types) + return result + + +def filter_metrics(metrics: dict[str, float], eval_type: str) -> dict[str, float]: + if eval_type == "all": + return metrics + keys = { + "velocity_error": ("mean_lin_vel_error", "mean_yaw_vel_error"), + "access_rate": ("success_count", "success_rate", "avg_final_progress_m"), + "time_consumption": ("success_count", "avg_time_consumption_s", "time_efficiency"), + "energy_consumption": ("avg_energy_j", "avg_cost_of_transport"), + "motion_stability": ("unstability_rate",), + "stability_margin": ("avg_stability_margin_m",), + "robustness": ("success_rate", "avg_max_tilt_rad", "avg_stable_time_s"), + }[eval_type] + return {key: metrics[key] for key in keys} + + +@hydra_task_config(args_cli.task, args_cli.agent) +def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agent_cfg: RslRlOnPolicyRunnerCfg): + agent_cfg = cli_args.update_rsl_rl_cfg(agent_cfg, args_cli) + + log_root_path = os.path.abspath(os.path.join("logs", "rsl_rl", agent_cfg.experiment_name)) + resume_path = resolve_checkpoint(log_root_path, agent_cfg.load_run, agent_cfg.load_checkpoint) + print(f"[INFO] Evaluating checkpoint: {resume_path}") + + if args_cli.eval_type == "upper_limits": + selected = run_upper_limits(env_cfg, agent_cfg, resume_path) + elif args_cli.eval_type == "all": + upper_limits = run_upper_limits(env_cfg, agent_cfg, resume_path) + selected = {"upper_limits": upper_limits} + print_rollout_summary(upper_limits["_summary"]) + if args_cli.include_summary: + summary_cfg = copy.deepcopy(env_cfg) + configure_eval_env(summary_cfg, "all") + env, policy = make_env_and_policy(summary_cfg, agent_cfg, resume_path) + summary = run_rollout(env, policy) + env.close() + selected["summary"] = summary + else: + eval_type_for_env = "robustness" if args_cli.eval_type == "robustness" else "all" + configure_eval_env(env_cfg, eval_type_for_env) + env, policy = make_env_and_policy(env_cfg, agent_cfg, resume_path) + metrics = run_rollout(env, policy) + selected = filter_metrics(metrics, args_cli.eval_type) + env.close() + print_metric_result(args_cli.eval_type, selected) + if args_cli.print_json: + print(json.dumps(selected, indent=2, sort_keys=True)) + + output = args_cli.output + if output is None: + eval_dir = os.path.join(os.path.dirname(resume_path), "eval") + os.makedirs(eval_dir, exist_ok=True) + suffix = datetime.now().strftime("%Y-%m-%d_%H-%M-%S") + output = os.path.join(eval_dir, f"{args_cli.eval_type}_{suffix}.json") + else: + output = os.path.abspath(os.path.expanduser(output)) + output_dir = os.path.dirname(output) + if output_dir: + os.makedirs(output_dir, exist_ok=True) + with open(output, "w", encoding="utf-8") as f: + json.dump({"eval_type": args_cli.eval_type, "checkpoint": resume_path, "metrics": selected}, f, indent=2) + print(f"[INFO] Saved evaluation result to: {output}") + + +if __name__ == "__main__": + main() + simulation_app.close() diff --git a/scripts/reinforcement_learning/rsl_rl/eval.sh b/scripts/reinforcement_learning/rsl_rl/eval.sh new file mode 100755 index 0000000..cf9730b --- /dev/null +++ b/scripts/reinforcement_learning/rsl_rl/eval.sh @@ -0,0 +1,23 @@ +#!/usr/bin/env bash +set -euo pipefail + +TASK=${TASK:-Rough-Deeprobotics-M20-v0} +CHECKPOINT=${1:-${CHECKPOINT:-}} +NUM_ENVS=${NUM_ENVS:-250} +DEVICE=${DEVICE:-cuda:0} +EVAL_TYPE=${EVAL_TYPE:-all} + +if [[ -z "${CHECKPOINT}" ]]; then + echo "Usage: $0 " + echo "Or set CHECKPOINT=/path/to/model.pt" + exit 1 +fi + +echo "Running eval_type=${EVAL_TYPE}" +python scripts/reinforcement_learning/rsl_rl/eval.py \ + --task="${TASK}" \ + --checkpoint="${CHECKPOINT}" \ + --num_envs="${NUM_ENVS}" \ + --device="${DEVICE}" \ + --eval_type="${EVAL_TYPE}" \ + --headless diff --git a/scripts/reinforcement_learning/rsl_rl/play.py b/scripts/reinforcement_learning/rsl_rl/play.py new file mode 100644 index 0000000..7270610 --- /dev/null +++ b/scripts/reinforcement_learning/rsl_rl/play.py @@ -0,0 +1,280 @@ +# Copyright (c) 2025 Deep Robotics +# SPDX-License-Identifier: BSD 3-Clause + +# Copyright (c) 2024-2025 Ziqi Fan +# SPDX-License-Identifier: Apache-2.0 + +# Copyright (c) 2024-2025, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Script to play a checkpoint if an RL agent from RSL-RL.""" + +"""Launch Isaac Sim Simulator first.""" + +import argparse +import os +import sys + +PROJECT_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "..")) +PROJECT_SOURCE_PATH = os.path.join(PROJECT_ROOT, "source", "rl_training") +if PROJECT_SOURCE_PATH not in sys.path: + sys.path.insert(0, PROJECT_SOURCE_PATH) + +from isaaclab.app import AppLauncher + +# local imports +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) +import cli_args +from rl_utils import camera_follow + +# add argparse arguments +parser = argparse.ArgumentParser(description="Train an RL agent with RSL-RL.") +parser.add_argument("--video", action="store_true", default=False, help="Record videos during training.") +parser.add_argument("--video_length", type=int, default=200, help="Length of the recorded video (in steps).") +parser.add_argument( + "--disable_fabric", action="store_true", default=False, help="Disable fabric and use USD I/O operations." +) +parser.add_argument("--num_envs", type=int, default=None, help="Number of environments to simulate.") +parser.add_argument("--task", type=str, default=None, help="Name of the task.") +parser.add_argument( + "--agent", type=str, default="rsl_rl_cfg_entry_point", help="Name of the RL agent configuration entry point." +) +parser.add_argument("--seed", type=int, default=None, help="Seed used for the environment") +parser.add_argument( + "--use_pretrained_checkpoint", + action="store_true", + help="Use the pre-trained checkpoint from Nucleus.", +) +parser.add_argument("--real-time", action="store_true", default=False, help="Run in real-time, if possible.") +parser.add_argument("--keyboard", action="store_true", default=False, help="Whether to use keyboard.") +parser.add_argument( + "--visual-only", + action="store_true", + default=False, + help="Display the environment with zero actions without loading an RL checkpoint.", +) +# append RSL-RL cli arguments +cli_args.add_rsl_rl_args(parser) +# append AppLauncher cli args +AppLauncher.add_app_launcher_args(parser) +# parse the arguments +args_cli, hydra_args = parser.parse_known_args() +# always enable cameras to record video +if args_cli.video: + args_cli.enable_cameras = True + +# clear out sys.argv for Hydra +sys.argv = [sys.argv[0]] + hydra_args + +# launch omniverse app +app_launcher = AppLauncher(args_cli) +simulation_app = app_launcher.app + +"""Rest everything follows.""" + +import gymnasium as gym +import time +import torch + +from rsl_rl.runners import OnPolicyRunner + +from isaaclab.devices import Se2Keyboard, Se2KeyboardCfg +from isaaclab.envs import ( + DirectMARLEnv, + DirectMARLEnvCfg, + DirectRLEnvCfg, + ManagerBasedRLEnvCfg, + multi_agent_to_single_agent, +) +from isaaclab.managers import ObservationTermCfg as ObsTerm +from isaaclab.utils.assets import retrieve_file_path +from isaaclab.utils.dict import print_dict +from isaaclab.utils.pretrained_checkpoint import get_published_pretrained_checkpoint +from isaaclab_rl.rsl_rl import RslRlOnPolicyRunnerCfg, RslRlVecEnvWrapper, export_policy_as_jit, export_policy_as_onnx +from isaaclab_tasks.utils import get_checkpoint_path +from isaaclab_tasks.utils.hydra import hydra_task_config + +import rl_training +import rl_training.tasks # noqa: F401 + +print(f"[INFO] Using rl_training package from: {rl_training.__file__}") + + +@hydra_task_config(args_cli.task, args_cli.agent) +def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agent_cfg: RslRlOnPolicyRunnerCfg): + """Play with RSL-RL agent.""" + task_name = args_cli.task.split(":")[-1] + # override configurations with non-hydra CLI arguments + agent_cfg = cli_args.update_rsl_rl_cfg(agent_cfg, args_cli) + env_cfg.scene.num_envs = args_cli.num_envs if args_cli.num_envs is not None else 50 + + # set the environment seed + # note: certain randomizations occur in the environment initialization so we set the seed here + env_cfg.seed = agent_cfg.seed + env_cfg.sim.device = args_cli.device if args_cli.device is not None else env_cfg.sim.device + + # spawn the robot randomly in the grid (instead of their terrain levels) + env_cfg.scene.terrain.max_init_terrain_level = None + # reduce the number of terrains to save memory + if env_cfg.scene.terrain.terrain_generator is not None: + env_cfg.scene.terrain.terrain_generator.num_rows = 5 + env_cfg.scene.terrain.terrain_generator.num_cols = 5 + env_cfg.scene.terrain.terrain_generator.curriculum = False + + # disable randomization for play + env_cfg.observations.policy.enable_corruption = False + # Remove perturbations and parameter randomization for reproducible playback. + for event_name in ( + "randomize_rigid_body_material", + "randomize_wheel_material", + "randomize_support_material", + "randomize_rigid_body_mass", + "randomize_rigid_body_mass_base", + "randomize_rigid_body_inertia", + "randomize_com_positions", + "randomize_apply_external_force_torque", + "randomize_actuator_gains", + "randomize_push_robot", + ): + if hasattr(env_cfg.events, event_name): + setattr(env_cfg.events, event_name, None) + env_cfg.curriculum.command_levels = None + + if args_cli.keyboard: + env_cfg.scene.num_envs = 1 + env_cfg.terminations.time_out = None + env_cfg.commands.base_velocity.debug_vis = False + config = Se2KeyboardCfg( + v_x_sensitivity=env_cfg.commands.base_velocity.ranges.lin_vel_x[1]/2, + v_y_sensitivity=env_cfg.commands.base_velocity.ranges.lin_vel_y[1], + omega_z_sensitivity=env_cfg.commands.base_velocity.ranges.ang_vel_z[1], + ) + controller = Se2Keyboard(config) + env_cfg.observations.policy.velocity_commands = ObsTerm( + func=lambda env: torch.tensor(controller.advance(), dtype=torch.float32).unsqueeze(0).to(env.device), + ) + + # specify directory for logging experiments + log_root_path = os.path.join("logs", "rsl_rl", agent_cfg.experiment_name) + log_root_path = os.path.abspath(log_root_path) + if args_cli.visual_only: + resume_path = None + log_dir = os.path.join(log_root_path, "visual_only") + print("[INFO] Visual-only mode: no policy checkpoint will be loaded.") + else: + print(f"[INFO] Loading experiment from directory: {log_root_path}") + if args_cli.use_pretrained_checkpoint: + resume_path = get_published_pretrained_checkpoint("rsl_rl", task_name) + if not resume_path: + print("[INFO] Unfortunately a pre-trained checkpoint is currently unavailable for this task.") + return + elif args_cli.load_run is not None: + resume_path = get_checkpoint_path(log_root_path, agent_cfg.load_run, agent_cfg.load_checkpoint) + elif args_cli.checkpoint: + resume_path = retrieve_file_path(args_cli.checkpoint) + else: + resume_path = get_checkpoint_path(log_root_path, agent_cfg.load_run, agent_cfg.load_checkpoint) + log_dir = os.path.dirname(resume_path) + + # create isaac environment + env = gym.make(args_cli.task, cfg=env_cfg, render_mode="rgb_array" if args_cli.video else None) + + # convert to single-agent instance if required by the RL algorithm + if isinstance(env.unwrapped, DirectMARLEnv): + env = multi_agent_to_single_agent(env) + + # wrap for video recording + if args_cli.video: + video_kwargs = { + "video_folder": os.path.join(log_dir, "videos", "play"), + "step_trigger": lambda step: step == 0, + "video_length": args_cli.video_length, + "disable_logger": True, + } + print("[INFO] Recording videos during training.") + print_dict(video_kwargs, nesting=4) + env = gym.wrappers.RecordVideo(env, **video_kwargs) + + # wrap around environment for rsl-rl + env = RslRlVecEnvWrapper(env, clip_actions=agent_cfg.clip_actions) + + if args_cli.visual_only: + policy = None + else: + print(f"[INFO]: Loading model checkpoint from: {resume_path}") + # load previously trained model + ppo_runner = OnPolicyRunner(env, agent_cfg.to_dict(), log_dir=None, device=agent_cfg.device) + ppo_runner.load(resume_path) + + # obtain the trained policy for inference + policy = ppo_runner.get_inference_policy(device=env.unwrapped.device) + + # extract the neural network module + # we do this in a try-except to maintain backwards compatibility. + try: + # version 2.3 onwards + policy_nn = ppo_runner.alg.policy + except AttributeError: + # version 2.2 and below + policy_nn = ppo_runner.alg.actor_critic + + # export policy to onnx/jit + export_model_dir = os.path.join(os.path.dirname(resume_path), "exported") + export_policy_as_onnx( + policy=policy_nn, + normalizer=None, + path=export_model_dir, + filename="policy.onnx", + ) + export_policy_as_jit( + policy=policy_nn, + normalizer=None, + path=export_model_dir, + filename="policy.pt", + ) + + dt = env.unwrapped.step_dt + # print(dt, "dt") + # reset environment + + obs = env.get_observations() + + timestep = 0 + # simulate environment + while simulation_app.is_running(): + start_time = time.time() + # run everything in inference mode + with torch.inference_mode(): + # agent stepping + if policy is None: + actions = torch.zeros((env.num_envs, env.num_actions), device=env.unwrapped.device) + else: + actions = policy(obs) + + # env stepping + obs, _, _, _ = env.step(actions) + if args_cli.video: + timestep += 1 + # Exit the play loop after recording one video + if timestep == args_cli.video_length: + break + + if args_cli.keyboard: + camera_follow(env) + + # time delay for real-time evaluation + sleep_time = dt - (time.time() - start_time) + if args_cli.real_time and sleep_time > 0: + time.sleep(sleep_time) + + # close the simulator + env.close() + + +if __name__ == "__main__": + # run the main function + main() + # close sim app + simulation_app.close() diff --git a/scripts/reinforcement_learning/rsl_rl/train.py b/scripts/reinforcement_learning/rsl_rl/train.py new file mode 100644 index 0000000..b42ed3a --- /dev/null +++ b/scripts/reinforcement_learning/rsl_rl/train.py @@ -0,0 +1,214 @@ +# Copyright (c) 2025 Deep Robotics +# SPDX-License-Identifier: BSD 3-Clause + +# Copyright (c) 2024-2025 Ziqi Fan +# SPDX-License-Identifier: Apache-2.0 + +# Copyright (c) 2024-2025, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Script to train RL agent with RSL-RL.""" + +"""Launch Isaac Sim Simulator first.""" + +import argparse +import os +import sys + +PROJECT_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "..")) +PROJECT_SOURCE_PATH = os.path.join(PROJECT_ROOT, "source", "rl_training") +if PROJECT_SOURCE_PATH not in sys.path: + sys.path.insert(0, PROJECT_SOURCE_PATH) + +from isaaclab.app import AppLauncher + +# local imports +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) +import cli_args + +# add argparse arguments +parser = argparse.ArgumentParser(description="Train an RL agent with RSL-RL.") +parser.add_argument("--video", action="store_true", default=False, help="Record videos during training.") +parser.add_argument("--video_length", type=int, default=200, help="Length of the recorded video (in steps).") +parser.add_argument("--video_interval", type=int, default=2000, help="Interval between video recordings (in steps).") +parser.add_argument("--num_envs", type=int, default=None, help="Number of environments to simulate.") +parser.add_argument("--task", type=str, default=None, help="Name of the task.") +parser.add_argument( + "--agent", type=str, default="rsl_rl_cfg_entry_point", help="Name of the RL agent configuration entry point." +) +parser.add_argument("--seed", type=int, default=None, help="Seed used for the environment") +parser.add_argument("--max_iterations", type=int, default=None, help="RL Policy training iterations.") +parser.add_argument( + "--distributed", action="store_true", default=False, help="Run training with multiple GPUs or nodes." +) +# append RSL-RL cli arguments +cli_args.add_rsl_rl_args(parser) +# append AppLauncher cli args +AppLauncher.add_app_launcher_args(parser) +args_cli, hydra_args = parser.parse_known_args() + +# always enable cameras to record video +if args_cli.video: + args_cli.enable_cameras = True + +# clear out sys.argv for Hydra +sys.argv = [sys.argv[0]] + hydra_args + +# launch omniverse app +app_launcher = AppLauncher(args_cli) +simulation_app = app_launcher.app + +"""Rest everything follows.""" + +import gymnasium as gym +import os +import torch +from datetime import datetime + +from rsl_rl.runners import OnPolicyRunner + +from isaaclab.envs import ( + DirectMARLEnv, + DirectMARLEnvCfg, + DirectRLEnvCfg, + ManagerBasedRLEnvCfg, + multi_agent_to_single_agent, +) +from isaaclab.utils.dict import print_dict +from isaaclab.utils.io import dump_yaml +from isaaclab_rl.rsl_rl import RslRlOnPolicyRunnerCfg, RslRlVecEnvWrapper +from isaaclab_tasks.utils import get_checkpoint_path +from isaaclab_tasks.utils.hydra import hydra_task_config + +import rl_training +import rl_training.tasks # noqa: F401 + +print(f"[INFO] Using rl_training package from: {rl_training.__file__}") + +torch.backends.cuda.matmul.allow_tf32 = True +torch.backends.cudnn.allow_tf32 = True +torch.backends.cudnn.deterministic = False +torch.backends.cudnn.benchmark = False + + +def resolve_resume_path(log_root_path: str, load_run: str | None, load_checkpoint: str | None) -> str: + """Resolve a resume checkpoint, allowing explicit paths across experiment folders.""" + if load_checkpoint is not None: + checkpoint_path = os.path.expanduser(load_checkpoint) + if os.path.isabs(checkpoint_path) or os.path.dirname(checkpoint_path): + checkpoint_path = os.path.abspath(checkpoint_path) + if not os.path.isfile(checkpoint_path): + raise FileNotFoundError(f"Unable to find the checkpoint file: {checkpoint_path}") + return checkpoint_path + + if load_run is not None: + run_path = os.path.expanduser(load_run) + if os.path.isabs(run_path) or os.path.dirname(run_path): + run_path = os.path.abspath(run_path) + if not os.path.isdir(run_path): + raise FileNotFoundError(f"Unable to find the run directory: {run_path}") + return get_checkpoint_path(os.path.dirname(run_path), os.path.basename(run_path), load_checkpoint) + + return get_checkpoint_path(log_root_path, load_run, load_checkpoint) + + +@hydra_task_config(args_cli.task, args_cli.agent) +def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agent_cfg: RslRlOnPolicyRunnerCfg): + """Train with RSL-RL agent.""" + # override configurations with non-hydra CLI arguments + agent_cfg = cli_args.update_rsl_rl_cfg(agent_cfg, args_cli) + env_cfg.scene.num_envs = args_cli.num_envs if args_cli.num_envs is not None else env_cfg.scene.num_envs + agent_cfg.max_iterations = ( + args_cli.max_iterations if args_cli.max_iterations is not None else agent_cfg.max_iterations + ) + + # set the environment seed + # note: certain randomizations occur in the environment initialization so we set the seed here + env_cfg.seed = agent_cfg.seed + env_cfg.sim.device = args_cli.device if args_cli.device is not None else env_cfg.sim.device + + for reward_name in ( + "track_diff_wheel_yaw_exp", + "lin_vel_xy_yaw_command_l2", + "hipx_joint_pos_penalty", + "hipy_joint_pos_penalty", + ): + reward_term = getattr(env_cfg.rewards, reward_name, None) + reward_weight = None if reward_term is None else reward_term.weight + print(f"[INFO] Reward term {reward_name}: weight={reward_weight}") + + # multi-gpu training configuration + if args_cli.distributed: + env_cfg.sim.device = f"cuda:{app_launcher.local_rank}" + agent_cfg.device = f"cuda:{app_launcher.local_rank}" + + # set seed to have diversity in different threads + seed = agent_cfg.seed + app_launcher.local_rank + env_cfg.seed = seed + agent_cfg.seed = seed + + # specify directory for logging experiments + log_root_path = os.path.join("logs", "rsl_rl", agent_cfg.experiment_name) + log_root_path = os.path.abspath(log_root_path) + print(f"[INFO] Logging experiment in directory: {log_root_path}") + # specify directory for logging runs: {time-stamp}_{run_name} + log_dir = datetime.now().strftime("%Y-%m-%d_%H-%M-%S") + # The Ray Tune workflow extracts experiment name using the logging line below, hence, do not change it (see PR #2346, comment-2819298849) + print(f"Exact experiment name requested from command line: {log_dir}") + if agent_cfg.run_name: + log_dir += f"_{agent_cfg.run_name}" + log_dir = os.path.join(log_root_path, log_dir) + + # create isaac environment + env = gym.make(args_cli.task, cfg=env_cfg, render_mode="rgb_array" if args_cli.video else None) + + # convert to single-agent instance if required by the RL algorithm + if isinstance(env.unwrapped, DirectMARLEnv): + env = multi_agent_to_single_agent(env) + + # save resume path before creating a new log_dir + if agent_cfg.resume or agent_cfg.algorithm.class_name == "Distillation": + resume_path = resolve_resume_path(log_root_path, agent_cfg.load_run, agent_cfg.load_checkpoint) + + # wrap for video recording + if args_cli.video: + video_kwargs = { + "video_folder": os.path.join(log_dir, "videos", "train"), + "step_trigger": lambda step: step % args_cli.video_interval == 0, + "video_length": args_cli.video_length, + "disable_logger": True, + } + print("[INFO] Recording videos during training.") + print_dict(video_kwargs, nesting=4) + env = gym.wrappers.RecordVideo(env, **video_kwargs) + + # wrap around environment for rsl-rl + env = RslRlVecEnvWrapper(env, clip_actions=agent_cfg.clip_actions) + + # create runner from rsl-rl + runner = OnPolicyRunner(env, agent_cfg.to_dict(), log_dir=log_dir, device=agent_cfg.device) + # write git state to logs + runner.add_git_repo_to_log(__file__) + # load the checkpoint + if agent_cfg.resume or agent_cfg.algorithm.class_name == "Distillation": + print(f"[INFO]: Loading model checkpoint from: {resume_path}") + # load previously trained model + runner.load(resume_path) + + # dump the configuration into log-directory + dump_yaml(os.path.join(log_dir, "params", "env.yaml"), env_cfg) + dump_yaml(os.path.join(log_dir, "params", "agent.yaml"), agent_cfg) + # run training + runner.learn(num_learning_iterations=agent_cfg.max_iterations, init_at_random_ep_len=True) + + # close the simulator + env.close() + + +if __name__ == "__main__": + # run the main function + main() + # close sim app + simulation_app.close() diff --git a/scripts/tools/generate_solenopsis_cad_meshes.py b/scripts/tools/generate_solenopsis_cad_meshes.py new file mode 100644 index 0000000..9d10e06 --- /dev/null +++ b/scripts/tools/generate_solenopsis_cad_meshes.py @@ -0,0 +1,306 @@ +#!/usr/bin/env python3 +"""Generate rigid-link STL meshes for the Solenopsis (fire-ant) URDF. + +The SolidWorks STEP export uses millimetres and this script intentionally keeps +the STL coordinates in millimetres. The URDF applies a 0.001 mesh scale. + +Rigid grouping used here: + * base: root-level fixed parts plus each P65 fixed part and motor plate; + * knee: arm, P65 rotating part, W190 harness and module fasteners; + * wheel: W190 motor/wheel appearance assembly. + +The grouping is explicit so it can be audited when the CAD assembly changes. +""" + +from __future__ import annotations + +import argparse +import json +import math +from pathlib import Path + +from OCP.BRep import BRep_Builder +from OCP.BRepBuilderAPI import BRepBuilderAPI_Transform +from OCP.BRepMesh import BRepMesh_IncrementalMesh +from OCP.STEPCAFControl import STEPCAFControl_Reader +from OCP.StlAPI import StlAPI_Writer +from OCP.TCollection import TCollection_ExtendedString +from OCP.TDataStd import TDataStd_Name +from OCP.TDF import TDF_Label +from OCP.TDocStd import TDocStd_Document +from OCP.TopoDS import TopoDS_Compound +from OCP.XCAFApp import XCAFApp_Application +from OCP.XCAFDoc import XCAFDoc_DocumentTool, XCAFDoc_ShapeTool +from OCP.collections import Sequence_TDF_Label +from OCP.gp import gp_Trsf + + +# SolidWorks root axes -> ROS/URDF axes: (x, y, z)_urdf = (z, x, y)_sw. +SW_TO_URDF = ( + (0.0, 0.0, 1.0, 0.0), + (1.0, 0.0, 0.0, 0.0), + (0.0, 1.0, 0.0, 0.0), + (0.0, 0.0, 0.0, 1.0), +) + +IDENTITY = ( + (1.0, 0.0, 0.0, 0.0), + (0.0, 1.0, 0.0, 0.0), + (0.0, 0.0, 1.0, 0.0), + (0.0, 0.0, 0.0, 1.0), +) + +# Standardised URDF moving-link frames. In both cases the arm points along +# link -Z at q=0 and the URDF joint axis is link -Y, matching the established +# controller convention. The right-hand transform mirrors the CAD module in +# the physically correct way while remaining a proper rotation (determinant 1). +MODULE_TO_LEFT_LINK = ( + (0.0, 0.0, -1.0, 0.0), + (-1.0, 0.0, 0.0, 0.0), + (0.0, 1.0, 0.0, 0.0), + (0.0, 0.0, 0.0, 1.0), +) +MODULE_TO_RIGHT_LINK = ( + (0.0, 0.0, 1.0, 0.0), + (1.0, 0.0, 0.0, 0.0), + (0.0, 1.0, 0.0, 0.0), + (0.0, 0.0, 0.0, 1.0), +) + +# The two named components are fixed to the chassis side of the P65 joint. +BASE_COMPONENT_MARKERS = ( + "Solenopsis-02-002 关节电机固定板", + "P65MAX-S外观模型-A0-20260205-固定组件", +) +WHEEL_COMPONENT_MARKER = "W190-S 电机外观模型" + + +def matmul(a, b): + return tuple( + tuple(sum(a[i][k] * b[k][j] for k in range(4)) for j in range(4)) + for i in range(4) + ) + + +def transform_point(t, p): + q = (*p, 1.0) + return tuple(sum(t[i][k] * q[k] for k in range(4)) for i in range(3)) + + +def translation(x, y, z): + return ( + (1.0, 0.0, 0.0, x), + (0.0, 1.0, 0.0, y), + (0.0, 0.0, 1.0, z), + (0.0, 0.0, 0.0, 1.0), + ) + + +def location_matrix(label): + trsf = XCAFDoc_ShapeTool.GetLocation_s(label).Transformation() + return tuple( + tuple(trsf.Value(i, j) for j in range(1, 5)) + if i < 4 + else (0.0, 0.0, 0.0, 1.0) + for i in range(1, 5) + ) + + +def gp_transform(matrix): + t = gp_Trsf() + t.SetValues(*(matrix[i][j] for i in range(3) for j in range(4))) + return t + + +class CadDocument: + def __init__(self, step_path: Path): + self.document = TDocStd_Document(TCollection_ExtendedString("solenopsis")) + application = XCAFApp_Application.GetApplication_s() + application.NewDocument(TCollection_ExtendedString("MDTV-XCAF"), self.document) + + reader = STEPCAFControl_Reader() + reader.SetNameMode(True) + reader.SetColorMode(True) + reader.SetPropsMode(True) + status = reader.ReadFile(str(step_path)) + if "RetDone" not in str(status): + raise RuntimeError(f"STEP read failed: {status}") + if not reader.Transfer(self.document): + raise RuntimeError("STEP transfer into XCAF failed") + + self.shape_tool = XCAFDoc_DocumentTool.ShapeTool_s(self.document.Main()) + free = Sequence_TDF_Label() + self.shape_tool.GetFreeShapes(free) + if free.Length() != 1: + raise RuntimeError(f"Expected one root assembly, found {free.Length()}") + self.root = free.Value(1) + + @staticmethod + def label_name(label): + attribute = TDataStd_Name() + if label.FindAttribute(TDataStd_Name.GetID_s(), attribute): + return attribute.Get().ToExtString() + return "" + + @staticmethod + def referred(label): + target = TDF_Label() + if XCAFDoc_ShapeTool.GetReferredShape_s(label, target): + return target + return label + + def referred_name(self, label): + return self.label_name(self.referred(label)) + + def components(self, label): + target = self.referred(label) + result = Sequence_TDF_Label() + XCAFDoc_ShapeTool.GetComponents_s(target, result, False) + return [result.Value(i) for i in range(1, result.Length() + 1)] + + def referred_shape(self, component): + return XCAFDoc_ShapeTool.GetShape_s(self.referred(component)) + + +class ShapeGroup: + def __init__(self): + self.builder = BRep_Builder() + self.compound = TopoDS_Compound() + self.builder.MakeCompound(self.compound) + self.members = [] + + def add(self, cad: CadDocument, component, output_from_parent, member_name=None): + output_from_component = matmul(output_from_parent, location_matrix(component)) + shape = cad.referred_shape(component) + moved = BRepBuilderAPI_Transform( + shape, gp_transform(output_from_component), True, False + ).Shape() + self.builder.Add(self.compound, moved) + self.members.append(member_name or cad.referred_name(component)) + + def write_stl(self, output_path: Path, linear_deflection_mm: float): + output_path.parent.mkdir(parents=True, exist_ok=True) + BRepMesh_IncrementalMesh( + self.compound, linear_deflection_mm, False, 0.35, True + ).Perform() + writer = StlAPI_Writer() + # Binary STL is lossless relative to the tessellation and is typically + # five to eight times smaller than the default ASCII representation. + writer.ASCIIMode = False + if not writer.Write(self.compound, str(output_path)): + raise RuntimeError(f"Failed to write {output_path}") + + +def module_name_from_hip(hip_sw): + # URDF forward = SW +Z, URDF left = SW +X. + fore_aft = "f" if hip_sw[2] > 0.0 else "h" + left_right = "l" if hip_sw[0] > 0.0 else "r" + return fore_aft + left_right + + +def rotation_rpy(matrix): + """Return fixed-axis URDF roll/pitch/yaw for a proper rotation matrix.""" + pitch = math.asin(max(-1.0, min(1.0, -matrix[2][0]))) + if abs(math.cos(pitch)) > 1.0e-10: + roll = math.atan2(matrix[2][1], matrix[2][2]) + yaw = math.atan2(matrix[1][0], matrix[0][0]) + else: + roll = math.atan2(-matrix[0][1], matrix[1][1]) + yaw = 0.0 + return roll, pitch, yaw + + +def generate(step_path: Path, output_dir: Path, deflection_mm: float): + cad = CadDocument(step_path) + base = ShapeGroup() + knees = {} + wheels = {} + manifest = { + "source_step": str(step_path), + "stl_units": "millimetre", + "linear_deflection_mm": deflection_mm, + "axis_map": "URDF (x,y,z) = SolidWorks (z,x,y)", + "modules": {}, + } + + module_components = [] + for component in cad.components(cad.root): + name = cad.referred_name(component) + if name == "Solenopsis-02 运动模组总成": + module_components.append(component) + else: + base.add(cad, component, SW_TO_URDF, name) + + if len(module_components) != 4: + raise RuntimeError(f"Expected four motion modules, found {len(module_components)}") + + for module in module_components: + root_from_module = location_matrix(module) + hip_sw = transform_point(root_from_module, (0.0, 300.0, 0.0)) + leg = module_name_from_hip(hip_sw) + if leg in knees: + raise RuntimeError(f"Duplicate leg classification: {leg}") + + knees[leg] = ShapeGroup() + wheels[leg] = ShapeGroup() + + # The STEP assembly stores the four legs at four arbitrary joint + # angles. A URDF needs a common neutral pose, so q=0 is defined with + # every 300 mm arm vertical and all established joint axes at -Y. + module_to_link = ( + MODULE_TO_LEFT_LINK if leg.endswith("l") else MODULE_TO_RIGHT_LINK + ) + knee_from_module = matmul( + module_to_link, translation(0.0, -300.0, 0.0) + ) + wheel_from_module = module_to_link + + for child in cad.components(module): + name = cad.referred_name(child) + if any(marker in name for marker in BASE_COMPONENT_MARKERS): + base.add(cad, child, matmul(SW_TO_URDF, root_from_module), f"{leg}:{name}") + elif WHEEL_COMPONENT_MARKER in name: + wheels[leg].add(cad, child, wheel_from_module, name) + else: + knees[leg].add(cad, child, knee_from_module, name) + + urdf_from_module = matmul(SW_TO_URDF, root_from_module) + hip_urdf_mm = transform_point(SW_TO_URDF, hip_sw) + manifest["modules"][leg] = { + "hip_urdf_m": [value / 1000.0 for value in hip_urdf_mm], + "joint_rpy_rad": [0.0, 0.0, 0.0], + "joint_axis": [0.0, -1.0, 0.0], + "cad_assembly_rpy_rad": list(rotation_rpy(urdf_from_module)), + "base_members": [ + name for name in base.members if name.startswith(f"{leg}:") + ], + "knee_members": knees[leg].members, + "wheel_members": wheels[leg].members, + } + + base.write_stl(output_dir / "base_link.STL", deflection_mm) + for leg in ("fl", "fr", "hl", "hr"): + knees[leg].write_stl(output_dir / f"{leg}_knee.STL", deflection_mm) + wheels[leg].write_stl(output_dir / f"{leg}_wheel.STL", deflection_mm) + + (output_dir / "mesh_manifest.json").write_text( + json.dumps(manifest, ensure_ascii=False, indent=2) + "\n", encoding="utf-8" + ) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("step", type=Path) + parser.add_argument("output_dir", type=Path) + parser.add_argument( + "--deflection-mm", + type=float, + default=0.2, + help="Maximum linear tessellation deflection in millimetres (default: 0.2)", + ) + args = parser.parse_args() + generate(args.step.resolve(), args.output_dir.resolve(), args.deflection_mm) + + +if __name__ == "__main__": + main() diff --git a/scripts/tools/list_envs.py b/scripts/tools/list_envs.py new file mode 100644 index 0000000..7061935 --- /dev/null +++ b/scripts/tools/list_envs.py @@ -0,0 +1,80 @@ +# Copyright (c) 2025 Deep Robotics +# SPDX-License-Identifier: BSD 3-Clause + +# Copyright (c) 2024-2025 Ziqi Fan +# SPDX-License-Identifier: Apache-2.0 + +# Copyright (c) 2024-2025, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +""" +Script to print all the available environments in Isaac Lab. + +The script iterates over all registered environments and stores the details in a table. +It prints the name of the environment, the entry point and the config file. + +All the environments are registered in the `isaaclab_tasks` extension. They start +with `Isaac` in their name. +""" + +"""Launch Isaac Sim Simulator first.""" + +from isaaclab.app import AppLauncher + +# launch omniverse app +app_launcher = AppLauncher(headless=True) +simulation_app = app_launcher.app + + +"""Rest everything follows.""" + +import gymnasium as gym +import textwrap +from prettytable import PrettyTable + +import rl_training.tasks # noqa: F401 + + +def main(): + """Print all environments registered in `isaaclab_tasks` extension.""" + # print all the available environments + table = PrettyTable(["S. No.", "Task Name", "Entry Point", "Config"]) + table.title = "Available Environments in Isaac Lab" + # set alignment of table columns + table.align["Task Name"] = "l" + table.align["Entry Point"] = "l" + table.align["Config"] = "l" + table.hrules = 1 + + # set max width for text wrapping + max_width = 50 + + # count of environments + index = 0 + # acquire all Isaac environments names + for task_spec in gym.registry.values(): + if "Deeprobotics" in task_spec.id: + # wrap long text in each column before adding it to the table + task_name = textwrap.fill(task_spec.id, max_width) + entry_point = textwrap.fill(task_spec.entry_point, max_width) + config = textwrap.fill(task_spec.kwargs["env_cfg_entry_point"], max_width) + + # add details to table + table.add_row([index + 1, task_name, entry_point, config]) + # increment count + index += 1 + + print(table) + + +if __name__ == "__main__": + try: + # run the main function + main() + except Exception as e: + raise e + finally: + # close the app + simulation_app.close() diff --git a/scripts/tools/optimize_solenopsis_stl.py b/scripts/tools/optimize_solenopsis_stl.py new file mode 100644 index 0000000..536088f --- /dev/null +++ b/scripts/tools/optimize_solenopsis_stl.py @@ -0,0 +1,100 @@ +#!/usr/bin/env python3 +"""Create lightweight visual STL meshes by deterministic vertex clustering. + +The source meshes are left untouched. This optimization is intended only for +URDF ``visual`` geometry; collision geometry and all physical properties stay +in the URDF. Vertices are snapped to a regular grid, so the maximum movement +of any source vertex is bounded by ``sqrt(3) * grid_mm / 2``. +""" + +from __future__ import annotations + +import argparse +import json +import math +from pathlib import Path + +import numpy as np +import trimesh + + +def optimize_mesh(source: Path, destination: Path, grid_mm: float) -> dict[str, object]: + mesh = trimesh.load_mesh(source, process=True) + source_vertices = len(mesh.vertices) + source_faces = len(mesh.faces) + + # Round each vertex to its nearest grid point. The rounding makes the + # geometric error bounded and gives deterministic output across runs. + clustered_vertices = np.rint(mesh.vertices / grid_mm).astype(np.int64) + unique_grid_vertices, inverse = np.unique(clustered_vertices, axis=0, return_inverse=True) + faces = inverse[mesh.faces] + + # Remove triangles collapsed by clustering. + valid = ( + (faces[:, 0] != faces[:, 1]) + & (faces[:, 1] != faces[:, 2]) + & (faces[:, 2] != faces[:, 0]) + ) + faces = faces[valid] + + # Remove duplicate triangles while preserving the winding of the first + # occurrence. Sorting is used only to detect duplicates. + canonical_faces = np.sort(faces, axis=1) + _, first_indices = np.unique(canonical_faces, axis=0, return_index=True) + faces = faces[np.sort(first_indices)] + + vertices = unique_grid_vertices.astype(np.float64) * grid_mm + result = trimesh.Trimesh(vertices=vertices, faces=faces, process=False) + result.remove_unreferenced_vertices() + + destination.parent.mkdir(parents=True, exist_ok=True) + result.export(destination, file_type="stl") + + return { + "file": source.name, + "source_vertices": source_vertices, + "source_faces": source_faces, + "optimized_vertices": len(result.vertices), + "optimized_faces": len(result.faces), + "source_bytes": source.stat().st_size, + "optimized_bytes": destination.stat().st_size, + } + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("source_dir", type=Path) + parser.add_argument("output_dir", type=Path) + parser.add_argument("--grid-mm", type=float, default=1.0) + args = parser.parse_args() + + if args.grid_mm <= 0.0: + raise ValueError("--grid-mm must be positive") + + source_dir = args.source_dir.resolve() + output_dir = args.output_dir.resolve() + records = [ + optimize_mesh(source, output_dir / source.name, args.grid_mm) + for source in sorted(source_dir.glob("*.STL")) + ] + if not records: + raise RuntimeError(f"No STL files found in {source_dir}") + + manifest = { + "source_dir": str(source_dir), + "grid_mm": args.grid_mm, + "maximum_vertex_displacement_mm": math.sqrt(3.0) * args.grid_mm / 2.0, + "source_faces": sum(int(record["source_faces"]) for record in records), + "optimized_faces": sum(int(record["optimized_faces"]) for record in records), + "source_bytes": sum(int(record["source_bytes"]) for record in records), + "optimized_bytes": sum(int(record["optimized_bytes"]) for record in records), + "meshes": records, + } + (output_dir / "optimization_manifest.json").write_text( + json.dumps(manifest, ensure_ascii=False, indent=2) + "\n", encoding="utf-8" + ) + print(json.dumps(manifest, ensure_ascii=False, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/source/rl_training/config/extension.toml b/source/rl_training/config/extension.toml new file mode 100644 index 0000000..23c2839 --- /dev/null +++ b/source/rl_training/config/extension.toml @@ -0,0 +1,35 @@ +[package] + +# Semantic Versioning is used: https://semver.org/ +version = "1.0.0" + +# Description +category = "isaaclab" +readme = "README.md" + +title = "RL Training Repo for DeepRobotics" +author = "Bo Peng" +maintainer = "Bo Peng" +description="RL Training Repo for DeepRobotic, Based on IsaacLab." +repository = "https://github.com/DeepRoboticsLab/rl_training.git" +keywords = ["extension", "rl_training", "isaaclab"] + +[dependencies] +"isaaclab" = {} +"isaaclab_assets" = {} +"isaaclab_mimic" = {} +"isaaclab_rl" = {} +"isaaclab_tasks" = {} +# NOTE: Add additional dependencies here + +[[python.module]] +name = "rl_training" + +[isaaclab_settings] +# TODO: Uncomment and list any apt dependencies here. +# If none, leave it commented out. +# apt_deps = ["example_package"] +# TODO: Uncomment and provide path to a ros_ws +# with rosdeps to be installed. If none, +# leave it commented out. +# ros_ws = "path/from/extension_root/to/ros_ws" diff --git a/source/rl_training/pyproject.toml b/source/rl_training/pyproject.toml new file mode 100644 index 0000000..d90ac35 --- /dev/null +++ b/source/rl_training/pyproject.toml @@ -0,0 +1,3 @@ +[build-system] +requires = ["setuptools", "wheel", "toml"] +build-backend = "setuptools.build_meta" diff --git a/source/rl_training/rl_training/__init__.py b/source/rl_training/rl_training/__init__.py new file mode 100644 index 0000000..27e3905 --- /dev/null +++ b/source/rl_training/rl_training/__init__.py @@ -0,0 +1,15 @@ +# Copyright (c) 2025 Deep Robotics +# SPDX-License-Identifier: BSD 3-Clause +# +# # Copyright (c) 2024-2025 Ziqi Fan +# SPDX-License-Identifier: Apache-2.0 + +""" +Python module serving as a project/extension template. +""" + +# Register Gym environments. +from .tasks import * + +# Register UI extensions. +from .ui_extension_example import * diff --git a/source/rl_training/rl_training/assets/__init__.py b/source/rl_training/rl_training/assets/__init__.py new file mode 100644 index 0000000..c970e4f --- /dev/null +++ b/source/rl_training/rl_training/assets/__init__.py @@ -0,0 +1,32 @@ +# Copyright (c) 2025 Deep Robotics +# SPDX-License-Identifier: BSD 3-Clause +# +# # Copyright (c) 2024-2025 Ziqi Fan +# SPDX-License-Identifier: Apache-2.0 + +# Copyright (c) 2022-2025, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Package containing asset and sensor configurations.""" + +import os +import toml + +## +# Configuration for different assets. +## + +# Conveniences to other module directories via relative paths +ISAACLAB_ASSETS_EXT_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../")) +"""Path to the extension source directory.""" + +ISAACLAB_ASSETS_DATA_DIR = os.path.join(ISAACLAB_ASSETS_EXT_DIR, "../../deep_robotics_model") +"""Path to the extension data directory.""" + +ISAACLAB_ASSETS_METADATA = toml.load(os.path.join(ISAACLAB_ASSETS_EXT_DIR, "config", "extension.toml")) +"""Extension metadata dictionary parsed from the extension.toml file.""" + +# Configure the module-level variables +__version__ = ISAACLAB_ASSETS_METADATA["package"]["version"] diff --git a/source/rl_training/rl_training/assets/deeprobotics.py b/source/rl_training/rl_training/assets/deeprobotics.py new file mode 100644 index 0000000..20ac20c --- /dev/null +++ b/source/rl_training/rl_training/assets/deeprobotics.py @@ -0,0 +1,124 @@ +# Copyright (c) 2025 Deep Robotics +# SPDX-License-Identifier: BSD 3-Clause + +# Copyright (c) 2024-2025 Ziqi Fan +# SPDX-License-Identifier: Apache-2.0 + +import isaaclab.sim as sim_utils +from isaaclab.actuators import DCMotorCfg, DelayedPDActuatorCfg +from isaaclab.assets.articulation import ArticulationCfg + +from rl_training.assets import ISAACLAB_ASSETS_DATA_DIR + +DEEPROBOTICS_LITE3_CFG = ArticulationCfg( + spawn=sim_utils.UsdFileCfg( + usd_path=f"{ISAACLAB_ASSETS_DATA_DIR}/Lite3/Lite3_usd/Lite3.usd", + activate_contact_sensors=True, + rigid_props=sim_utils.RigidBodyPropertiesCfg( + disable_gravity=False, + retain_accelerations=False, + linear_damping=0.0, + angular_damping=0.0, + max_linear_velocity=1000.0, + max_angular_velocity=1000.0, + max_depenetration_velocity=1.0, + ), + articulation_props=sim_utils.ArticulationRootPropertiesCfg( + enabled_self_collisions=False, solver_position_iteration_count=4, solver_velocity_iteration_count=1 + ), + ), + init_state=ArticulationCfg.InitialStateCfg( + pos=(0.0, 0.0, 0.35), + joint_pos={ + + ".*Knee_joint": 0.0, + }, + joint_vel={".*": 0.0}, + ), + soft_joint_pos_limit_factor=0.99, + actuators={ + # "Hip": DelayedPDActuatorCfg( + # joint_names_expr=[".*_Hip[X,Y]_joint"], + # effort_limit=24.0, + # velocity_limit=26.2, + # stiffness=30.0, + # damping=1.0, + # friction=0.0, + # armature=0.0, + # min_delay=0, + # max_delay=5, + # ), + "Knee": DelayedPDActuatorCfg( + joint_names_expr=[".*_Knee_joint"], + effort_limit=36.0, + velocity_limit=26.2, + stiffness=30.0, + damping=1.0, + friction=0.0, + armature=0.0, + min_delay=0, + max_delay=5, + ), + }, +) + +DEEPROBOTICS_M20_CFG = ArticulationCfg( + spawn=sim_utils.UrdfFileCfg( + asset_path=f"{ISAACLAB_ASSETS_DATA_DIR}/M20/M20_urdf/urdf/Solenopsis.urdf", + fix_base=False, + merge_fixed_joints=True, + joint_drive=None, + activate_contact_sensors=True, + visual_material=sim_utils.PreviewSurfaceCfg( + diffuse_color=(0.42, 0.45, 0.50), + roughness=0.65, + metallic=0.15, + ), + rigid_props=sim_utils.RigidBodyPropertiesCfg( + disable_gravity=False, + retain_accelerations=False, + linear_damping=0.0, + angular_damping=0.0, + max_linear_velocity=1000.0, + max_angular_velocity=1000.0, + max_depenetration_velocity=1.0, + ), + articulation_props=sim_utils.ArticulationRootPropertiesCfg( + enabled_self_collisions=False, solver_position_iteration_count=4, solver_velocity_iteration_count=1 + ), + ), + init_state=ArticulationCfg.InitialStateCfg( + pos=(0.0, 0.0, 0.306), + joint_pos={ + "f[l,r]_knee_joint": 0.5, + "h[l,r]_knee_joint": -0.5, + ".*wheel_joint": 0.0, + }, + joint_vel={".*": 0.0}, + ), + soft_joint_pos_limit_factor=0.9, + actuators={ + "joint": DelayedPDActuatorCfg( + joint_names_expr=[".*knee_joint"], + effort_limit=75.0, + velocity_limit=17.8024, + stiffness=80.0, + damping=2.0, + friction=0.0, + armature=0.0, + min_delay=0, + max_delay=5, + ), + "wheel": DelayedPDActuatorCfg( + joint_names_expr=[".*_wheel_joint"], + effort_limit=28.68, + velocity_limit=109.956, + stiffness=0.0, + damping=0.6, + friction=1.0, + armature=0.00243216, + min_delay=0, + max_delay=5, + ), + }, +) diff --git a/source/rl_training/rl_training/tasks/__init__.py b/source/rl_training/rl_training/tasks/__init__.py new file mode 100644 index 0000000..e13871c --- /dev/null +++ b/source/rl_training/rl_training/tasks/__init__.py @@ -0,0 +1,27 @@ +# Copyright (c) 2025 Deep Robotics +# SPDX-License-Identifier: BSD 3-Clause +# +# # Copyright (c) 2024-2025 Ziqi Fan +# SPDX-License-Identifier: Apache-2.0 + +# Copyright (c) 2022-2025, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Package containing task implementations for various robotic environments.""" + +import os +import toml + +from isaaclab_tasks.utils import import_packages + +## +# Register Gym environments. +## + + +# The blacklist is used to prevent importing configs from sub-packages +_BLACKLIST_PKGS = ["utils"] +# Import all configs in this package +import_packages(__name__, _BLACKLIST_PKGS) diff --git a/source/rl_training/rl_training/tasks/manager_based/__init__.py b/source/rl_training/rl_training/tasks/manager_based/__init__.py new file mode 100644 index 0000000..804027a --- /dev/null +++ b/source/rl_training/rl_training/tasks/manager_based/__init__.py @@ -0,0 +1,16 @@ +# Copyright (c) 2025 Deep Robotics +# SPDX-License-Identifier: BSD 3-Clause +# +# # Copyright (c) 2024-2025 Ziqi Fan +# SPDX-License-Identifier: Apache-2.0 + +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +""" +Config-based workflow environments. +""" + +import gymnasium as gym diff --git a/source/rl_training/rl_training/tasks/manager_based/locomotion/__init__.py b/source/rl_training/rl_training/tasks/manager_based/locomotion/__init__.py new file mode 100644 index 0000000..6d8a370 --- /dev/null +++ b/source/rl_training/rl_training/tasks/manager_based/locomotion/__init__.py @@ -0,0 +1,14 @@ +# Copyright (c) 2025 Deep Robotics +# SPDX-License-Identifier: BSD 3-Clause +# +# # Copyright (c) 2024-2025 Ziqi Fan +# SPDX-License-Identifier: Apache-2.0 + +# Copyright (c) 2022-2025, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Locomotion environments for legged robots.""" + +from .velocity import * # noqa diff --git a/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/__init__.py b/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/__init__.py new file mode 100644 index 0000000..dd54bb1 --- /dev/null +++ b/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/__init__.py @@ -0,0 +1,18 @@ +# Copyright (c) 2025 Deep Robotics +# SPDX-License-Identifier: BSD 3-Clause +# +# # Copyright (c) 2024-2025 Ziqi Fan +# SPDX-License-Identifier: Apache-2.0 + +# Copyright (c) 2022-2025, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Locomotion environments with velocity-tracking commands. + +These environments are based on the `legged_gym` environments provided by Rudin et al. + +Reference: + https://github.com/leggedrobotics/legged_gym +""" diff --git a/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/config/__init__.py b/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/config/__init__.py new file mode 100644 index 0000000..6a7892a --- /dev/null +++ b/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/config/__init__.py @@ -0,0 +1,15 @@ +# Copyright (c) 2025 Deep Robotics +# SPDX-License-Identifier: BSD 3-Clause +# +# # Copyright (c) 2024-2025 Ziqi Fan +# SPDX-License-Identifier: Apache-2.0 + +# Copyright (c) 2022-2025, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Configurations for velocity-based locomotion environments.""" + +# We leave this file empty since we don't want to expose any configs in this package directly. +# We still need this file to import the "config" module in the parent package. diff --git a/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/config/quadruped/__init__.py b/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/config/quadruped/__init__.py new file mode 100644 index 0000000..6a7892a --- /dev/null +++ b/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/config/quadruped/__init__.py @@ -0,0 +1,15 @@ +# Copyright (c) 2025 Deep Robotics +# SPDX-License-Identifier: BSD 3-Clause +# +# # Copyright (c) 2024-2025 Ziqi Fan +# SPDX-License-Identifier: Apache-2.0 + +# Copyright (c) 2022-2025, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Configurations for velocity-based locomotion environments.""" + +# We leave this file empty since we don't want to expose any configs in this package directly. +# We still need this file to import the "config" module in the parent package. diff --git a/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/config/quadruped/deeprobotics_lite3/__init__.py b/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/config/quadruped/deeprobotics_lite3/__init__.py new file mode 100644 index 0000000..8bdfbf1 --- /dev/null +++ b/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/config/quadruped/deeprobotics_lite3/__init__.py @@ -0,0 +1,36 @@ +# Copyright (c) 2025 Deep Robotics +# SPDX-License-Identifier: BSD 3-Clause + +# Copyright (c) 2024-2025 Ziqi Fan +# SPDX-License-Identifier: Apache-2.0 + +import gymnasium as gym + +from . import agents + +## +# Register Gym environments. +## + +gym.register( + id="Flat-Deeprobotics-Lite3-v0", + entry_point="isaaclab.envs:ManagerBasedRLEnv", + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": f"{__name__}.flat_env_cfg:DeeproboticsLite3FlatEnvCfg", + "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:DeeproboticsLite3FlatPPORunnerCfg", + "cusrl_cfg_entry_point": f"{agents.__name__}.cusrl_ppo_cfg:DeeproboticsLite3FlatTrainerCfg", + }, +) + +gym.register( + id="Rough-Deeprobotics-Lite3-v0", + entry_point="isaaclab.envs:ManagerBasedRLEnv", + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": f"{__name__}.rough_env_cfg:DeeproboticsLite3RoughEnvCfg", + "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:DeeproboticsLite3RoughPPORunnerCfg", + "cusrl_cfg_entry_point": f"{agents.__name__}.cusrl_ppo_cfg:DeeproboticsLite3RoughTrainerCfg", + }, +) + diff --git a/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/config/quadruped/deeprobotics_lite3/agents/__init__.py b/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/config/quadruped/deeprobotics_lite3/agents/__init__.py new file mode 100644 index 0000000..26c100d --- /dev/null +++ b/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/config/quadruped/deeprobotics_lite3/agents/__init__.py @@ -0,0 +1,5 @@ +# Copyright (c) 2025 Deep Robotics +# SPDX-License-Identifier: BSD 3-Clause +# +# # Copyright (c) 2024-2025 Ziqi Fan +# SPDX-License-Identifier: Apache-2.0 diff --git a/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/config/quadruped/deeprobotics_lite3/agents/rsl_rl_ppo_cfg.py b/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/config/quadruped/deeprobotics_lite3/agents/rsl_rl_ppo_cfg.py new file mode 100644 index 0000000..b259309 --- /dev/null +++ b/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/config/quadruped/deeprobotics_lite3/agents/rsl_rl_ppo_cfg.py @@ -0,0 +1,48 @@ +# Copyright (c) 2025 Deep Robotics +# SPDX-License-Identifier: BSD 3-Clause + +# Copyright (c) 2024-2025 Ziqi Fan +# SPDX-License-Identifier: Apache-2.0 + +from isaaclab.utils import configclass +from isaaclab_rl.rsl_rl import RslRlOnPolicyRunnerCfg, RslRlPpoActorCriticCfg, RslRlPpoAlgorithmCfg + + +@configclass +class DeeproboticsLite3RoughPPORunnerCfg(RslRlOnPolicyRunnerCfg): + num_steps_per_env = 24 + max_iterations = 10000 + save_interval = 100 + experiment_name = "deeprobotics_lite3_rough" + empirical_normalization = False + clip_actions = 100 + policy = RslRlPpoActorCriticCfg( + init_noise_std=1.0, + noise_std_type="log", + actor_hidden_dims=[512, 256, 128], + critic_hidden_dims=[512, 256, 128], + activation="elu", + ) + algorithm = RslRlPpoAlgorithmCfg( + value_loss_coef=1.0, + use_clipped_value_loss=True, + clip_param=0.2, + entropy_coef=0.01, + num_learning_epochs=5, + num_mini_batches=4, + learning_rate=1.0e-3, + schedule="adaptive", + gamma=0.99, + lam=0.95, + desired_kl=0.01, + max_grad_norm=1.0, + ) + + +@configclass +class DeeproboticsLite3FlatPPORunnerCfg(DeeproboticsLite3RoughPPORunnerCfg): + def __post_init__(self): + super().__post_init__() + + self.max_iterations = 10000 + self.experiment_name = "deeprobotics_lite3_flat" diff --git a/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/config/quadruped/deeprobotics_lite3/flat_env_cfg.py b/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/config/quadruped/deeprobotics_lite3/flat_env_cfg.py new file mode 100644 index 0000000..448cb82 --- /dev/null +++ b/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/config/quadruped/deeprobotics_lite3/flat_env_cfg.py @@ -0,0 +1,32 @@ +# Copyright (c) 2025 Deep Robotics +# SPDX-License-Identifier: BSD 3-Clause + +# Copyright (c) 2024-2025 Ziqi Fan +# SPDX-License-Identifier: Apache-2.0 + +from isaaclab.utils import configclass + +from .rough_env_cfg import DeeproboticsLite3RoughEnvCfg + + +@configclass +class DeeproboticsLite3FlatEnvCfg(DeeproboticsLite3RoughEnvCfg): + def __post_init__(self): + # post init of parent + super().__post_init__() + + # override rewards + # self.rewards.base_height_l2.params["sensor_cfg"] = None + # change terrain to flat + self.scene.terrain.terrain_type = "plane" + self.scene.terrain.terrain_generator = None + # no height scan + # self.scene.height_scanner = None + self.observations.policy.height_scan = None + # self.observations.critic.height_scan = None + # no terrain curriculum + self.curriculum.terrain_levels = None + + # If the weight of rewards is 0, set rewards to None + if self.__class__.__name__ == "DeeproboticsLite3FlatEnvCfg": + self.disable_zero_weight_rewards() diff --git a/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/config/quadruped/deeprobotics_lite3/rough_env_cfg.py b/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/config/quadruped/deeprobotics_lite3/rough_env_cfg.py new file mode 100644 index 0000000..5841730 --- /dev/null +++ b/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/config/quadruped/deeprobotics_lite3/rough_env_cfg.py @@ -0,0 +1,183 @@ +# Copyright (c) 2025 Deep Robotics +# SPDX-License-Identifier: BSD 3-Clause + +# Copyright (c) 2024-2025 Ziqi Fan +# SPDX-License-Identifier: Apache-2.0 + +from isaaclab.utils import configclass + +from rl_training.tasks.manager_based.locomotion.velocity.velocity_env_cfg import LocomotionVelocityRoughEnvCfg +# from isaaclab.sensors.ray_caster import GridPatternCfg +## +# Pre-defined configs +## +from rl_training.assets.deeprobotics import DEEPROBOTICS_LITE3_CFG # isort: skip + + +@configclass +class DeeproboticsLite3RoughEnvCfg(LocomotionVelocityRoughEnvCfg): + base_link_name = "TORSO" + foot_link_name = ".*_FOOT" + # fmt: off + joint_names = [ + "FL_HipX_joint", "FL_HipY_joint", "FL_Knee_joint", + "FR_HipX_joint", "FR_HipY_joint", "FR_Knee_joint", + "HL_HipX_joint", "HL_HipY_joint", "HL_Knee_joint", + "HR_HipX_joint", "HR_HipY_joint", "HR_Knee_joint", + ] + + link_names = [ + 'TORSO', + 'FL_HIP', 'FR_HIP', 'HL_HIP', 'HR_HIP', + 'FL_THIGH', 'FR_THIGH', 'HL_THIGH', 'HR_THIGH', + 'FL_SHANK', 'FR_SHANK', 'HL_SHANK', 'HR_SHANK', + 'FL_FOOT', 'FR_FOOT', 'HL_FOOT', 'HR_FOOT', + ] + # fmt: on + + def __post_init__(self): + # post init of parent + super().__post_init__() + + # ------------------------------Sence------------------------------ + self.scene.robot = DEEPROBOTICS_LITE3_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") + self.scene.height_scanner.prim_path = "{ENV_REGEX_NS}/Robot/" + self.base_link_name + self.scene.height_scanner_base.prim_path = "{ENV_REGEX_NS}/Robot/" + self.base_link_name + self.scene.height_scanner.pattern_cfg.resolution = 0.07 # = GridPatternCfg(resolution=0.07, size=[1.6, 1.0]), + + # ------------------------------Observations------------------------------ + self.observations.policy.base_lin_vel = None # type: ignore + self.observations.policy.height_scan = None # type: ignore + self.observations.policy.base_ang_vel.scale = 0.25 + self.observations.policy.joint_pos.scale = 1.0 + self.observations.policy.joint_vel.scale = 0.05 + self.observations.policy.joint_pos.params["asset_cfg"].joint_names = self.joint_names + self.observations.policy.joint_vel.params["asset_cfg"].joint_names = self.joint_names + + # ------------------------------Actions------------------------------ + # reduce action scale + self.actions.joint_pos.scale = {".*_HipX_joint": 0.125, "^(?!.*_HipX_joint).*": 0.25} + self.actions.joint_pos.clip = {".*": (-100.0, 100.0)} + self.actions.joint_pos.joint_names = self.joint_names + + # ------------------------------Events------------------------------ + self.events.randomize_reset_base.params = { + "pose_range": { + "x": (-1.0, 1.0), + "y": (-1.0, 1.0), + "z": (0.0, 0.0), + "roll": (-0.3, 0.3), + "pitch": (-0.3, 0.3), + "yaw": (-3.14, 3.14), + }, + "velocity_range": { + "x": (-0.2, 0.2), + "y": (-0.2, 0.2), + "z": (-0.2, 0.2), + "roll": (-0.05, 0.05), + "pitch": (-0.05, 0.05), + "yaw": (-0.0, 0.0), + }, + } + + + self.events.randomize_rigid_body_mass.params["asset_cfg"].body_names = self.link_names # [self.base_link_name] + self.events.randomize_rigid_body_mass_base = None + self.events.randomize_com_positions.params["asset_cfg"].body_names = self.base_link_name # [self.base_link_name] + # self.events.randomize_com_positions = None + self.events.randomize_apply_external_force_torque = None + self.events.randomize_push_robot = None + self.events.randomize_actuator_gains.params["asset_cfg"].joint_names = self.joint_names + + # set terrain generation probability to 0 for boxes and stairs + self.scene.terrain.terrain_generator.sub_terrains["random_rough"].proportion = 0.4 + self.scene.terrain.terrain_generator.sub_terrains["hf_pyramid_slope"].proportion = 0.3 + self.scene.terrain.terrain_generator.sub_terrains["hf_pyramid_slope_inv"].proportion = 0.3 + self.scene.terrain.terrain_generator.sub_terrains["boxes"].proportion = 0.0 + self.scene.terrain.terrain_generator.sub_terrains["pyramid_stairs"].proportion = 0.0 + self.scene.terrain.terrain_generator.sub_terrains["pyramid_stairs_inv"].proportion = 0.0 + # scale down the terrains because the robot is small + # self.scene.terrain.terrain_generator.sub_terrains["boxes"].grid_height_range = (0.025, 0.1) + # self.scene.terrain.terrain_generator.sub_terrains["boxes"].grid_width = 0.8 + self.scene.terrain.terrain_generator.sub_terrains["random_rough"].noise_range = (0.01, 0.06) + self.scene.terrain.terrain_generator.sub_terrains["random_rough"].noise_step = 0.01 + + # ------------------------------Rewards------------------------------ + self.rewards.action_rate_l2.weight = -0.02 #-0.02 + # self.rewards.smoothness_2.weight = -0.0075 + + self.rewards.base_height_l2.weight = -10.0 + self.rewards.base_height_l2.params["target_height"] = 0.35 + self.rewards.base_height_l2.params["asset_cfg"].body_names = [self.base_link_name] + + self.rewards.feet_air_time.weight = 5.0 # 5.0 + self.rewards.feet_air_time.params["threshold"] = 0.5 + self.rewards.feet_air_time.params["sensor_cfg"].body_names = [self.foot_link_name] + self.rewards.feet_air_time_variance.weight = -8.0 # -8.0 + self.rewards.feet_air_time_variance.params["sensor_cfg"].body_names = [self.foot_link_name] + self.rewards.feet_slide.weight = -0.05 + self.rewards.feet_slide.params["sensor_cfg"].body_names = [self.foot_link_name] + self.rewards.feet_slide.params["asset_cfg"].body_names = [self.foot_link_name] + self.rewards.stand_still.weight = -0.5 # -1.0 + self.rewards.stand_still.params["asset_cfg"].joint_names = self.joint_names + self.rewards.stand_still.params["command_threshold"] = 0.1 + self.rewards.feet_height_body.weight = -2.5 # -2.5 + self.rewards.feet_height_body.params["target_height"] = -0.35 + self.rewards.feet_height_body.params["asset_cfg"].body_names = [self.foot_link_name] + self.rewards.feet_height.weight = -0.2 # -0.2 + self.rewards.feet_height.params["asset_cfg"].body_names = [self.foot_link_name] + self.rewards.feet_height.params["target_height"] = 0.05 + self.rewards.contact_forces.weight = -1e-1 # -2e-2 + self.rewards.contact_forces.params["sensor_cfg"].body_names = [self.foot_link_name] + + self.rewards.lin_vel_z_l2.weight = -2.0 #-2.0 + self.rewards.ang_vel_xy_l2.weight = -0.05 # -0.05 + + self.rewards.track_lin_vel_xy_exp.weight = 3.0 + self.rewards.track_ang_vel_z_exp.weight = 1.5 + + self.rewards.undesired_contacts.weight = -0.5 + self.rewards.undesired_contacts.params["sensor_cfg"].body_names = [f"^(?!.*{self.foot_link_name}).*"] + + self.rewards.joint_torques_l2.weight = -2.5e-5 + self.rewards.joint_acc_l2.weight = -1e-8 + self.rewards.joint_deviation_l1.weight = -0.5 + self.rewards.joint_deviation_l1.params["asset_cfg"].joint_names = [".*HipX.*"] + self.rewards.joint_power.weight = -2e-5 + self.rewards.flat_orientation_l2.weight = -5.0 + + # add the following rewards to improve the gait + self.rewards.feet_gait.weight = 0.5 + self.rewards.feet_gait.params["synced_feet_pair_names"] = [ + ["FL_FOOT", "HR_FOOT"], + ["FR_FOOT", "HL_FOOT"] + ] + + self.rewards.joint_mirror.weight = -0.05 + self.rewards.joint_mirror.params["mirror_joints"] = [ + ["FL_(HipX|HipY|Knee).*", "HR_(HipX|HipY|Knee).*"], + ["FR_(HipX|HipY|Knee).*", "HL_(HipX|HipY|Knee).*"], + ] + + self.rewards.joint_pos_limits.weight = -5.0 + # self.rewards.joint_pos_penalty.weight = -1.0 + self.rewards.feet_contact_without_cmd.weight = 0.1 + self.rewards.feet_contact_without_cmd.params["sensor_cfg"].body_names = [self.foot_link_name] + + + # If the weight of rewards is 0, set rewards to None + if self.__class__.__name__ == "DeeproboticsLite3RoughEnvCfg": + self.disable_zero_weight_rewards() + + # ------------------------------Terminations------------------------------ + self.terminations.illegal_contact = None + # self.terminations.bad_orientation_2 = None + + # ------------------------------Curriculums------------------------------ + # self.curriculum.command_levels.params["range_multiplier"] = (0.2, 1.0) + self.curriculum.command_levels = None + + # ------------------------------Commands------------------------------ + self.commands.base_velocity.ranges.lin_vel_x = (-1.5, 1.5) + self.commands.base_velocity.ranges.lin_vel_y = (-0.8, 0.8) + self.commands.base_velocity.ranges.ang_vel_z = (-0.8, 0.8) \ No newline at end of file diff --git a/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/config/wheeled/__init__.py b/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/config/wheeled/__init__.py new file mode 100644 index 0000000..6a7892a --- /dev/null +++ b/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/config/wheeled/__init__.py @@ -0,0 +1,15 @@ +# Copyright (c) 2025 Deep Robotics +# SPDX-License-Identifier: BSD 3-Clause +# +# # Copyright (c) 2024-2025 Ziqi Fan +# SPDX-License-Identifier: Apache-2.0 + +# Copyright (c) 2022-2025, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Configurations for velocity-based locomotion environments.""" + +# We leave this file empty since we don't want to expose any configs in this package directly. +# We still need this file to import the "config" module in the parent package. diff --git a/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/config/wheeled/deeprobotics_m20/__init__.py b/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/config/wheeled/deeprobotics_m20/__init__.py new file mode 100644 index 0000000..9acad03 --- /dev/null +++ b/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/config/wheeled/deeprobotics_m20/__init__.py @@ -0,0 +1,46 @@ +# Copyright (c) 2025 Deep Robotics +# SPDX-License-Identifier: BSD 3-Clause +# +# # Copyright (c) 2024-2025 Ziqi Fan +# SPDX-License-Identifier: Apache-2.0 + +import gymnasium as gym + +from . import agents + +## +# Register Gym environments. +## + +gym.register( + id="Flat-Deeprobotics-M20-v0", + entry_point="isaaclab.envs:ManagerBasedRLEnv", + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": f"{__name__}.flat_env_cfg:DeeproboticsM20FlatEnvCfg", + "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:DeeproboticsM20FlatPPORunnerCfg", + "cusrl_cfg_entry_point": f"{agents.__name__}.cusrl_ppo_cfg:DeeproboticsM20FlatTrainerCfg", + }, +) + +gym.register( + id="Rough-Deeprobotics-M20-v0", + entry_point="isaaclab.envs:ManagerBasedRLEnv", + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": f"{__name__}.rough_env_cfg:DeeproboticsM20RoughEnvCfg", + "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:DeeproboticsM20RoughPPORunnerCfg", + "cusrl_cfg_entry_point": f"{agents.__name__}.cusrl_ppo_cfg:DeeproboticsM20RoughTrainerCfg", + }, +) + +gym.register( + id="Rough-Deeprobotics-M20-Support-v0", + entry_point="isaaclab.envs:ManagerBasedRLEnv", + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": f"{__name__}.rough_env_cfg:DeeproboticsM20SupportRoughEnvCfg", + "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:DeeproboticsM20SupportRoughPPORunnerCfg", + "cusrl_cfg_entry_point": f"{agents.__name__}.cusrl_ppo_cfg:DeeproboticsM20RoughTrainerCfg", + }, +) diff --git a/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/config/wheeled/deeprobotics_m20/agents/__init__.py b/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/config/wheeled/deeprobotics_m20/agents/__init__.py new file mode 100644 index 0000000..26c100d --- /dev/null +++ b/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/config/wheeled/deeprobotics_m20/agents/__init__.py @@ -0,0 +1,5 @@ +# Copyright (c) 2025 Deep Robotics +# SPDX-License-Identifier: BSD 3-Clause +# +# # Copyright (c) 2024-2025 Ziqi Fan +# SPDX-License-Identifier: Apache-2.0 diff --git a/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/config/wheeled/deeprobotics_m20/agents/rsl_rl_ppo_cfg.py b/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/config/wheeled/deeprobotics_m20/agents/rsl_rl_ppo_cfg.py new file mode 100644 index 0000000..8a6e0c9 --- /dev/null +++ b/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/config/wheeled/deeprobotics_m20/agents/rsl_rl_ppo_cfg.py @@ -0,0 +1,56 @@ +# Copyright (c) 2025 Deep Robotics +# SPDX-License-Identifier: BSD 3-Clause + +# Copyright (c) 2024-2025 Ziqi Fan +# SPDX-License-Identifier: Apache-2.0 + +from isaaclab.utils import configclass +from isaaclab_rl.rsl_rl import RslRlOnPolicyRunnerCfg, RslRlPpoActorCriticCfg, RslRlPpoAlgorithmCfg + + +@configclass +class DeeproboticsM20RoughPPORunnerCfg(RslRlOnPolicyRunnerCfg): + num_steps_per_env = 24 + max_iterations = 20000 + save_interval = 100 + experiment_name = "deeprobotics_m20_rough" + empirical_normalization = False + clip_actions = 100 + policy = RslRlPpoActorCriticCfg( + init_noise_std=1.0, + noise_std_type="log", + actor_hidden_dims=[512, 256, 128], + critic_hidden_dims=[512, 256, 128], + activation="elu", + ) + algorithm = RslRlPpoAlgorithmCfg( + value_loss_coef=1.0, + use_clipped_value_loss=True, + clip_param=0.2, + entropy_coef=0.01, + num_learning_epochs=5, + num_mini_batches=4, + learning_rate=1.0e-3, + schedule="adaptive", + gamma=0.99, + lam=0.95, + desired_kl=0.01, + max_grad_norm=1.0, + ) + + +@configclass +class DeeproboticsM20FlatPPORunnerCfg(DeeproboticsM20RoughPPORunnerCfg): + def __post_init__(self): + super().__post_init__() + + self.max_iterations = 5000 + self.experiment_name = "deeprobotics_m20_flat" + + +@configclass +class DeeproboticsM20SupportRoughPPORunnerCfg(DeeproboticsM20RoughPPORunnerCfg): + def __post_init__(self): + super().__post_init__() + + self.experiment_name = "deeprobotics_m20_support" diff --git a/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/config/wheeled/deeprobotics_m20/flat_env_cfg.py b/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/config/wheeled/deeprobotics_m20/flat_env_cfg.py new file mode 100644 index 0000000..62b0c84 --- /dev/null +++ b/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/config/wheeled/deeprobotics_m20/flat_env_cfg.py @@ -0,0 +1,32 @@ +# Copyright (c) 2025 Deep Robotics +# SPDX-License-Identifier: BSD 3-Clause + +# Copyright (c) 2024-2025 Ziqi Fan +# SPDX-License-Identifier: Apache-2.0 + +from isaaclab.utils import configclass + +from .rough_env_cfg import DeeproboticsM20RoughEnvCfg + + +@configclass +class DeeproboticsM20FlatEnvCfg(DeeproboticsM20RoughEnvCfg): + def __post_init__(self): + # post init of parent + super().__post_init__() + + # override rewards + self.rewards.base_height_l2.params["sensor_cfg"] = None + # change terrain to flat + self.scene.terrain.terrain_type = "plane" + self.scene.terrain.terrain_generator = None + # no height scan + self.scene.height_scanner = None + self.observations.policy.height_scan = None + self.observations.critic.height_scan = None + # no terrain curriculum + self.curriculum.terrain_levels = None + + # If the weight of rewards is 0, set rewards to None + if self.__class__.__name__ == "DeeproboticsM20FlatEnvCfg": + self.disable_zero_weight_rewards() diff --git a/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/config/wheeled/deeprobotics_m20/rough_env_cfg.py b/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/config/wheeled/deeprobotics_m20/rough_env_cfg.py new file mode 100644 index 0000000..ab051e5 --- /dev/null +++ b/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/config/wheeled/deeprobotics_m20/rough_env_cfg.py @@ -0,0 +1,504 @@ +# Copyright (c) 2025 Deep Robotics +# SPDX-License-Identifier: BSD 3-Clause + +# Copyright (c) 2024-2025 Ziqi Fan +# SPDX-License-Identifier: Apache-2.0 + +from isaaclab.managers import CurriculumTermCfg as CurrTerm +from isaaclab.managers import EventTermCfg as EventTerm +from isaaclab.managers import RewardTermCfg as RewTerm +from isaaclab.managers import ObservationTermCfg as ObsTerm +from isaaclab.managers import SceneEntityCfg +from isaaclab.utils import configclass + +import rl_training.tasks.manager_based.locomotion.velocity.mdp as mdp +from rl_training.tasks.manager_based.locomotion.velocity.velocity_env_cfg import ( + ActionsCfg, + EventCfg, + LocomotionVelocityRoughEnvCfg, + RewardsCfg, +) + +## +# Pre-defined configs +## +from rl_training.assets.deeprobotics import DEEPROBOTICS_M20_CFG # isort: skip + + +@configclass +class DeeproboticsM20ActionsCfg(ActionsCfg): + """Action specifications for the MDP.""" + + joint_pos = mdp.JointPositionActionCfg( + asset_name="robot", joint_names=[""], scale=0.25, use_default_offset=True, clip=None, preserve_order=True + ) + + joint_vel = mdp.JointVelocityActionCfg( + asset_name="robot", joint_names=[""], scale=20.0, use_default_offset=True, clip=None, preserve_order=True + ) + + +@configclass +class DeeproboticsM20EventCfg(EventCfg): + """Domain randomization with tire-specific contact material.""" + + # Replace the inherited all-body material randomization with separate tire and support materials. + randomize_rigid_body_material = None + + randomize_wheel_material = EventTerm( + func=mdp.randomize_rigid_body_material, + mode="startup", + params={ + "asset_cfg": SceneEntityCfg("robot", body_names=".*_wheel"), + "static_friction_range": (0.9, 1.4), + "dynamic_friction_range": (0.6, 0.9), + "restitution_range": (0.0, 0.05), + "num_buckets": 1024, + }, + ) + + randomize_support_material = EventTerm( + func=mdp.randomize_rigid_body_material, + mode="startup", + params={ + "asset_cfg": SceneEntityCfg("robot", body_names=["base_link", ".*_knee"]), + "static_friction_range": (0.4, 0.8), + "dynamic_friction_range": (0.25, 0.4), + "restitution_range": (0.0, 0.05), + "num_buckets": 1024, + }, + ) + + +@configclass +class DeeproboticsM20RewardsCfg(RewardsCfg): + """Reward terms for the MDP.""" + + joint_vel_wheel_l2 = RewTerm( + func=mdp.joint_vel_l2, weight=0.0, params={"asset_cfg": SceneEntityCfg("robot", joint_names="")} + ) + + joint_acc_wheel_l2 = RewTerm( + func=mdp.joint_acc_l2, weight=0.0, params={"asset_cfg": SceneEntityCfg("robot", joint_names="")} + ) + + joint_torques_wheel_l2 = RewTerm( + func=mdp.joint_torques_l2, weight=0.0, params={"asset_cfg": SceneEntityCfg("robot", joint_names="")} + ) + + track_diff_wheel_yaw_exp = RewTerm( + func=mdp.track_diff_wheel_yaw_exp, + weight=1.0, + params={ + "command_name": "base_velocity", + "std": 0.6, + "wheel_radius": 0.096, + "track_width": 0.50, + "left_wheel_indices": (0, 2), + "right_wheel_indices": (1, 3), + "command_threshold": 0.05, + "asset_cfg": SceneEntityCfg("robot", joint_names=""), + }, + ) + + lin_vel_xy_yaw_command_l2 = RewTerm( + func=mdp.lin_vel_xy_yaw_command_l2, + weight=-0.4, + params={ + "command_name": "base_velocity", + "yaw_command_threshold": 0.25, + "lin_command_threshold": 0.15, + "asset_cfg": SceneEntityCfg("robot"), + }, + ) + + command_progress = RewTerm( + func=mdp.command_progress, + weight=0.0, + params={ + "command_name": "base_velocity", + "min_command": 0.12, + "max_reward_speed": 0.8, + "asset_cfg": SceneEntityCfg("robot"), + }, + ) + + support_contact_progress = RewTerm( + func=mdp.ObstacleSupportEscapeReward, + weight=0.0, + params={ + "command_name": "base_velocity", + "sensor_cfg": SceneEntityCfg("contact_forces", body_names=""), + "wheel_sensor_cfg": SceneEntityCfg("contact_forces", body_names=""), + "asset_cfg": SceneEntityCfg("robot"), + "wheel_asset_cfg": SceneEntityCfg("robot", joint_names=""), + "wheel_radius": 0.096, + "support_force_threshold": 15.0, + "wheel_obstacle_force_threshold": 35.0, + "max_contact_time": 0.45, + "min_command": 0.15, + "stall_speed_ratio": 0.35, + "min_stall_speed": 0.12, + "slip_threshold": 0.30, + "obstacle_memory_time": 0.60, + "support_window_time": 0.80, + "recovery_speed_ratio": 0.60, + "min_recovery_speed": 0.18, + "recovery_reward_time": 0.25, + "cooldown_time": 0.60, + "recovery_bonus": 0.25, + "max_reward_speed": 0.8, + }, + ) + + support_contact_excess = RewTerm( + func=mdp.support_contact_excess, + weight=0.0, + params={ + "sensor_cfg": SceneEntityCfg("contact_forces", body_names=""), + "force_threshold": 120.0, + "max_contact_time": 0.6, + }, + ) + + +@configclass +class DeeproboticsM20RoughEnvCfg(LocomotionVelocityRoughEnvCfg): + """Stage 1: blind tactile locomotion without body-support rewards.""" + + actions: DeeproboticsM20ActionsCfg = DeeproboticsM20ActionsCfg() + events: DeeproboticsM20EventCfg = DeeproboticsM20EventCfg() + rewards: DeeproboticsM20RewardsCfg = DeeproboticsM20RewardsCfg() + + base_link_name = "base_link" + foot_link_name = ".*_wheel" + support_link_names = ["base_link", ".*_knee"] + + # fmt: off + leg_joint_names = [ + "fl_knee_joint", + "fr_knee_joint", + "hl_knee_joint", + "hr_knee_joint", + ] + wheel_joint_names = [ + "fl_wheel_joint", "fr_wheel_joint", "hl_wheel_joint", "hr_wheel_joint", + ] + + knee_joint_names = [ + "fl_knee_joint", "fr_knee_joint", "hl_knee_joint", "hr_knee_joint", + ] + joint_names = leg_joint_names + wheel_joint_names + # fmt: on + + def __post_init__(self): + # post init of parent + super().__post_init__() + + # ------------------------------Sence------------------------------ + self.scene.robot = DEEPROBOTICS_M20_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") + self.scene.height_scanner.prim_path = "{ENV_REGEX_NS}/Robot/" + self.base_link_name + self.scene.height_scanner_base.prim_path = "{ENV_REGEX_NS}/Robot/" + self.base_link_name + self.scene.terrain.physics_material.restitution = 0.0 + self.sim.physics_material = self.scene.terrain.physics_material + + # ------------------------------Observations------------------------------ + self.observations.policy.joint_pos.func = mdp.joint_pos_rel_without_wheel + self.observations.policy.joint_pos.params["wheel_asset_cfg"] = SceneEntityCfg( + "robot", joint_names=self.wheel_joint_names + ) + self.observations.critic.joint_pos.func = mdp.joint_pos_rel_without_wheel + self.observations.critic.joint_pos.params["wheel_asset_cfg"] = SceneEntityCfg( + "robot", joint_names=self.wheel_joint_names + ) + self.observations.policy.base_lin_vel.scale = 2.0 + self.observations.policy.base_ang_vel.scale = 0.25 + self.observations.policy.joint_pos.scale = 1.0 + self.observations.policy.joint_vel.scale = 0.05 + self.observations.policy.base_lin_vel = None + self.observations.policy.height_scan = None + self.observations.policy.joint_pos.params["asset_cfg"].joint_names = self.joint_names + self.observations.policy.joint_vel.params["asset_cfg"].joint_names = self.joint_names + self.observations.policy.wheel_contact = ObsTerm( + func=mdp.contact_state, + params={ + "sensor_cfg": SceneEntityCfg("contact_forces", body_names=[self.foot_link_name]), + "threshold": 2.0, + }, + clip=(0.0, 1.0), + scale=1.0, + ) + self.observations.policy.support_contact = ObsTerm( + func=mdp.contact_state, + params={ + "sensor_cfg": SceneEntityCfg("contact_forces", body_names=self.support_link_names), + "threshold": 2.0, + }, + clip=(0.0, 1.0), + scale=1.0, + ) + self.observations.policy.wheel_contact_force = ObsTerm( + func=mdp.contact_force_norm, + params={ + "sensor_cfg": SceneEntityCfg("contact_forces", body_names=[self.foot_link_name]), + "scale": 100.0, + "clip": 2.0, + }, + clip=(0.0, 2.0), + scale=1.0, + ) + self.observations.policy.support_contact_force = ObsTerm( + func=mdp.contact_force_norm, + params={ + "sensor_cfg": SceneEntityCfg("contact_forces", body_names=self.support_link_names), + "scale": 100.0, + "clip": 2.0, + }, + clip=(0.0, 2.0), + scale=1.0, + ) + self.observations.policy.wheel_slip = ObsTerm( + func=mdp.wheel_slip_proxy, + params={ + "asset_cfg": SceneEntityCfg("robot", joint_names=self.wheel_joint_names, preserve_order=True), + "wheel_radius": 0.096, + }, + clip=(0.0, 3.0), + scale=0.5, + ) + self.observations.policy.applied_torque = ObsTerm( + func=mdp.applied_torque, + params={ + "asset_cfg": SceneEntityCfg("robot", joint_names=self.joint_names, preserve_order=True), + "scale": 80.0, + "clip": 2.0, + }, + clip=(-2.0, 2.0), + scale=1.0, + ) + + # ------------------------------Actions------------------------------ + # reduce action scale + self.actions.joint_pos.scale = 0.25 + self.actions.joint_vel.scale = 5.0 + self.actions.joint_pos.clip = {".*": (-100.0, 100.0)} + self.actions.joint_vel.clip = {".*": (-100.0, 100.0)} + self.actions.joint_pos.joint_names = self.leg_joint_names + self.actions.joint_vel.joint_names = self.wheel_joint_names + + # ------------------------------Events------------------------------ + self.events.randomize_reset_base.params = { + "pose_range": { + "x": (-1.0, 1.0), + "y": (-1.0, 1.0), + "z": (0.0, 0.0), + "roll": (-0.3, 0.3), + "pitch": (-0.3, 0.3), + "yaw": (-3.14, 3.14), + }, + "velocity_range": { + "x": (-0.2, 0.2), + "y": (-0.2, 0.2), + "z": (-0.2, 0.2), + "roll": (-0.05, 0.05), + "pitch": (-0.05, 0.05), + "yaw": (-0.0, 0.0), + }, + } + self.events.randomize_rigid_body_mass_base.params["asset_cfg"].body_names = [self.base_link_name] + self.events.randomize_rigid_body_mass.params["asset_cfg"].body_names = [ + f"^(?!.*{self.base_link_name}).*" + ] + self.events.randomize_com_positions.params["asset_cfg"].body_names = [self.base_link_name] + self.events.randomize_apply_external_force_torque.params["asset_cfg"].body_names = [self.base_link_name] + + self.scene.terrain.max_init_terrain_level = 2 + self.scene.terrain.terrain_generator.sub_terrains["boxes"].grid_height_range = (0.02, 0.14) + self.scene.terrain.terrain_generator.sub_terrains["pyramid_stairs"].step_height_range = (0.03, 0.16) + self.scene.terrain.terrain_generator.sub_terrains["pyramid_stairs_inv"].step_height_range = (0.03, 0.16) + self.scene.terrain.terrain_generator.sub_terrains["random_rough"].noise_range = (0.01, 0.10) + self.scene.terrain.terrain_generator.sub_terrains["random_rough"].noise_step = 0.01 + + # only randomize actuator gains for actuated joints (knees + wheels), exclude hipx/hipy completely + self.events.randomize_actuator_gains.params["asset_cfg"].joint_names = self.joint_names + + # ------------------------------Rewards------------------------------ + # General + self.rewards.is_terminated.weight = 0 + + # Root penalties + self.rewards.lin_vel_z_l2.weight = -2.0 + self.rewards.ang_vel_xy_l2.weight = -0.02 + self.rewards.flat_orientation_l2.weight = -0.5 + self.rewards.base_height_l2.weight = -0.5 + self.rewards.base_height_l2.params["target_height"] = 0.45 + self.rewards.base_height_l2.params["asset_cfg"].body_names = [self.base_link_name] + self.rewards.body_lin_acc_l2.weight = 0 + self.rewards.body_lin_acc_l2.params["asset_cfg"].body_names = [self.base_link_name] + + # Joint penalties + self.rewards.joint_torques_l2.weight = -2.5e-5 + self.rewards.joint_torques_l2.params["asset_cfg"].joint_names = self.leg_joint_names + self.rewards.joint_torques_wheel_l2.weight = 0 + self.rewards.joint_torques_wheel_l2.params["asset_cfg"].joint_names = self.wheel_joint_names + self.rewards.joint_vel_l2.weight = 0 + self.rewards.joint_vel_l2.params["asset_cfg"].joint_names = self.leg_joint_names + self.rewards.joint_vel_wheel_l2.weight = 0 + self.rewards.joint_vel_wheel_l2.params["asset_cfg"].joint_names = self.wheel_joint_names + self.rewards.joint_acc_l2.weight = -2e-7 + self.rewards.joint_acc_l2.params["asset_cfg"].joint_names = self.leg_joint_names + self.rewards.joint_acc_wheel_l2.weight = -1e-7 + self.rewards.joint_acc_wheel_l2.params["asset_cfg"].joint_names = self.wheel_joint_names + # self.rewards.create_joint_deviation_l1_rewterm("joint_deviation_hip_l1", -0.2, [".*_hip_joint"]) + self.rewards.joint_pos_limits.weight = -5.0 + self.rewards.joint_pos_limits.params["asset_cfg"].joint_names = self.leg_joint_names + self.rewards.joint_vel_limits.weight = 0 + self.rewards.joint_vel_limits.params["asset_cfg"].joint_names = self.wheel_joint_names + self.rewards.joint_power.weight = -2e-5 + self.rewards.joint_power.params["asset_cfg"].joint_names = self.leg_joint_names + self.rewards.stand_still.weight = -2.0 + self.rewards.stand_still.params["asset_cfg"].joint_names = self.leg_joint_names + + self.rewards.knee_joint_pos_penalty.weight = -0.1 + self.rewards.knee_joint_pos_penalty.params["asset_cfg"].joint_names = self.knee_joint_names + self.rewards.wheel_vel_penalty.weight = 0 + self.rewards.wheel_vel_penalty.params["sensor_cfg"].body_names = self.foot_link_name + self.rewards.wheel_vel_penalty.params["asset_cfg"].joint_names = self.wheel_joint_names + self.rewards.joint_mirror.weight = -0.03 + self.rewards.joint_mirror.params["mirror_joints"] = [ + ["fl_knee_joint", "hr_knee_joint"], + ["fr_knee_joint", "hl_knee_joint"], + ] + + # Action penalties + self.rewards.action_rate_l2.weight = -0.01 + + # Contact sensor + self.rewards.undesired_contacts.weight = -1.0 + self.rewards.undesired_contacts.params["sensor_cfg"].body_names = [f"^(?!.*{self.foot_link_name}).*"] + self.rewards.undesired_contacts.params["threshold"] = 1.0 + self.rewards.contact_forces.weight = -1.5e-4 + self.rewards.contact_forces.params["sensor_cfg"].body_names = [self.foot_link_name] + self.rewards.support_contact_progress.weight = 0.0 + self.rewards.support_contact_progress.params["sensor_cfg"].body_names = self.support_link_names + self.rewards.support_contact_progress.params["wheel_sensor_cfg"].body_names = [self.foot_link_name] + self.rewards.support_contact_progress.params["wheel_asset_cfg"].joint_names = self.wheel_joint_names + self.rewards.support_contact_excess.weight = 0.0 + self.rewards.support_contact_excess.params["sensor_cfg"].body_names = self.support_link_names + + # Velocity-tracking rewards + self.rewards.track_lin_vel_xy_exp.weight = 2.0 # 1.8 + self.rewards.track_ang_vel_z_exp.weight = 1.5 # 1.2 + self.rewards.track_ang_vel_z_exp.params["std"] = 0.5 + self.rewards.track_diff_wheel_yaw_exp.weight = 1.0 + self.rewards.track_diff_wheel_yaw_exp.params["asset_cfg"].joint_names = self.wheel_joint_names + self.rewards.lin_vel_xy_yaw_command_l2.weight = -0.4 + + # Others + self.rewards.feet_air_time.weight = 0 + self.rewards.feet_air_time.params["threshold"] = 0.5 + self.rewards.feet_air_time.params["sensor_cfg"].body_names = [self.foot_link_name] + self.rewards.feet_contact.weight = 0 + self.rewards.feet_contact.params["sensor_cfg"].body_names = [self.foot_link_name] + self.rewards.feet_contact_without_cmd.weight = 0.1 + self.rewards.feet_contact_without_cmd.params["sensor_cfg"].body_names = [self.foot_link_name] + self.rewards.feet_stumble.weight = 0 + self.rewards.feet_stumble.params["sensor_cfg"].body_names = [self.foot_link_name] + self.rewards.feet_slide.weight = 0 + self.rewards.feet_slide.params["sensor_cfg"].body_names = [self.foot_link_name] + self.rewards.feet_slide.params["asset_cfg"].body_names = [self.foot_link_name] + self.rewards.feet_height.weight = 0 + self.rewards.feet_height.params["target_height"] = 0.1 + self.rewards.feet_height.params["asset_cfg"].body_names = [self.foot_link_name] + self.rewards.feet_height_body.weight = 0 + self.rewards.feet_height_body.params["target_height"] = -0.4 + self.rewards.feet_height_body.params["asset_cfg"].body_names = [self.foot_link_name] + self.rewards.feet_gait.weight = 0 + self.rewards.feet_gait.params["synced_feet_pair_names"] = (("fl_wheel", "hr_wheel"), ("fr_wheel", "hl_wheel")) + self.rewards.upward.weight = 0.08 + + # If the weight of rewards is 0, set rewards to None + if self.__class__.__name__ == "DeeproboticsM20RoughEnvCfg": + self.disable_zero_weight_rewards() + + # ------------------------------Terminations------------------------------ + # self.terminations.illegal_contact.params["sensor_cfg"].body_names = [self.base_link_name] + self.terminations.illegal_contact = None + self.terminations.bad_orientation_2 = None + + # ------------------------------Curriculums------------------------------ + # self.curriculum.command_levels.params["range_multiplier"] = (0.2, 1.0) + self.curriculum.command_levels = None + + # ------------------------------Commands------------------------------ + self.commands.base_velocity.debug_vis = False + self.commands.base_velocity.ranges.lin_vel_x = (-2.0, 2.0) + self.commands.base_velocity.ranges.lin_vel_y = (0.0, 0.0) + self.commands.base_velocity.ranges.ang_vel_z = (-1.5, 1.5) + + +@configclass +class DeeproboticsM20SupportRoughEnvCfg(DeeproboticsM20RoughEnvCfg): + """Stage 2: resume from Stage 1 and enable brief body/knee support for obstacle crossing.""" + + def __post_init__(self): + super().__post_init__() + + # Keep the same policy/action dimensions as Stage 1 for checkpoint resume. + # Only rewards, command distribution, and terrain curriculum are changed here. + self.scene.terrain.max_init_terrain_level = 1 + self.scene.terrain.terrain_generator.sub_terrains["boxes"].proportion = 0.25 + self.scene.terrain.terrain_generator.sub_terrains["pyramid_stairs"].proportion = 0.25 + self.scene.terrain.terrain_generator.sub_terrains["pyramid_stairs_inv"].proportion = 0.15 + self.scene.terrain.terrain_generator.sub_terrains["random_rough"].proportion = 0.15 + self.scene.terrain.terrain_generator.sub_terrains["hf_pyramid_slope"].proportion = 0.10 + self.scene.terrain.terrain_generator.sub_terrains["hf_pyramid_slope_inv"].proportion = 0.10 + self.scene.terrain.terrain_generator.sub_terrains["boxes"].grid_height_range = (0.03, 0.22) + self.scene.terrain.terrain_generator.sub_terrains["pyramid_stairs"].step_height_range = (0.04, 0.24) + self.scene.terrain.terrain_generator.sub_terrains["pyramid_stairs_inv"].step_height_range = (0.04, 0.20) + self.scene.terrain.terrain_generator.sub_terrains["random_rough"].noise_range = (0.01, 0.14) + + # Keep privileged terrain information for the critic only. The actor remains blind/tactile. + self.observations.policy.height_scan = None + self.observations.critic.height_scan.params["sensor_cfg"] = SceneEntityCfg("height_scanner") + self.observations.critic.height_scan.clip = (-1.0, 1.0) + self.observations.critic.height_scan.scale = 1.0 + + # Body/knee contact is allowed only as short support during forward progress. + self.rewards.undesired_contacts.weight = -0.10 + self.rewards.undesired_contacts.params["sensor_cfg"].body_names = [f"^(?!.*{self.foot_link_name}).*"] + self.rewards.undesired_contacts.params["threshold"] = 35.0 + self.rewards.command_progress.weight = 0.8 + self.rewards.support_contact_progress.weight = 1.0 + self.rewards.support_contact_progress.params["sensor_cfg"].body_names = self.support_link_names + self.rewards.support_contact_progress.params["wheel_sensor_cfg"].body_names = [self.foot_link_name] + self.rewards.support_contact_progress.params["wheel_asset_cfg"].joint_names = self.wheel_joint_names + self.rewards.support_contact_excess.weight = -0.60 + self.rewards.support_contact_excess.params["sensor_cfg"].body_names = self.support_link_names + self.rewards.support_contact_excess.params["force_threshold"] = 100.0 + self.rewards.support_contact_excess.params["max_contact_time"] = 0.45 + + # Stage 2 should still avoid the bad solution of crawling on the body. + self.rewards.track_lin_vel_xy_exp.weight = 1.5 + self.rewards.track_ang_vel_z_exp.weight = 1.0 + self.rewards.flat_orientation_l2.weight = -0.45 + self.rewards.base_height_l2.weight = -0.35 + self.rewards.base_height_l2.params["target_height"] = 0.29 + + # Stage 2 first focuses on forward obstacle crossing, then expands the speed range. + self.commands.base_velocity.heading_command = False + self.commands.base_velocity.rel_standing_envs = 0.0 + self.commands.base_velocity.rel_heading_envs = 0.0 + self.commands.base_velocity.ranges.lin_vel_x = (0.25, 1.2) + self.commands.base_velocity.ranges.lin_vel_y = (0.0, 0.0) + self.commands.base_velocity.ranges.ang_vel_z = (-0.6, 0.6) + self.curriculum.command_levels = CurrTerm( + func=mdp.command_levels_positive_x, + params={ + "reward_term_name": "track_lin_vel_xy_exp", + "initial_max_ratio": 0.35, + "increment": 0.1, + }, + ) + + self.disable_zero_weight_rewards() diff --git a/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/mdp/__init__.py b/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/mdp/__init__.py new file mode 100644 index 0000000..af6c758 --- /dev/null +++ b/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/mdp/__init__.py @@ -0,0 +1,21 @@ +# Copyright (c) 2025 Deep Robotics +# SPDX-License-Identifier: BSD 3-Clause +# +# # Copyright (c) 2024-2025 Ziqi Fan +# SPDX-License-Identifier: Apache-2.0 + +# Copyright (c) 2022-2025, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""This sub-module contains the functions that are specific to the locomotion environments.""" + +from isaaclab.envs.mdp import * # noqa: F401, F403 +from isaaclab_tasks.manager_based.locomotion.velocity.mdp import * # noqa: F401, F403 + +from .commands import * # noqa: F401, F403 +from .curriculums import * # noqa: F401, F403 +from .events import * # noqa: F401, F403 +from .observations import * # noqa: F401, F403 +from .rewards import * # noqa: F401, F403 diff --git a/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/mdp/commands.py b/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/mdp/commands.py new file mode 100644 index 0000000..c94fdf6 --- /dev/null +++ b/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/mdp/commands.py @@ -0,0 +1,130 @@ +# Copyright (c) 2025 Deep Robotics +# SPDX-License-Identifier: BSD 3-Clause +# +# # Copyright (c) 2024-2025 Ziqi Fan +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import torch +from typing import TYPE_CHECKING, Sequence + +from isaaclab.managers import CommandTerm, CommandTermCfg +from isaaclab.utils import configclass + +import rl_training.tasks.manager_based.locomotion.velocity.mdp as mdp + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedEnv + + +class UniformThresholdVelocityCommand(mdp.UniformVelocityCommand): + """Command generator that generates a velocity command in SE(2) from uniform distribution with threshold.""" + + cfg: mdp.UniformThresholdVelocityCommandCfg + """The configuration of the command generator.""" + + def _resample_command(self, env_ids: Sequence[int]): + super()._resample_command(env_ids) + # set small commands to zero + self.vel_command_b[env_ids, :2] *= (torch.norm(self.vel_command_b[env_ids, :2], dim=1) > 0.2).unsqueeze(1) + + +@configclass +class UniformThresholdVelocityCommandCfg(mdp.UniformVelocityCommandCfg): + """Configuration for the uniform threshold velocity command generator.""" + + class_type: type = UniformThresholdVelocityCommand + + +class DiscreteCommandController(CommandTerm): + """ + Command generator that assigns discrete commands to environments. + + Commands are stored as a list of predefined integers. + The controller maps these commands by their indices (e.g., index 0 -> 10, index 1 -> 20). + """ + + cfg: DiscreteCommandControllerCfg + """Configuration for the command controller.""" + + def __init__(self, cfg: DiscreteCommandControllerCfg, env: ManagerBasedEnv): + """ + Initialize the command controller. + + Args: + cfg: The configuration of the command controller. + env: The environment object. + """ + # Initialize the base class + super().__init__(cfg, env) + + # Validate that available_commands is non-empty + if not self.cfg.available_commands: + raise ValueError("The available_commands list cannot be empty.") + + # Ensure all elements are integers + if not all(isinstance(cmd, int) for cmd in self.cfg.available_commands): + raise ValueError("All elements in available_commands must be integers.") + + # Store the available commands + self.available_commands = self.cfg.available_commands + + # Create buffers to store the command + # -- command buffer: stores discrete action indices for each environment + self.command_buffer = torch.zeros(self.num_envs, dtype=torch.int32, device=self.device) + + # -- current_commands: stores a snapshot of the current commands (as integers) + self.current_commands = [self.available_commands[0]] * self.num_envs # Default to the first command + + def __str__(self) -> str: + """Return a string representation of the command controller.""" + return ( + "DiscreteCommandController:\n" + f"\tNumber of environments: {self.num_envs}\n" + f"\tAvailable commands: {self.available_commands}\n" + ) + + """ + Properties + """ + + @property + def command(self) -> torch.Tensor: + """Return the current command buffer. Shape is (num_envs, 1).""" + return self.command_buffer + + """ + Implementation specific functions. + """ + + def _update_metrics(self): + """Update metrics for the command controller.""" + pass + + def _resample_command(self, env_ids: Sequence[int]): + """Resample commands for the given environments.""" + sampled_indices = torch.randint( + len(self.available_commands), (len(env_ids),), dtype=torch.int32, device=self.device + ) + sampled_commands = torch.tensor( + [self.available_commands[idx.item()] for idx in sampled_indices], dtype=torch.int32, device=self.device + ) + self.command_buffer[env_ids] = sampled_commands + + def _update_command(self): + """Update and store the current commands.""" + self.current_commands = self.command_buffer.tolist() + + +@configclass +class DiscreteCommandControllerCfg(CommandTermCfg): + """Configuration for the discrete command controller.""" + + class_type: type = DiscreteCommandController + + available_commands: list[int] = [] + """ + List of available discrete commands, where each element is an integer. + Example: [10, 20, 30, 40, 50] + """ diff --git a/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/mdp/curriculums.py b/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/mdp/curriculums.py new file mode 100644 index 0000000..a5b0a8f --- /dev/null +++ b/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/mdp/curriculums.py @@ -0,0 +1,95 @@ +# Copyright (c) 2025 Deep Robotics +# SPDX-License-Identifier: BSD 3-Clause +# +# # Copyright (c) 2024-2025 Ziqi Fan +# SPDX-License-Identifier: Apache-2.0 + +"""Common functions that can be used to create curriculum for the learning environment. + +The functions can be passed to the :class:`isaaclab.managers.CurriculumTermCfg` object to enable +the curriculum introduced by the function. +""" + +from __future__ import annotations + +import torch +from collections.abc import Sequence +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedRLEnv + + +def command_levels_vel( + env: ManagerBasedRLEnv, + env_ids: Sequence[int], + reward_term_name: str, + range_multiplier: Sequence[float] = (0.1, 1.0), +) -> None: + """command_levels_vel""" + base_velocity_ranges = env.command_manager.get_term("base_velocity").cfg.ranges + # Get original velocity ranges (ONLY ON FIRST EPISODE) + if env.common_step_counter == 0: + env._original_vel_x = torch.tensor(base_velocity_ranges.lin_vel_x, device=env.device) + env._original_vel_y = torch.tensor(base_velocity_ranges.lin_vel_y, device=env.device) + env._initial_vel_x = env._original_vel_x * range_multiplier[0] + env._final_vel_x = env._original_vel_x * range_multiplier[1] + env._initial_vel_y = env._original_vel_y * range_multiplier[0] + env._final_vel_y = env._original_vel_y * range_multiplier[1] + + # Initialize command ranges to initial values + base_velocity_ranges.lin_vel_x = env._initial_vel_x.tolist() + base_velocity_ranges.lin_vel_y = env._initial_vel_y.tolist() + + # avoid updating command curriculum at each step since the maximum command is common to all envs + if env.common_step_counter % env.max_episode_length == 0: + episode_sums = env.reward_manager._episode_sums[reward_term_name] + reward_term_cfg = env.reward_manager.get_term_cfg(reward_term_name) + delta_command = torch.tensor([-0.1, 0.1], device=env.device) + + # If the tracking reward is above 80% of the maximum, increase the range of commands + if torch.mean(episode_sums[env_ids]) / env.max_episode_length_s > 0.8 * reward_term_cfg.weight: + new_vel_x = torch.tensor(base_velocity_ranges.lin_vel_x, device=env.device) + delta_command + new_vel_y = torch.tensor(base_velocity_ranges.lin_vel_y, device=env.device) + delta_command + + # Clamp to ensure we don't exceed final ranges + new_vel_x = torch.clamp(new_vel_x, min=env._final_vel_x[0], max=env._final_vel_x[1]) + new_vel_y = torch.clamp(new_vel_y, min=env._final_vel_y[0], max=env._final_vel_y[1]) + + # Update ranges + base_velocity_ranges.lin_vel_x = new_vel_x.tolist() + base_velocity_ranges.lin_vel_y = new_vel_y.tolist() + + return torch.tensor(base_velocity_ranges.lin_vel_x[1], device=env.device) + + +def command_levels_positive_x( + env: ManagerBasedRLEnv, + env_ids: Sequence[int], + reward_term_name: str, + initial_max_ratio: float = 0.35, + increment: float = 0.1, +) -> torch.Tensor: + """Expand only the upper bound of a positive forward-speed command range. + + Unlike :func:`command_levels_vel`, this curriculum preserves the configured minimum forward speed. This avoids + producing near-zero commands in obstacle-focused stages where the policy needs enough momentum to learn climbing. + """ + base_velocity_ranges = env.command_manager.get_term("base_velocity").cfg.ranges + if env.common_step_counter == 0: + original_x = torch.tensor(base_velocity_ranges.lin_vel_x, device=env.device) + env._positive_x_min = original_x[0] + env._positive_x_max = original_x[1] + initial_x_max = torch.maximum(env._positive_x_min, env._positive_x_max * initial_max_ratio) + base_velocity_ranges.lin_vel_x = [env._positive_x_min.item(), initial_x_max.item()] + + if env.common_step_counter % env.max_episode_length == 0: + episode_sums = env.reward_manager._episode_sums[reward_term_name] + reward_term_cfg = env.reward_manager.get_term_cfg(reward_term_name) + mean_reward = torch.mean(episode_sums[env_ids]) / env.max_episode_length_s + if mean_reward > 0.8 * reward_term_cfg.weight: + current_x_max = torch.tensor(base_velocity_ranges.lin_vel_x[1], device=env.device) + new_x_max = torch.clamp(current_x_max + increment, max=env._positive_x_max) + base_velocity_ranges.lin_vel_x = [env._positive_x_min.item(), new_x_max.item()] + + return torch.tensor(base_velocity_ranges.lin_vel_x[1], device=env.device) diff --git a/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/mdp/events.py b/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/mdp/events.py new file mode 100644 index 0000000..a307b8e --- /dev/null +++ b/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/mdp/events.py @@ -0,0 +1,213 @@ +# Copyright (c) 2025 Deep Robotics +# SPDX-License-Identifier: BSD 3-Clause +# +# # Copyright (c) 2024-2025 Ziqi Fan +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import torch +from typing import TYPE_CHECKING, Literal + +import isaaclab.utils.math as math_utils +from isaaclab.assets import Articulation, RigidObject +from isaaclab.managers import SceneEntityCfg + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedEnv + + +def randomize_rigid_body_inertia( + env: ManagerBasedEnv, + env_ids: torch.Tensor | None, + asset_cfg: SceneEntityCfg, + inertia_distribution_params: tuple[float, float], + operation: Literal["add", "scale", "abs"], + distribution: Literal["uniform", "log_uniform", "gaussian"] = "uniform", +): + """Randomize the inertia tensors of the bodies by adding, scaling, or setting random values. + + This function allows randomizing only the diagonal inertia tensor components (xx, yy, zz) of the bodies. + The function samples random values from the given distribution parameters and adds, scales, or sets the values + into the physics simulation based on the operation. + + .. tip:: + This function uses CPU tensors to assign the body inertias. It is recommended to use this function + only during the initialization of the environment. + """ + # extract the used quantities (to enable type-hinting) + asset: RigidObject | Articulation = env.scene[asset_cfg.name] + + # resolve environment ids + if env_ids is None: + env_ids = torch.arange(env.scene.num_envs, device="cpu") + else: + env_ids = env_ids.cpu() + + # resolve body indices + if asset_cfg.body_ids == slice(None): + body_ids = torch.arange(asset.num_bodies, dtype=torch.int, device="cpu") + else: + body_ids = torch.tensor(asset_cfg.body_ids, dtype=torch.int, device="cpu") + + # get the current inertia tensors of the bodies (num_assets, num_bodies, 9 for articulations or 9 for rigid objects) + inertias = asset.root_physx_view.get_inertias() + + # apply randomization on default values + inertias[env_ids[:, None], body_ids, :] = asset.data.default_inertia[env_ids[:, None], body_ids, :].clone() + + # randomize each diagonal element (xx, yy, zz -> indices 0, 4, 8) + for idx in [0, 4, 8]: + # Extract and randomize the specific diagonal element + randomized_inertias = _randomize_prop_by_op( + inertias[:, :, idx], + inertia_distribution_params, + env_ids, + body_ids, + operation, + distribution, + ) + # Assign the randomized values back to the inertia tensor + inertias[env_ids[:, None], body_ids, idx] = randomized_inertias + + # set the inertia tensors into the physics simulation + asset.root_physx_view.set_inertias(inertias, env_ids) + + +def randomize_com_positions( + env: ManagerBasedEnv, + env_ids: torch.Tensor | None, + asset_cfg: SceneEntityCfg, + com_distribution_params: tuple[float, float], + operation: Literal["add", "scale", "abs"], + distribution: Literal["uniform", "log_uniform", "gaussian"] = "uniform", +): + """Randomize the center of mass (COM) positions for the rigid bodies. + + This function allows randomizing the COM positions of the bodies in the physics simulation. The positions can be + randomized by adding, scaling, or setting random values sampled from the specified distribution. + + .. tip:: + This function is intended for initialization or offline adjustments, as it modifies physics properties directly. + + Args: + env (ManagerBasedEnv): The simulation environment. + env_ids (torch.Tensor | None): Specific environment indices to apply randomization, or None for all environments. + asset_cfg (SceneEntityCfg): The configuration for the target asset whose COM will be randomized. + com_distribution_params (tuple[float, float]): Parameters of the distribution (e.g., min and max for uniform). + operation (Literal["add", "scale", "abs"]): The operation to apply for randomization. + distribution (Literal["uniform", "log_uniform", "gaussian"]): The distribution to sample random values from. + """ + # Extract the asset (Articulation or RigidObject) + asset: RigidObject | Articulation = env.scene[asset_cfg.name] + + # Resolve environment indices + if env_ids is None: + env_ids = torch.arange(env.scene.num_envs, device="cpu") + else: + env_ids = env_ids.cpu() + + # Resolve body indices + if asset_cfg.body_ids == slice(None): + body_ids = torch.arange(asset.num_bodies, dtype=torch.int, device="cpu") + else: + body_ids = torch.tensor(asset_cfg.body_ids, dtype=torch.int, device="cpu") + + # Get the current COM offsets (num_assets, num_bodies, 3) + com_offsets = asset.root_physx_view.get_coms() + + for dim_idx in range(3): # Randomize x, y, z independently + randomized_offset = _randomize_prop_by_op( + com_offsets[:, :, dim_idx], + com_distribution_params, + env_ids, + body_ids, + operation, + distribution, + ) + com_offsets[env_ids[:, None], body_ids, dim_idx] = randomized_offset[env_ids[:, None], body_ids] + + # Set the randomized COM offsets into the simulation + asset.root_physx_view.set_coms(com_offsets, env_ids) + + +""" +Internal helper functions. +""" + + +def _randomize_prop_by_op( + data: torch.Tensor, + distribution_parameters: tuple[float | torch.Tensor, float | torch.Tensor], + dim_0_ids: torch.Tensor | None, + dim_1_ids: torch.Tensor | slice, + operation: Literal["add", "scale", "abs"], + distribution: Literal["uniform", "log_uniform", "gaussian"], +) -> torch.Tensor: + """Perform data randomization based on the given operation and distribution. + + Args: + data: The data tensor to be randomized. Shape is (dim_0, dim_1). + distribution_parameters: The parameters for the distribution to sample values from. + dim_0_ids: The indices of the first dimension to randomize. + dim_1_ids: The indices of the second dimension to randomize. + operation: The operation to perform on the data. Options: 'add', 'scale', 'abs'. + distribution: The distribution to sample the random values from. Options: 'uniform', 'log_uniform'. + + Returns: + The data tensor after randomization. Shape is (dim_0, dim_1). + + Raises: + NotImplementedError: If the operation or distribution is not supported. + """ + # resolve shape + # -- dim 0 + if dim_0_ids is None: + n_dim_0 = data.shape[0] + dim_0_ids = slice(None) # type: ignore + else: + n_dim_0 = len(dim_0_ids) + if not isinstance(dim_1_ids, slice): + dim_0_ids = dim_0_ids[:, None] + # -- dim 1 + if isinstance(dim_1_ids, slice): + n_dim_1 = data.shape[1] + else: + n_dim_1 = len(dim_1_ids) + + # resolve the distribution + if distribution == "uniform": + dist_fn = math_utils.sample_uniform + elif distribution == "log_uniform": + dist_fn = math_utils.sample_log_uniform + elif distribution == "gaussian": + dist_fn = math_utils.sample_gaussian + else: + raise NotImplementedError( + f"Unknown distribution: '{distribution}' for joint properties randomization." + " Please use 'uniform', 'log_uniform', 'gaussian'." + ) + # perform the operation + if operation == "add": + data[dim_0_ids, dim_1_ids] += dist_fn(*distribution_parameters, (n_dim_0, n_dim_1), device=data.device) + elif operation == "scale": + data[dim_0_ids, dim_1_ids] *= dist_fn(*distribution_parameters, (n_dim_0, n_dim_1), device=data.device) + elif operation == "abs": + data[dim_0_ids, dim_1_ids] = dist_fn(*distribution_parameters, (n_dim_0, n_dim_1), device=data.device) + else: + raise NotImplementedError( + f"Unknown operation: '{operation}' for property randomization. Please use 'add', 'scale', or 'abs'." + ) + return data + + +def bad_orientation_2( + env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot") # type: ignore +) -> torch.Tensor: + """Terminate when the asset's orientation is too far from the desired orientation limits. + + This is computed by checking the angle between the projected gravity vector and the z-axis. + """ + # extract the used quantities (to enable type-hinting) + asset: RigidObject = env.scene[asset_cfg.name] + return (asset.data.projected_gravity_b[:, 2] > 0) | (asset.data.projected_gravity_b[:, :2].abs() > 0.7).any(-1) \ No newline at end of file diff --git a/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/mdp/observations.py b/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/mdp/observations.py new file mode 100644 index 0000000..c9d9681 --- /dev/null +++ b/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/mdp/observations.py @@ -0,0 +1,105 @@ +# Copyright (c) 2025 Deep Robotics +# SPDX-License-Identifier: BSD 3-Clause +# +# # Copyright (c) 2024-2025 Ziqi Fan +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import torch +from typing import TYPE_CHECKING + +from isaaclab.assets import Articulation +from isaaclab.managers import SceneEntityCfg +from isaaclab.sensors import ContactSensor + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedEnv, ManagerBasedRLEnv + + +# def joint_pos_rel_without_wheel( +# env: ManagerBasedEnv, +# asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), +# wheel_asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), +# ) -> torch.Tensor: +# """The joint positions of the asset w.r.t. the default joint positions.(Without the wheel joints)""" +# # extract the used quantities (to enable type-hinting) +# asset: Articulation = env.scene[asset_cfg.name] +# joint_pos_rel = asset.data.joint_pos[:, asset_cfg.joint_ids] - asset.data.default_joint_pos[:, asset_cfg.joint_ids] +# joint_pos_rel[:, wheel_asset_cfg.joint_ids] = 0 +# return joint_pos_rel + +def joint_pos_rel_without_wheel( + env: ManagerBasedEnv, + asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), + wheel_asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), +) -> torch.Tensor: + + asset: Articulation = env.scene[asset_cfg.name] + + # 1. 先在“完整 joint 空间”算相对位置 + joint_pos_rel_full = asset.data.joint_pos - asset.data.default_joint_pos + + # 2. 把 wheel joint 清零（这里的 index 是全局的，合法） + joint_pos_rel_full[:, wheel_asset_cfg.joint_ids] = 0.0 + + # 3. 最后再裁剪成 policy 需要的 joint + joint_pos_rel = joint_pos_rel_full[:, asset_cfg.joint_ids] + + return joint_pos_rel + + + +def phase(env: ManagerBasedRLEnv, cycle_time: float) -> torch.Tensor: + if not hasattr(env, "episode_length_buf") or env.episode_length_buf is None: + env.episode_length_buf = torch.zeros(env.num_envs, device=env.device, dtype=torch.long) + phase = env.episode_length_buf[:, None] * env.step_dt / cycle_time + phase_tensor = torch.cat([torch.sin(2 * torch.pi * phase), torch.cos(2 * torch.pi * phase)], dim=-1) + return phase_tensor + + +def contact_state( + env: ManagerBasedRLEnv, + sensor_cfg: SceneEntityCfg, + threshold: float = 1.0, +) -> torch.Tensor: + """Binary contact state for selected bodies.""" + contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name] + force_norm = torch.linalg.norm(contact_sensor.data.net_forces_w[:, sensor_cfg.body_ids, :], dim=-1) + return (force_norm > threshold).float() + + +def contact_force_norm( + env: ManagerBasedRLEnv, + sensor_cfg: SceneEntityCfg, + scale: float = 100.0, + clip: float = 2.0, +) -> torch.Tensor: + """Normalized contact force magnitude for selected bodies.""" + contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name] + force_norm = torch.linalg.norm(contact_sensor.data.net_forces_w[:, sensor_cfg.body_ids, :], dim=-1) + return torch.clamp(force_norm / scale, min=0.0, max=clip) + + +def wheel_slip_proxy( + env: ManagerBasedRLEnv, + wheel_radius: float, + asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), +) -> torch.Tensor: + """A proprioceptive slip proxy from wheel peripheral speed and base forward speed.""" + asset: Articulation = env.scene[asset_cfg.name] + wheel_speed = torch.abs(asset.data.joint_vel[:, asset_cfg.joint_ids] * wheel_radius) + base_forward_speed = torch.abs(asset.data.root_lin_vel_b[:, 0:1]) + return torch.clamp(torch.abs(wheel_speed - base_forward_speed), min=0.0, max=3.0) + + +def applied_torque( + env: ManagerBasedRLEnv, + asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), + scale: float = 80.0, + clip: float = 2.0, +) -> torch.Tensor: + """Normalized applied motor torques.""" + asset: Articulation = env.scene[asset_cfg.name] + torque = asset.data.applied_torque[:, asset_cfg.joint_ids] + return torch.clamp(torque / scale, min=-clip, max=clip) diff --git a/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/mdp/rewards.py b/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/mdp/rewards.py new file mode 100644 index 0000000..56638e9 --- /dev/null +++ b/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/mdp/rewards.py @@ -0,0 +1,985 @@ +# Copyright (c) 2025 Deep Robotics +# SPDX-License-Identifier: BSD 3-Clause + +# Copyright (c) 2024-2025 Ziqi Fan +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import torch +from typing import TYPE_CHECKING + +import isaaclab.utils.math as math_utils +from isaaclab.assets import Articulation, RigidObject +from isaaclab.managers import ManagerTermBase +from isaaclab.managers import RewardTermCfg as RewTerm +from isaaclab.managers import SceneEntityCfg +from isaaclab.sensors import ContactSensor, RayCaster +from isaaclab.utils.math import quat_apply_inverse, yaw_quat + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedRLEnv + + +def track_lin_vel_xy_exp( + env: ManagerBasedRLEnv, std: float, command_name: str, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot") +) -> torch.Tensor: + """Reward tracking of linear velocity commands (xy axes) using exponential kernel.""" + # extract the used quantities (to enable type-hinting) + asset: RigidObject = env.scene[asset_cfg.name] + # compute the error + lin_vel_error = torch.sum( + torch.square(env.command_manager.get_command(command_name)[:, :2] - asset.data.root_lin_vel_b[:, :2]), + dim=1, + ) + reward = torch.exp(-lin_vel_error / std**2) + # reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7 + return reward + + +def track_ang_vel_z_exp( + env: ManagerBasedRLEnv, std: float, command_name: str, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot") +) -> torch.Tensor: + """Reward tracking of angular velocity commands (yaw) using exponential kernel.""" + # extract the used quantities (to enable type-hinting) + asset: RigidObject = env.scene[asset_cfg.name] + # compute the error + ang_vel_error = torch.square(env.command_manager.get_command(command_name)[:, 2] - asset.data.root_ang_vel_b[:, 2]) + reward = torch.exp(-ang_vel_error / std**2) + # reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7 + return reward + + +def track_diff_wheel_yaw_exp( + env: ManagerBasedRLEnv, + std: float, + command_name: str, + wheel_radius: float, + track_width: float, + left_wheel_indices: tuple[int, ...], + right_wheel_indices: tuple[int, ...], + command_threshold: float, + asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), +) -> torch.Tensor: + """Reward differential wheel speed matching the commanded yaw rate.""" + asset: Articulation = env.scene[asset_cfg.name] + joint_vel = asset.data.joint_vel[:, asset_cfg.joint_ids] + + left_wheel_vel = torch.mean(joint_vel[:, left_wheel_indices], dim=1) + right_wheel_vel = torch.mean(joint_vel[:, right_wheel_indices], dim=1) + estimated_yaw_rate = wheel_radius * (right_wheel_vel - left_wheel_vel) / track_width + yaw_command = env.command_manager.get_command(command_name)[:, 2] + + yaw_error = torch.square(yaw_command - estimated_yaw_rate) + reward = torch.exp(-yaw_error / std**2) + return torch.where(torch.abs(yaw_command) > command_threshold, reward, torch.zeros_like(reward)) + + +def lin_vel_xy_yaw_command_l2( + env: ManagerBasedRLEnv, + command_name: str, + yaw_command_threshold: float, + lin_command_threshold: float, + asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), +) -> torch.Tensor: + """Penalize linear drift when the command asks mostly for in-place yaw.""" + asset: Articulation = env.scene[asset_cfg.name] + command = env.command_manager.get_command(command_name) + is_turn_in_place = ( + torch.abs(command[:, 2]) > yaw_command_threshold + ) & (torch.linalg.norm(command[:, :2], dim=1) < lin_command_threshold) + lin_vel_error = torch.sum(torch.square(asset.data.root_lin_vel_b[:, :2]), dim=1) + return torch.where(is_turn_in_place, lin_vel_error, torch.zeros_like(lin_vel_error)) + + +def command_progress( + env: ManagerBasedRLEnv, + command_name: str, + min_command: float, + max_reward_speed: float, + asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), +) -> torch.Tensor: + """Reward positive progress along the commanded planar direction.""" + asset: Articulation = env.scene[asset_cfg.name] + command_xy = env.command_manager.get_command(command_name)[:, :2] + command_norm = torch.linalg.norm(command_xy, dim=1) + command_dir = command_xy / torch.clamp(command_norm.unsqueeze(1), min=1e-6) + progress_speed = torch.sum(asset.data.root_lin_vel_b[:, :2] * command_dir, dim=1) + reward = torch.clamp(progress_speed, min=0.0, max=max_reward_speed) / max_reward_speed + return torch.where(command_norm > min_command, reward, torch.zeros_like(reward)) + + +def support_contact_excess( + env: ManagerBasedRLEnv, + sensor_cfg: SceneEntityCfg, + force_threshold: float = 120.0, + max_contact_time: float = 0.6, +) -> torch.Tensor: + """Penalize dragging or high-force support contacts.""" + contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name] + force_norm = torch.linalg.norm(contact_sensor.data.net_forces_w[:, sensor_cfg.body_ids, :], dim=-1) + contact_time = contact_sensor.data.current_contact_time[:, sensor_cfg.body_ids] + + force_excess = torch.clamp(force_norm - force_threshold, min=0.0) / force_threshold + time_excess = torch.clamp(contact_time - max_contact_time, min=0.0) + return torch.sum(force_excess + time_excess, dim=1) + + +class ObstacleSupportEscapeReward(ManagerTermBase): + """Reward a blocked -> brief support contact -> recovered progress sequence. + + The state machine prevents ordinary forward motion with incidental support contact from receiving the full reward. + Obstacle evidence is inferred from low commanded progress together with either horizontal wheel force or wheel slip, + so the actor can remain blind while the training reward uses contact information already available on the robot. + """ + + def __init__(self, cfg: RewTerm, env: ManagerBasedRLEnv): + super().__init__(cfg, env) + self._blocked_memory = torch.zeros(self.num_envs, device=self.device) + self._support_window = torch.zeros(self.num_envs, device=self.device) + self._recovery_reward_window = torch.zeros(self.num_envs, device=self.device) + self._cooldown = torch.zeros(self.num_envs, device=self.device) + self._previous_support_active = torch.zeros(self.num_envs, dtype=torch.bool, device=self.device) + + def reset(self, env_ids=None) -> None: + if env_ids is None: + env_ids = slice(None) + self._blocked_memory[env_ids] = 0.0 + self._support_window[env_ids] = 0.0 + self._recovery_reward_window[env_ids] = 0.0 + self._cooldown[env_ids] = 0.0 + self._previous_support_active[env_ids] = False + + def __call__( + self, + env: ManagerBasedRLEnv, + command_name: str, + sensor_cfg: SceneEntityCfg, + wheel_sensor_cfg: SceneEntityCfg, + asset_cfg: SceneEntityCfg, + wheel_asset_cfg: SceneEntityCfg, + wheel_radius: float, + support_force_threshold: float, + wheel_obstacle_force_threshold: float, + max_contact_time: float, + min_command: float, + stall_speed_ratio: float, + min_stall_speed: float, + slip_threshold: float, + obstacle_memory_time: float, + support_window_time: float, + recovery_speed_ratio: float, + min_recovery_speed: float, + recovery_reward_time: float, + cooldown_time: float, + recovery_bonus: float, + max_reward_speed: float, + ) -> torch.Tensor: + asset: Articulation = env.scene[asset_cfg.name] + contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name] + dt = float(env.step_dt) + + command_x = env.command_manager.get_command(command_name)[:, 0] + command_speed = torch.abs(command_x) + direction = torch.where(command_x >= 0.0, torch.ones_like(command_x), -torch.ones_like(command_x)) + progress_speed = direction * asset.data.root_lin_vel_b[:, 0] + commanded_motion = command_speed > min_command + + wheel_forces = contact_sensor.data.net_forces_w[:, wheel_sensor_cfg.body_ids, :] + max_horizontal_wheel_force = torch.linalg.norm(wheel_forces[..., :2], dim=-1).amax(dim=1) + wheel_joint_vel = asset.data.joint_vel[:, wheel_asset_cfg.joint_ids] + mean_wheel_surface_speed = torch.mean(torch.abs(wheel_joint_vel), dim=1) * wheel_radius + wheel_slip = torch.clamp(mean_wheel_surface_speed - torch.clamp(progress_speed, min=0.0), min=0.0) + + stall_limit = torch.maximum(stall_speed_ratio * command_speed, torch.full_like(command_speed, min_stall_speed)) + obstacle_evidence = commanded_motion & (progress_speed < stall_limit) & ( + (max_horizontal_wheel_force > wheel_obstacle_force_threshold) | (wheel_slip > slip_threshold) + ) + + self._blocked_memory = torch.clamp(self._blocked_memory - dt, min=0.0) + self._support_window = torch.clamp(self._support_window - dt, min=0.0) + self._recovery_reward_window = torch.clamp(self._recovery_reward_window - dt, min=0.0) + self._cooldown = torch.clamp(self._cooldown - dt, min=0.0) + self._blocked_memory = torch.where( + obstacle_evidence, torch.full_like(self._blocked_memory, obstacle_memory_time), self._blocked_memory + ) + + support_force = torch.linalg.norm(contact_sensor.data.net_forces_w[:, sensor_cfg.body_ids, :], dim=-1) + support_time = contact_sensor.data.current_contact_time[:, sensor_cfg.body_ids] + support_active = torch.any(support_force > support_force_threshold, dim=1) + brief_support = torch.any( + (support_force > support_force_threshold) & (support_time <= max_contact_time), dim=1 + ) + support_onset = brief_support & ~self._previous_support_active + valid_support_onset = support_onset & (self._blocked_memory > 0.0) & (self._cooldown <= 0.0) + + was_waiting_for_recovery = self._support_window > 0.0 + self._support_window = torch.where( + valid_support_onset, torch.full_like(self._support_window, support_window_time), self._support_window + ) + + recovery_threshold = torch.maximum( + recovery_speed_ratio * command_speed, torch.full_like(command_speed, min_recovery_speed) + ) + excessive_support = torch.any( + (support_force > support_force_threshold) & (support_time > max_contact_time), dim=1 + ) + recovered = ( + was_waiting_for_recovery + & commanded_motion + & (progress_speed >= recovery_threshold) + & ~excessive_support + ) + self._recovery_reward_window = torch.where( + recovered, + torch.full_like(self._recovery_reward_window, recovery_reward_time), + self._recovery_reward_window, + ) + self._support_window = torch.where(recovered, torch.zeros_like(self._support_window), self._support_window) + self._cooldown = torch.where(recovered, torch.full_like(self._cooldown, cooldown_time), self._cooldown) + + recovery_progress = torch.clamp(progress_speed, min=0.0, max=max_reward_speed) / max_reward_speed + # Delay the discrete bonus until recovery is observed. This makes support contact alone worthless and + # prevents the policy from farming reward by repeatedly tapping the body against an obstacle. + reward = recovery_bonus * recovered.float() / dt + reward += torch.where( + self._recovery_reward_window > 0.0, recovery_progress, torch.zeros_like(recovery_progress) + ) + + self._previous_support_active = support_active + return reward + + +def track_lin_vel_xy_yaw_frame_exp( + env, std: float, command_name: str, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot") +) -> torch.Tensor: + """Reward tracking of linear velocity commands (xy axes) in the gravity aligned robot frame using exponential kernel.""" + # extract the used quantities (to enable type-hinting) + asset = env.scene[asset_cfg.name] + vel_yaw = quat_apply_inverse(yaw_quat(asset.data.root_quat_w), asset.data.root_lin_vel_w[:, :3]) + lin_vel_error = torch.sum( + torch.square(env.command_manager.get_command(command_name)[:, :2] - vel_yaw[:, :2]), dim=1 + ) + reward = torch.exp(-lin_vel_error / std**2) + # reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7 + return reward + + +def track_ang_vel_z_world_exp( + env, command_name: str, std: float, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot") +) -> torch.Tensor: + """Reward tracking of angular velocity commands (yaw) in world frame using exponential kernel.""" + # extract the used quantities (to enable type-hinting) + asset = env.scene[asset_cfg.name] + ang_vel_error = torch.square(env.command_manager.get_command(command_name)[:, 2] - asset.data.root_ang_vel_w[:, 2]) + reward = torch.exp(-ang_vel_error / std**2) + reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7 + return reward + + +def joint_power(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: + """Reward joint_power""" + # extract the used quantities (to enable type-hinting) + asset: Articulation = env.scene[asset_cfg.name] + # compute the reward + reward = torch.sum( + torch.abs(asset.data.joint_vel[:, asset_cfg.joint_ids] * asset.data.applied_torque[:, asset_cfg.joint_ids]), + dim=1, + ) + return reward + + +def stand_still_without_cmd( + env: ManagerBasedRLEnv, + command_name: str, + command_threshold: float, + asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), +) -> torch.Tensor: + """Penalize joint positions that deviate from the default one when no command.""" + # extract the used quantities (to enable type-hinting) + asset: Articulation = env.scene[asset_cfg.name] + # compute out of limits constraints + diff_angle = asset.data.joint_pos[:, asset_cfg.joint_ids] - asset.data.default_joint_pos[:, asset_cfg.joint_ids] + reward = torch.sum(torch.abs(diff_angle), dim=1) + reward *= torch.linalg.norm(env.command_manager.get_command(command_name), dim=1) < command_threshold + # reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7 + return reward + +def joint_pos_penalty( + env: ManagerBasedRLEnv, + command_name: str, + asset_cfg: SceneEntityCfg, + stand_still_scale: float, + velocity_threshold: float, + command_threshold: float, +) -> torch.Tensor: + """Penalize joint position error from default on the articulation.""" + # extract the used quantities (to enable type-hinting) + asset: Articulation = env.scene[asset_cfg.name] + cmd = torch.linalg.norm(env.command_manager.get_command(command_name), dim=1) + body_vel = torch.linalg.norm(asset.data.root_lin_vel_b[:, :2], dim=1) + running_reward = torch.linalg.norm( + (asset.data.joint_pos[:, asset_cfg.joint_ids] - asset.data.default_joint_pos[:, asset_cfg.joint_ids]), dim=1 + ) + reward = torch.where( + torch.logical_or(cmd > command_threshold, body_vel > velocity_threshold), + running_reward, + stand_still_scale * running_reward, + ) + # reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7 + return reward + + +def wheel_vel_penalty( + env: ManagerBasedRLEnv, + sensor_cfg: SceneEntityCfg, + command_name: str, + velocity_threshold: float, + command_threshold: float, + asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), +) -> torch.Tensor: + asset: Articulation = env.scene[asset_cfg.name] + cmd = torch.linalg.norm(env.command_manager.get_command(command_name), dim=1) + body_vel = torch.linalg.norm(asset.data.root_lin_vel_b[:, :2], dim=1) + joint_vel = torch.abs(asset.data.joint_vel[:, asset_cfg.joint_ids]) + contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name] + in_air = contact_sensor.compute_first_air(env.step_dt)[:, sensor_cfg.body_ids] + running_reward = torch.sum(in_air * joint_vel, dim=1) + standing_reward = torch.sum(joint_vel, dim=1) + reward = torch.where( + torch.logical_or(cmd > command_threshold, body_vel > velocity_threshold), + running_reward, + standing_reward, + ) + return reward + + +class GaitReward(ManagerTermBase): + """Gait enforcing reward term for quadrupeds. + + This reward penalizes contact timing differences between selected foot pairs defined in :attr:`synced_feet_pair_names` + to bias the policy towards a desired gait, i.e trotting, bounding, or pacing. Note that this reward is only for + quadrupedal gaits with two pairs of synchronized feet. + """ + + def __init__(self, cfg: RewTerm, env: ManagerBasedRLEnv): + """Initialize the term. + + Args: + cfg: The configuration of the reward. + env: The RL environment instance. + """ + super().__init__(cfg, env) + self.std: float = cfg.params["std"] + self.command_name: str = cfg.params["command_name"] + self.max_err: float = cfg.params["max_err"] + self.velocity_threshold: float = cfg.params["velocity_threshold"] + self.command_threshold: float = cfg.params["command_threshold"] + self.contact_sensor: ContactSensor = env.scene.sensors[cfg.params["sensor_cfg"].name] + self.asset: Articulation = env.scene[cfg.params["asset_cfg"].name] + # match foot body names with corresponding foot body ids + synced_feet_pair_names = cfg.params["synced_feet_pair_names"] + if ( + len(synced_feet_pair_names) != 2 + or len(synced_feet_pair_names[0]) != 2 + or len(synced_feet_pair_names[1]) != 2 + ): + raise ValueError("This reward only supports gaits with two pairs of synchronized feet, like trotting.") + synced_feet_pair_0 = self.contact_sensor.find_bodies(synced_feet_pair_names[0])[0] + synced_feet_pair_1 = self.contact_sensor.find_bodies(synced_feet_pair_names[1])[0] + self.synced_feet_pairs = [synced_feet_pair_0, synced_feet_pair_1] + + def __call__( + self, + env: ManagerBasedRLEnv, + std: float, + command_name: str, + max_err: float, + velocity_threshold: float, + command_threshold: float, + synced_feet_pair_names, + asset_cfg: SceneEntityCfg, + sensor_cfg: SceneEntityCfg, + ) -> torch.Tensor: + """Compute the reward. + + This reward is defined as a multiplication between six terms where two of them enforce pair feet + being in sync and the other four rewards if all the other remaining pairs are out of sync + + Args: + env: The RL environment instance. + Returns: + The reward value. + """ + # for synchronous feet, the contact (air) times of two feet should match + sync_reward_0 = self._sync_reward_func(self.synced_feet_pairs[0][0], self.synced_feet_pairs[0][1]) + sync_reward_1 = self._sync_reward_func(self.synced_feet_pairs[1][0], self.synced_feet_pairs[1][1]) + sync_reward = sync_reward_0 * sync_reward_1 + # for asynchronous feet, the contact time of one foot should match the air time of the other one + async_reward_0 = self._async_reward_func(self.synced_feet_pairs[0][0], self.synced_feet_pairs[1][0]) + async_reward_1 = self._async_reward_func(self.synced_feet_pairs[0][1], self.synced_feet_pairs[1][1]) + async_reward_2 = self._async_reward_func(self.synced_feet_pairs[0][0], self.synced_feet_pairs[1][1]) + async_reward_3 = self._async_reward_func(self.synced_feet_pairs[1][0], self.synced_feet_pairs[0][1]) + async_reward = async_reward_0 * async_reward_1 * async_reward_2 * async_reward_3 + # only enforce gait if cmd > 0 + cmd = torch.linalg.norm(env.command_manager.get_command(self.command_name), dim=1) + body_vel = torch.linalg.norm(self.asset.data.root_com_lin_vel_b[:, :2], dim=1) + reward = torch.where( + torch.logical_or(cmd > self.command_threshold, body_vel > self.velocity_threshold), + sync_reward * async_reward, + 0.0, + ) + # reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7 + return reward + + """ + Helper functions. + """ + + def _sync_reward_func(self, foot_0: int, foot_1: int) -> torch.Tensor: + """Reward synchronization of two feet.""" + air_time = self.contact_sensor.data.current_air_time + contact_time = self.contact_sensor.data.current_contact_time + # penalize the difference between the most recent air time and contact time of synced feet pairs. + se_air = torch.clip(torch.square(air_time[:, foot_0] - air_time[:, foot_1]), max=self.max_err**2) + se_contact = torch.clip(torch.square(contact_time[:, foot_0] - contact_time[:, foot_1]), max=self.max_err**2) + return torch.exp(-(se_air + se_contact) / self.std) + + def _async_reward_func(self, foot_0: int, foot_1: int) -> torch.Tensor: + """Reward anti-synchronization of two feet.""" + air_time = self.contact_sensor.data.current_air_time + contact_time = self.contact_sensor.data.current_contact_time + # penalize the difference between opposing contact modes air time of feet 1 to contact time of feet 2 + # and contact time of feet 1 to air time of feet 2) of feet pairs that are not in sync with each other. + se_act_0 = torch.clip(torch.square(air_time[:, foot_0] - contact_time[:, foot_1]), max=self.max_err**2) + se_act_1 = torch.clip(torch.square(contact_time[:, foot_0] - air_time[:, foot_1]), max=self.max_err**2) + return torch.exp(-(se_act_0 + se_act_1) / self.std) + + +def joint_mirror(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg, mirror_joints: list[list[str]]) -> torch.Tensor: + # extract the used quantities (to enable type-hinting) + asset: Articulation = env.scene[asset_cfg.name] + if not hasattr(env, "joint_mirror_joints_cache") or env.joint_mirror_joints_cache is None: + # Cache joint positions for all pairs + env.joint_mirror_joints_cache = [ + [asset.find_joints(joint_name) for joint_name in joint_pair] for joint_pair in mirror_joints + ] + reward = torch.zeros(env.num_envs, device=env.device) + # Iterate over all joint pairs + for joint_pair in env.joint_mirror_joints_cache: + # Calculate the difference for each pair and add to the total reward + diff = torch.sum( + torch.square(asset.data.joint_pos[:, joint_pair[0][0]] - asset.data.joint_pos[:, joint_pair[1][0]]), + dim=-1, + ) + reward += diff + reward *= 1 / len(mirror_joints) if len(mirror_joints) > 0 else 0 + reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7 + return reward + + +def action_mirror(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg, mirror_joints: list[list[str]]) -> torch.Tensor: + # extract the used quantities (to enable type-hinting) + asset: Articulation = env.scene[asset_cfg.name] + if not hasattr(env, "action_mirror_joints_cache") or env.action_mirror_joints_cache is None: + # Cache joint positions for all pairs + env.action_mirror_joints_cache = [ + [asset.find_joints(joint_name) for joint_name in joint_pair] for joint_pair in mirror_joints + ] + reward = torch.zeros(env.num_envs, device=env.device) + # Iterate over all joint pairs + for joint_pair in env.action_mirror_joints_cache: + # Calculate the difference for each pair and add to the total reward + diff = torch.sum( + torch.square( + torch.abs(env.action_manager.action[:, joint_pair[0][0]]) + - torch.abs(env.action_manager.action[:, joint_pair[1][0]]) + ), + dim=-1, + ) + reward += diff + reward *= 1 / len(mirror_joints) if len(mirror_joints) > 0 else 0 + # reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7 + return reward + + +def action_sync(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg, joint_groups: list[list[str]]) -> torch.Tensor: + # extract the used quantities (to enable type-hinting) + asset: Articulation = env.scene[asset_cfg.name] + + # Cache joint indices if not already done + if not hasattr(env, "action_sync_joint_cache") or env.action_sync_joint_cache is None: + env.action_sync_joint_cache = [ + [asset.find_joints(joint_name) for joint_name in joint_group] for joint_group in joint_groups + ] + + reward = torch.zeros(env.num_envs, device=env.device) + # Iterate over each joint group + for joint_group in env.action_sync_joint_cache: + if len(joint_group) < 2: + continue # need at least 2 joints to compare + + # Get absolute actions for all joints in this group + actions = torch.stack( + [torch.abs(env.action_manager.action[:, joint[0]]) for joint in joint_group], dim=1 + ) # shape: (num_envs, num_joints_in_group) + + # Calculate mean action for each environment + mean_actions = torch.mean(actions, dim=1, keepdim=True) + + # Calculate variance from mean for each joint + variance = torch.mean(torch.square(actions - mean_actions), dim=1) + + # Add to reward (we want to minimize this variance) + reward += variance.squeeze() + reward *= 1 / len(joint_groups) if len(joint_groups) > 0 else 0 + # reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7 + return reward + + +# def feet_air_time( +# env: ManagerBasedRLEnv, command_name: str, sensor_cfg: SceneEntityCfg, threshold: float +# ) -> torch.Tensor: +# """Reward long steps taken by the feet using L2-kernel. + +# This function rewards the agent for taking steps that are longer than a threshold. This helps ensure +# that the robot lifts its feet off the ground and takes steps. The reward is computed as the sum of +# the time for which the feet are in the air. + +# If the commands are small (i.e. the agent is not supposed to take a step), then the reward is zero. +# """ +# # extract the used quantities (to enable type-hinting) +# contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name] +# # compute the reward +# first_contact = contact_sensor.compute_first_contact(env.step_dt)[:, sensor_cfg.body_ids] +# last_air_time = contact_sensor.data.last_air_time[:, sensor_cfg.body_ids] +# reward = torch.sum((last_air_time - threshold) * first_contact, dim=1) +# # no reward for zero command +# reward *= torch.norm(env.command_manager.get_command(command_name)[:, :2], dim=1) > 0.1 +# # print(torch.norm(env.command_manager.get_command(command_name)[:, :2], dim=1), "command norm") +# reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7 + # return reward + +# def feet_air_time( +# env: ManagerBasedRLEnv, +# asset_cfg: SceneEntityCfg, +# sensor_cfg: SceneEntityCfg, +# mode_time: float, +# velocity_threshold: float, +# ) -> torch.Tensor: +# """Reward longer feet air and contact time.""" +# # extract the used quantities (to enable type-hinting) +# contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name] +# asset: Articulation = env.scene[asset_cfg.name] +# if contact_sensor.cfg.track_air_time is False: +# raise RuntimeError("Activate ContactSensor's track_air_time!") +# # compute the reward +# current_air_time = contact_sensor.data.current_air_time[:, sensor_cfg.body_ids] +# current_contact_time = contact_sensor.data.current_contact_time[:, sensor_cfg.body_ids] + +# t_max = torch.max(current_air_time, current_contact_time) +# t_min = torch.clip(t_max, max=mode_time) +# stance_cmd_reward = torch.clip(current_contact_time - current_air_time, -mode_time, mode_time) +# cmd = torch.norm(env.command_manager.get_command("base_velocity"), dim=1).unsqueeze(dim=1).expand(-1, 4) +# body_vel = torch.linalg.norm(asset.data.root_lin_vel_b[:, :2], dim=1).unsqueeze(dim=1).expand(-1, 4) +# reward = torch.where( +# torch.logical_or(cmd > 0.0, body_vel > velocity_threshold), +# torch.where(t_max < mode_time, t_min, 0), +# stance_cmd_reward, +# ) +# return torch.sum(reward, dim=1) + + +def feet_air_time_positive_biped(env, command_name: str, threshold: float, sensor_cfg: SceneEntityCfg) -> torch.Tensor: + """Reward long steps taken by the feet for bipeds. + + This function rewards the agent for taking steps up to a specified threshold and also keep one foot at + a time in the air. + + If the commands are small (i.e. the agent is not supposed to take a step), then the reward is zero. + """ + contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name] + # compute the reward + air_time = contact_sensor.data.current_air_time[:, sensor_cfg.body_ids] + contact_time = contact_sensor.data.current_contact_time[:, sensor_cfg.body_ids] + in_contact = contact_time > 0.0 + in_mode_time = torch.where(in_contact, contact_time, air_time) + single_stance = torch.sum(in_contact.int(), dim=1) == 1 + reward = torch.min(torch.where(single_stance.unsqueeze(-1), in_mode_time, 0.0), dim=1)[0] + reward = torch.clamp(reward, max=threshold) + # no reward for zero command + reward *= torch.norm(env.command_manager.get_command(command_name)[:, :2], dim=1) > 0.1 + # reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7 + return reward + + +def feet_air_time_variance_penalty(env: ManagerBasedRLEnv, sensor_cfg: SceneEntityCfg) -> torch.Tensor: + """Penalize variance in the amount of time each foot spends in the air/on the ground relative to each other""" + # extract the used quantities (to enable type-hinting) + contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name] + # compute the reward + last_air_time = contact_sensor.data.last_air_time[:, sensor_cfg.body_ids] + last_contact_time = contact_sensor.data.last_contact_time[:, sensor_cfg.body_ids] + reward = torch.var(torch.clip(last_air_time, max=0.5), dim=1) + torch.var( + torch.clip(last_contact_time, max=0.5), dim=1) + # print(last_air_time, "last air time") + # print(last_contact_time, "last contact time") + # reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7 + return reward + + + + +def feet_contact( + env: ManagerBasedRLEnv, command_name: str, expect_contact_num: int, sensor_cfg: SceneEntityCfg +) -> torch.Tensor: + """Reward feet contact""" + # extract the used quantities (to enable type-hinting) + contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name] + # compute the reward + contact = contact_sensor.compute_first_contact(env.step_dt)[:, sensor_cfg.body_ids] + contact_num = torch.sum(contact, dim=1) + reward = (contact_num != expect_contact_num).float() + # no reward for zero command + reward *= torch.linalg.norm(env.command_manager.get_command(command_name), dim=1) > 0.5 + reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7 + return reward + + +def feet_contact_without_cmd(env: ManagerBasedRLEnv, command_name: str, sensor_cfg: SceneEntityCfg) -> torch.Tensor: + """Reward feet contact""" + # extract the used quantities (to enable type-hinting) + contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name] + # compute the reward + contact = contact_sensor.compute_first_contact(env.step_dt)[:, sensor_cfg.body_ids] + # print(contact, "contact") + reward = torch.sum(contact, dim=-1).float() + # print(reward, "reward after sum") + reward *= torch.linalg.norm(env.command_manager.get_command(command_name), dim=1) < 0.5 + # print(env.command_manager.get_command(command_name), "env.command_manager.get_command(command_name)") + # print(reward, "reward after multiply") + # reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7 + return reward + + +def feet_stumble(env: ManagerBasedRLEnv, sensor_cfg: SceneEntityCfg) -> torch.Tensor: + # extract the used quantities (to enable type-hinting) + contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name] + forces_z = torch.abs(contact_sensor.data.net_forces_w[:, sensor_cfg.body_ids, 2]) + forces_xy = torch.linalg.norm(contact_sensor.data.net_forces_w[:, sensor_cfg.body_ids, :2], dim=2) + # Penalize feet hitting vertical surfaces + reward = torch.any(forces_xy > 4 * forces_z, dim=1).float() + # reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7 + return reward + + +def feet_distance_y_exp( + env: ManagerBasedRLEnv, stance_width: float, std: float, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot") +) -> torch.Tensor: + asset: RigidObject = env.scene[asset_cfg.name] + cur_footsteps_translated = asset.data.body_link_pos_w[:, asset_cfg.body_ids, :] - asset.data.root_link_pos_w[ + :, : + ].unsqueeze(1) + n_feet = len(asset_cfg.body_ids) + footsteps_in_body_frame = torch.zeros(env.num_envs, n_feet, 3, device=env.device) + for i in range(n_feet): + footsteps_in_body_frame[:, i, :] = math_utils.quat_apply( + math_utils.quat_conjugate(asset.data.root_link_quat_w), cur_footsteps_translated[:, i, :] + ) + side_sign = torch.tensor( + [1.0 if i % 2 == 0 else -1.0 for i in range(n_feet)], + device=env.device, + ) + stance_width_tensor = stance_width * torch.ones([env.num_envs, 1], device=env.device) + desired_ys = stance_width_tensor / 2 * side_sign.unsqueeze(0) + stance_diff = torch.square(desired_ys - footsteps_in_body_frame[:, :, 1]) + reward = torch.exp(-torch.sum(stance_diff, dim=1) / (std**2)) + # reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7 + return reward + + +def feet_distance_xy_exp( + env: ManagerBasedRLEnv, + stance_width: float, + stance_length: float, + std: float, + asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), +) -> torch.Tensor: + asset: RigidObject = env.scene[asset_cfg.name] + + # Compute the current footstep positions relative to the root + cur_footsteps_translated = asset.data.body_link_pos_w[:, asset_cfg.body_ids, :] - asset.data.root_link_pos_w[ + :, : + ].unsqueeze(1) + + footsteps_in_body_frame = torch.zeros(env.num_envs, 4, 3, device=env.device) + for i in range(4): + footsteps_in_body_frame[:, i, :] = math_utils.quat_apply( + math_utils.quat_conjugate(asset.data.root_link_quat_w), cur_footsteps_translated[:, i, :] + ) + + # Desired x and y positions for each foot + stance_width_tensor = stance_width * torch.ones([env.num_envs, 1], device=env.device) + stance_length_tensor = stance_length * torch.ones([env.num_envs, 1], device=env.device) + + desired_xs = torch.cat( + [stance_length_tensor / 2, stance_length_tensor / 2, -stance_length_tensor / 2, -stance_length_tensor / 2], + dim=1, + ) + desired_ys = torch.cat( + [stance_width_tensor / 2, -stance_width_tensor / 2, stance_width_tensor / 2, -stance_width_tensor / 2], dim=1 + ) + + # Compute differences in x and y + stance_diff_x = torch.square(desired_xs - footsteps_in_body_frame[:, :, 0]) + stance_diff_y = torch.square(desired_ys - footsteps_in_body_frame[:, :, 1]) + + # Combine x and y differences and compute the exponential penalty + stance_diff = stance_diff_x + stance_diff_y + reward = torch.exp(-torch.sum(stance_diff, dim=1) / std**2) + # reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7 + return reward + + +def feet_height( + env: ManagerBasedRLEnv, + command_name: str, + asset_cfg: SceneEntityCfg, + target_height: float, + tanh_mult: float, +) -> torch.Tensor: + """Reward the swinging feet for clearing a specified height off the ground""" + asset: RigidObject = env.scene[asset_cfg.name] + foot_z_target_error = torch.square(asset.data.body_pos_w[:, asset_cfg.body_ids, 2] - target_height) + # foot_velocity_tanh = torch.tanh( + # tanh_mult * torch.linalg.norm(asset.data.body_lin_vel_w[:, asset_cfg.body_ids, :2], dim=2) + # ) + # reward = torch.sum(foot_z_target_error * foot_velocity_tanh, dim=1) + reward = torch.sum(foot_z_target_error, dim=1) + # print(foot_z_target_error, "foot_z_target_error") + # no reward for zero command + reward *= torch.linalg.norm(env.command_manager.get_command(command_name), dim=1) > 0.2 + # reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7 + return reward + + +def feet_height_body( + env: ManagerBasedRLEnv, + command_name: str, + asset_cfg: SceneEntityCfg, + target_height: float, + tanh_mult: float, +) -> torch.Tensor: + """Reward the swinging feet for clearing a specified height off the ground""" + asset: RigidObject = env.scene[asset_cfg.name] + cur_footpos_translated = asset.data.body_pos_w[:, asset_cfg.body_ids, :] - asset.data.root_pos_w[:, :].unsqueeze(1) + footpos_in_body_frame = torch.zeros(env.num_envs, len(asset_cfg.body_ids), 3, device=env.device) + cur_footvel_translated = asset.data.body_lin_vel_w[:, asset_cfg.body_ids, :] - asset.data.root_lin_vel_w[ + :, : + ].unsqueeze(1) + footvel_in_body_frame = torch.zeros(env.num_envs, len(asset_cfg.body_ids), 3, device=env.device) + for i in range(len(asset_cfg.body_ids)): + footpos_in_body_frame[:, i, :] = math_utils.quat_apply_inverse( + asset.data.root_quat_w, cur_footpos_translated[:, i, :] + ) + footvel_in_body_frame[:, i, :] = math_utils.quat_apply_inverse( + asset.data.root_quat_w, cur_footvel_translated[:, i, :] + ) + foot_z_target_error = torch.square(footpos_in_body_frame[:, :, 2] - target_height).view(env.num_envs, -1) + foot_velocity_tanh = torch.tanh(tanh_mult * torch.norm(footvel_in_body_frame[:, :, :2], dim=2)) + reward = torch.sum(foot_z_target_error * foot_velocity_tanh, dim=1) + reward *= torch.linalg.norm(env.command_manager.get_command(command_name), dim=1) > 0.1 + # reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7 + return reward + + +def feet_slide( + env: ManagerBasedRLEnv, sensor_cfg: SceneEntityCfg, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot") +) -> torch.Tensor: + """Penalize feet sliding. + + This function penalizes the agent for sliding its feet on the ground. The reward is computed as the + norm of the linear velocity of the feet multiplied by a binary contact sensor. This ensures that the + agent is penalized only when the feet are in contact with the ground. + """ + # Penalize feet sliding + contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name] + contacts = contact_sensor.data.net_forces_w_history[:, :, sensor_cfg.body_ids, :].norm(dim=-1).max(dim=1)[0] > 1.0 + asset: RigidObject = env.scene[asset_cfg.name] + + # feet_vel = asset.data.body_lin_vel_w[:, asset_cfg.body_ids, :2] + # reward = torch.sum(feet_vel.norm(dim=-1) * contacts, dim=1) + + cur_footvel_translated = asset.data.body_lin_vel_w[:, asset_cfg.body_ids, :] - asset.data.root_lin_vel_w[ + :, : + ].unsqueeze(1) + footvel_in_body_frame = torch.zeros(env.num_envs, len(asset_cfg.body_ids), 3, device=env.device) + for i in range(len(asset_cfg.body_ids)): + footvel_in_body_frame[:, i, :] = math_utils.quat_apply_inverse( + asset.data.root_quat_w, cur_footvel_translated[:, i, :] + ) + foot_leteral_vel = torch.sqrt(torch.sum(torch.square(footvel_in_body_frame[:, :, :2]), dim=2)).view( + env.num_envs, -1 + ) + reward = torch.sum(foot_leteral_vel * contacts, dim=1) + # reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7 + return reward + +# def stand_still_joint_deviation_l1( +# env, command_name: str, command_threshold: float = 0.06, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot") +# ) -> torch.Tensor: +# """Penalize offsets from the default joint positions when the command is very small.""" + # command = env.command_manager.get_command(command_name) +# # Penalize motion when command is nearly zero. +# return joint_deviation_l1(env, asset_cfg) * (torch.norm(command[:, :], dim=1) < command_threshold) + +# def joint_deviation_l1(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: +# """Penalize joint positions that deviate from the default one.""" +# # extract the used quantities (to enable type-hinting) +# asset: Articulation = env.scene[asset_cfg.name] +# # compute out of limits constraints +# angle = asset.data.joint_pos[:, asset_cfg.joint_ids] - asset.data.default_joint_pos[:, asset_cfg.joint_ids] +# return torch.sum(torch.abs(angle), dim=1) + + +# def smoothness_1(env: ManagerBasedRLEnv) -> torch.Tensor: +# # Penalize changes in actions +# diff = torch.square(env.action_manager.action - env.action_manager.prev_action) +# diff = diff * (env.action_manager.prev_action[:, :] != 0) # ignore first step +# return torch.sum(diff, dim=1) + + +# def joint_acc_l2_new(env: ManagerBasedRLEnv) -> torch.Tensor: + +# def smoothness_2(env: ManagerBasedRLEnv) -> torch.Tensor: +# # Penalize changes in actions +# diff = torch.square(env.action_manager.action - 2 * env.action_manager.prev_action + env.action_manager.prev_prev_action) +# diff = diff * (env.action_manager.prev_action[:, :] != 0) # ignore first step +# diff = diff * (env.action_manager.prev_prev_action[:, :] != 0) # ignore second step +# # print(torch.sum(diff, dim=1), "smoothness l2") +# return torch.sum(diff, dim=1) + + +def upward(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: + """Penalize z-axis base linear velocity using L2 squared kernel.""" + # extract the used quantities (to enable type-hinting) + asset: RigidObject = env.scene[asset_cfg.name] + reward = torch.square(1 - asset.data.projected_gravity_b[:, 2]) + return reward + + + + + +def base_height_l2( + env: ManagerBasedRLEnv, + target_height: float, + asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), + sensor_cfg: SceneEntityCfg | None = None, +) -> torch.Tensor: + """Penalize asset height from its target using L2 squared kernel. + + Note: + For flat terrain, target height is in the world frame. For rough terrain, + sensor readings can adjust the target height to account for the terrain. + """ + # extract the used quantities (to enable type-hinting) + asset: RigidObject = env.scene[asset_cfg.name] + if sensor_cfg is not None: + sensor: RayCaster = env.scene[sensor_cfg.name] + # Adjust the target height using the sensor data + ray_hits = sensor.data.ray_hits_w[..., 2] + if torch.isnan(ray_hits).any() or torch.isinf(ray_hits).any() or torch.max(torch.abs(ray_hits)) > 1e6: + adjusted_target_height = asset.data.root_link_pos_w[:, 2] + else: + adjusted_target_height = target_height + torch.mean(ray_hits, dim=1) + else: + # Use the provided target height directly for flat terrain + adjusted_target_height = target_height + # Compute the L2 squared penalty + reward = torch.square(asset.data.root_pos_w[:, 2] - adjusted_target_height) + # reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7 + return reward + + +def lin_vel_z_l2(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: + """Penalize z-axis base linear velocity using L2 squared kernel.""" + # extract the used quantities (to enable type-hinting) + asset: RigidObject = env.scene[asset_cfg.name] + reward = torch.square(asset.data.root_lin_vel_b[:, 2]) + # reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7 + return reward + + +def ang_vel_xy_l2(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: + """Penalize xy-axis base angular velocity using L2 squared kernel.""" + # extract the used quantities (to enable type-hinting) + asset: RigidObject = env.scene[asset_cfg.name] + reward = torch.sum(torch.square(asset.data.root_ang_vel_b[:, :2]), dim=1) + # reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7 + return reward + + +def undesired_contacts(env: ManagerBasedRLEnv, threshold: float, sensor_cfg: SceneEntityCfg) -> torch.Tensor: + """Penalize undesired contacts as the number of violations that are above a threshold.""" + # extract the used quantities (to enable type-hinting) + contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name] + # check if contact force is above threshold + net_contact_forces = contact_sensor.data.net_forces_w_history + is_contact = torch.max(torch.norm(net_contact_forces[:, :, sensor_cfg.body_ids], dim=-1), dim=1)[0] > threshold + # sum over contacts for each environment + reward = torch.sum(is_contact, dim=1).float() + # reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7 + return reward + + +def flat_orientation_l2(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: + """Penalize non-flat base orientation using L2 squared kernel. + + This is computed by penalizing the xy-components of the projected gravity vector. + """ + # extract the used quantities (to enable type-hinting) + asset: RigidObject = env.scene[asset_cfg.name] + reward = torch.sum(torch.square(asset.data.projected_gravity_b[:, :2]), dim=1) + # reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7 + return reward + +def feet_air_time_including_ang_z( + env: ManagerBasedRLEnv, command_name: str, sensor_cfg: SceneEntityCfg, threshold: float +) -> torch.Tensor: + """Reward long steps taken by the feet using L2-kernel. + + This function rewards the agent for taking steps that are longer than a threshold. This helps ensure + that the robot lifts its feet off the ground and takes steps. The reward is computed as the sum of + the time for which the feet are in the air. + + If the commands are small (i.e. the agent is not supposed to take a step), then the reward is zero. + """ + # extract the used quantities (to enable type-hinting) + contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name] + # compute the reward + first_contact = contact_sensor.compute_first_contact(env.step_dt)[:, sensor_cfg.body_ids] + last_air_time = contact_sensor.data.last_air_time[:, sensor_cfg.body_ids] + reward = torch.sum((last_air_time - threshold) * first_contact, dim=1) + # no reward for zero command + reward *= torch.norm(env.command_manager.get_command(command_name), dim=1) > 0.1 + # reward *= torch.norm(env.command_manager.get_command(command_name)[:, :3], dim=1) > 0.1 + return reward + +def lin_vel_xy_l2_with_ang_z_command( + env: ManagerBasedRLEnv, + command_name: str, + command_threshold: float, + asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), + ) -> torch.Tensor: + """Penalize xy-axis base linear velocity using L2 squared kernel if command is ang_vel_z.""" + # extract the used quantities (to enable type-hinting) + asset: RigidObject = env.scene[asset_cfg.name] + # reward = torch.square(asset.data.root_lin_vel_b[:, 2]) + reward = torch.sum(torch.square(asset.data.root_lin_vel_b[:, :2]), dim=1) + command = env.command_manager.get_command(command_name) + reward *= (torch.sum(torch.square(command[:, 2:]), dim=1) > command_threshold) & \ + (torch.sum(torch.square(command[:, :2]), dim=1) < command_threshold) + # reward *= torch.sum(torch.square(env.command_manager.get_command(command_name)[:, 2:]), dim=1) > command_threshold + # reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7 + return reward diff --git a/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/velocity_env_cfg.py b/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/velocity_env_cfg.py new file mode 100644 index 0000000..a8e0671 --- /dev/null +++ b/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/velocity_env_cfg.py @@ -0,0 +1,866 @@ +# Copyright (c) 2025 Deep Robotics +# SPDX-License-Identifier: BSD 3-Clause + +# Copyright (c) 2024-2025 Ziqi Fan +# SPDX-License-Identifier: Apache-2.0 + +# Copyright (c) 2022-2025, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import inspect +import math +import sys +from dataclasses import MISSING + +import isaaclab.sim as sim_utils +from isaaclab.assets import ArticulationCfg, AssetBaseCfg +from isaaclab.envs import ManagerBasedRLEnvCfg +from isaaclab.managers import CurriculumTermCfg as CurrTerm +from isaaclab.managers import EventTermCfg as EventTerm +from isaaclab.managers import ObservationGroupCfg as ObsGroup +from isaaclab.managers import ObservationTermCfg as ObsTerm +from isaaclab.managers import RewardTermCfg as RewTerm +from isaaclab.managers import SceneEntityCfg +from isaaclab.managers import TerminationTermCfg as DoneTerm +from isaaclab.scene import InteractiveSceneCfg +from isaaclab.sensors import ContactSensorCfg, RayCasterCfg, patterns +from isaaclab.terrains import TerrainImporterCfg +from isaaclab.utils import configclass +from isaaclab.utils.noise import AdditiveUniformNoiseCfg as Unoise + +import rl_training.tasks.manager_based.locomotion.velocity.mdp as mdp + +## +# Pre-defined configs +## +from isaaclab.terrains.config.rough import ROUGH_TERRAINS_CFG # isort: skip + + +## +# Scene definition +## + + +@configclass +class MySceneCfg(InteractiveSceneCfg): + """Configuration for the terrain scene with a legged robot.""" + + # ground terrain + terrain = TerrainImporterCfg( + prim_path="/World/ground", + terrain_type="generator", + terrain_generator=ROUGH_TERRAINS_CFG, + max_init_terrain_level=5, + collision_group=-1, + physics_material=sim_utils.RigidBodyMaterialCfg( + friction_combine_mode="multiply", + restitution_combine_mode="multiply", + static_friction=1.0, + dynamic_friction=1.0, + restitution=1.0, + ), + # Keep scene startup independent of NVIDIA's online asset server. + # The robot visual geometry is still loaded from the STL meshes in its URDF. + visual_material=sim_utils.PreviewSurfaceCfg( + diffuse_color=(0.18, 0.20, 0.22), + roughness=0.9, + ), + debug_vis=False, + ) + # robots + robot: ArticulationCfg = MISSING + # sensors + height_scanner = RayCasterCfg( + prim_path="{ENV_REGEX_NS}/Robot/base", + offset=RayCasterCfg.OffsetCfg(pos=(0.0, 0.0, 20.0)), + ray_alignment="yaw", + pattern_cfg=patterns.GridPatternCfg(resolution=0.1, size=[1.6, 1.0]), + debug_vis=False, + mesh_prim_paths=["/World/ground"], + ) + height_scanner_base = RayCasterCfg( + prim_path="{ENV_REGEX_NS}/Robot/base", + offset=RayCasterCfg.OffsetCfg(pos=(0.0, 0.0, 20.0)), + ray_alignment="yaw", + pattern_cfg=patterns.GridPatternCfg(resolution=0.05, size=(0.1, 0.1)), + debug_vis=False, + mesh_prim_paths=["/World/ground"], + ) + contact_forces = ContactSensorCfg(prim_path="{ENV_REGEX_NS}/Robot/.*", history_length=3, track_air_time=True) + # lights + sky_light = AssetBaseCfg( + prim_path="/World/skyLight", + spawn=sim_utils.DomeLightCfg( + intensity=1200.0, + color=(0.75, 0.75, 0.75), + ), + ) + + +## +# MDP settings +## + + +@configclass +class CommandsCfg: + """Command specifications for the MDP.""" + + base_velocity = mdp.UniformThresholdVelocityCommandCfg( + asset_name="robot", + resampling_time_range=(10.0, 10.0), + rel_standing_envs=0.02, + rel_heading_envs=1.0, + heading_command=True, + heading_control_stiffness=0.5, + debug_vis=True, + ranges=mdp.UniformThresholdVelocityCommandCfg.Ranges( + lin_vel_x=(-1.0, 1.0), lin_vel_y=(-1.0, 1.0), ang_vel_z=(-1.0, 1.0), heading=(-math.pi, math.pi) + ), + ) + + +@configclass +class ActionsCfg: + """Action specifications for the MDP.""" + + joint_pos = mdp.JointPositionActionCfg( + asset_name="robot", joint_names=[".*"], scale=0.5, use_default_offset=True, clip=None, preserve_order=True + ) + + +@configclass +class ObservationsCfg: + """Observation specifications for the MDP.""" + + @configclass + class PolicyCfg(ObsGroup): + """Observations for policy group.""" + + # observation terms (order preserved) + base_lin_vel = ObsTerm( + func=mdp.base_lin_vel, + noise=Unoise(n_min=-0.1, n_max=0.1), + clip=(-100.0, 100.0), + scale=1.0, + ) + base_ang_vel = ObsTerm( + func=mdp.base_ang_vel, + noise=Unoise(n_min=-0.2, n_max=0.2), + clip=(-100.0, 100.0), + scale=1.0, + ) + projected_gravity = ObsTerm( + func=mdp.projected_gravity, + noise=Unoise(n_min=-0.05, n_max=0.05), + clip=(-100.0, 100.0), + scale=1.0, + ) + velocity_commands = ObsTerm( + func=mdp.generated_commands, + params={"command_name": "base_velocity"}, + clip=(-100.0, 100.0), + scale=1.0, + ) + joint_pos = ObsTerm( + func=mdp.joint_pos_rel, + params={"asset_cfg": SceneEntityCfg("robot", joint_names=".*", preserve_order=True)}, + noise=Unoise(n_min=-0.01, n_max=0.01), + clip=(-100.0, 100.0), + scale=1.0, + ) + joint_vel = ObsTerm( + func=mdp.joint_vel_rel, + params={"asset_cfg": SceneEntityCfg("robot", joint_names=".*", preserve_order=True)}, + noise=Unoise(n_min=-1.5, n_max=1.5), + clip=(-100.0, 100.0), + scale=1.0, + ) + actions = ObsTerm( + func=mdp.last_action, + clip=(-100.0, 100.0), + scale=1.0, + ) + height_scan = ObsTerm( + func=mdp.height_scan, + params={"sensor_cfg": SceneEntityCfg("height_scanner")}, + noise=Unoise(n_min=-0.1, n_max=0.1), + clip=(-1.0, 1.0), + scale=1.0, + ) + + def __post_init__(self): + self.enable_corruption = True + self.concatenate_terms = True + + @configclass + class CriticCfg(ObsGroup): + """Observations for critic group.""" + + # observation terms (order preserved) + base_lin_vel = ObsTerm( + func=mdp.base_lin_vel, + clip=(-100.0, 100.0), + scale=1.0, + ) + base_ang_vel = ObsTerm( + func=mdp.base_ang_vel, + clip=(-100.0, 100.0), + scale=1.0, + ) + projected_gravity = ObsTerm( + func=mdp.projected_gravity, + clip=(-100.0, 100.0), + scale=1.0, + ) + velocity_commands = ObsTerm( + func=mdp.generated_commands, + params={"command_name": "base_velocity"}, + clip=(-100.0, 100.0), + scale=1.0, + ) + joint_pos = ObsTerm( + func=mdp.joint_pos_rel, + params={"asset_cfg": SceneEntityCfg("robot", joint_names=".*", preserve_order=True)}, + clip=(-100.0, 100.0), + scale=1.0, + ) + joint_vel = ObsTerm( + func=mdp.joint_vel_rel, + params={"asset_cfg": SceneEntityCfg("robot", joint_names=".*", preserve_order=True)}, + clip=(-100.0, 100.0), + scale=1.0, + ) + actions = ObsTerm( + func=mdp.last_action, + clip=(-100.0, 100.0), + scale=1.0, + ) + height_scan = ObsTerm( + func=mdp.height_scan, + params={"sensor_cfg": SceneEntityCfg("height_scanner")}, + clip=(-1.0, 1.0), + scale=1.0, + ) + # joint_effort = ObsTerm( + # func=mdp.joint_effort, + # clip=(-100, 100), + # scale=0.01, + # ) + + def __post_init__(self): + self.enable_corruption = False + self.concatenate_terms = True + + # observation groups + policy: PolicyCfg = PolicyCfg() + critic: CriticCfg = CriticCfg() + + +@configclass +class EventCfg: + """Configuration for events.""" + + # startup + randomize_rigid_body_material = EventTerm( + func=mdp.randomize_rigid_body_material, + mode="startup", + params={ + "asset_cfg": SceneEntityCfg("robot", body_names=".*"), + "static_friction_range": (0.3, 1.0), + "dynamic_friction_range": (0.3, 0.8), + "restitution_range": (0.0, 0.4), + "num_buckets": 1024, + }, + ) + + randomize_rigid_body_mass = EventTerm( + func=mdp.randomize_rigid_body_mass, + mode="startup", + params={ + "asset_cfg": SceneEntityCfg("robot", body_names=""), + "mass_distribution_params": (0.85, 1.15), + "operation": "scale", + "recompute_inertia": True, + }, + ) + + randomize_rigid_body_mass_base = EventTerm( + func=mdp.randomize_rigid_body_mass, + mode="startup", + params={ + "asset_cfg": SceneEntityCfg("robot", body_names=""), + "mass_distribution_params": (-1.0, 3.0), + "operation": "add", + "recompute_inertia": True, + }, + ) + + randomize_rigid_body_inertia = EventTerm( + func=mdp.randomize_rigid_body_inertia, + mode="startup", + params={ + "asset_cfg": SceneEntityCfg("robot", body_names=".*"), + "inertia_distribution_params": (0.85, 1.15), + "operation": "scale", + }, + ) + + randomize_com_positions = EventTerm( + func=mdp.randomize_rigid_body_com, + mode="startup", + params={ + "asset_cfg": SceneEntityCfg("robot", body_names=".*"), + "com_range": {"x": (-0.03, 0.03), "y": (-0.03, 0.03), "z": (-0.02, 0.02)}, + }, + ) + + # reset + randomize_apply_external_force_torque = EventTerm( + func=mdp.apply_external_force_torque, + mode="reset", + params={ + "asset_cfg": SceneEntityCfg("robot", body_names=""), + "force_range": (-10.0, 10.0), + "torque_range": (-10.0, 10.0), + }, + ) + + randomize_reset_joints = EventTerm( + func=mdp.reset_joints_by_scale, + # func=mdp.reset_joints_by_offset, + mode="reset", + params={ + "position_range": (1.0, 1.0), + "velocity_range": (0.0, 0.0), + }, + ) + + randomize_actuator_gains = EventTerm( + func=mdp.randomize_actuator_gains, + mode="reset", + params={ + "asset_cfg": SceneEntityCfg("robot", joint_names=".*"), + "stiffness_distribution_params": (0.85, 1.15), + "damping_distribution_params": (0.85, 1.15), + "operation": "scale", + "distribution": "uniform", + }, + ) + + randomize_reset_base = EventTerm( + func=mdp.reset_root_state_uniform, + mode="reset", + params={ + "pose_range": {"x": (-0.5, 0.5), "y": (-0.5, 0.5), "yaw": (-3.14, 3.14)}, + "velocity_range": { + "x": (-0.5, 0.5), + "y": (-0.5, 0.5), + "z": (-0.5, 0.5), + "roll": (-0.5, 0.5), + "pitch": (-0.5, 0.5), + "yaw": (-0.5, 0.5), + }, + }, + ) + + # interval + randomize_push_robot = EventTerm( + func=mdp.push_by_setting_velocity, + mode="interval", + interval_range_s=(10.0, 15.0), + params={"velocity_range": {"x": (-0.5, 0.5), "y": (-0.5, 0.5)}}, + ) + + +@configclass +class RewardsCfg: + """Reward terms for the MDP.""" + + # General + is_terminated = RewTerm(func=mdp.is_terminated, weight=0.0) + + # Root penalties + lin_vel_z_l2 = RewTerm(func=mdp.lin_vel_z_l2, weight=0.0) + ang_vel_xy_l2 = RewTerm(func=mdp.ang_vel_xy_l2, weight=0.0) + flat_orientation_l2 = RewTerm(func=mdp.flat_orientation_l2, weight=0.0) + base_height_l2 = RewTerm( + func=mdp.base_height_l2, + weight=0.0, + params={ + "asset_cfg": SceneEntityCfg("robot", body_names=""), + "sensor_cfg": SceneEntityCfg("height_scanner_base"), + "target_height": 0.0, + }, + ) + body_lin_acc_l2 = RewTerm( + func=mdp.body_lin_acc_l2, + weight=0.0, + params={"asset_cfg": SceneEntityCfg("robot", body_names="")}, + ) + + # Joint penalties + joint_torques_l2 = RewTerm( + func=mdp.joint_torques_l2, weight=0.0, params={"asset_cfg": SceneEntityCfg("robot", joint_names=".*")} + ) + joint_vel_l2 = RewTerm( + func=mdp.joint_vel_l2, weight=0.0, params={"asset_cfg": SceneEntityCfg("robot", joint_names=".*")} + ) + joint_acc_l2 = RewTerm( + func=mdp.joint_acc_l2, weight=0.0, params={"asset_cfg": SceneEntityCfg("robot", joint_names=".*")} + ) + + joint_deviation_l1 = RewTerm( + func=mdp.joint_deviation_l1, + weight=0.0, + params={"asset_cfg": SceneEntityCfg("robot", joint_names=".*")}, + ) + + def create_joint_deviation_l1_rewterm(self, attr_name, weight, joint_names_pattern): + rew_term = RewTerm( + func=mdp.joint_deviation_l1, + weight=weight, + params={"asset_cfg": SceneEntityCfg("robot", joint_names=joint_names_pattern)}, + ) + setattr(self, attr_name, rew_term) + + joint_pos_limits = RewTerm( + func=mdp.joint_pos_limits, weight=0.0, params={"asset_cfg": SceneEntityCfg("robot", joint_names=".*")} + ) + joint_vel_limits = RewTerm( + func=mdp.joint_vel_limits, + weight=0.0, + params={"asset_cfg": SceneEntityCfg("robot", joint_names=".*"), "soft_ratio": 1.0}, + ) + joint_power = RewTerm( + func=mdp.joint_power, + weight=0.0, + params={ + "asset_cfg": SceneEntityCfg("robot", joint_names=".*"), + }, + ) + + stand_still_without_cmd = RewTerm( + func=mdp.stand_still_without_cmd, + weight=0.0, + params={ + "command_name": "base_velocity", + "command_threshold": 0.1, + "asset_cfg": SceneEntityCfg("robot", joint_names=".*"), + }, + ) + + joint_pos_penalty = RewTerm( + func=mdp.joint_pos_penalty, + weight=0.0, + params={ + "command_name": "base_velocity", + "asset_cfg": SceneEntityCfg("robot", joint_names=".*"), + "stand_still_scale": 5.0, + "velocity_threshold": 0.5, + "command_threshold": 0.1, + }, + ) + + hipx_joint_pos_penalty = RewTerm( + func=mdp.joint_pos_penalty, + weight=0.0, + params={ + "command_name": "base_velocity", + "asset_cfg": SceneEntityCfg("robot", joint_names=".*"), + "stand_still_scale": 5.0, + "velocity_threshold": 0.5, + "command_threshold": 0.1, + }, + ) + + hipy_joint_pos_penalty = RewTerm( + func=mdp.joint_pos_penalty, + weight=0.0, + params={ + "command_name": "base_velocity", + "asset_cfg": SceneEntityCfg("robot", joint_names=".*"), + "stand_still_scale": 5.0, + "velocity_threshold": 0.5, + "command_threshold": 0.1, + }, + ) + + knee_joint_pos_penalty = RewTerm( + func=mdp.joint_pos_penalty, + weight=0.0, + params={ + "command_name": "base_velocity", + "asset_cfg": SceneEntityCfg("robot", joint_names=".*"), + "stand_still_scale": 5.0, + "velocity_threshold": 0.5, + "command_threshold": 0.1, + }, + ) + + + + + wheel_vel_penalty = RewTerm( + func=mdp.wheel_vel_penalty, + weight=0.0, + params={ + "asset_cfg": SceneEntityCfg("robot", joint_names=""), + "sensor_cfg": SceneEntityCfg("contact_forces", body_names=""), + "command_name": "base_velocity", + "velocity_threshold": 0.5, + "command_threshold": 0.1, + }, + ) + + joint_mirror = RewTerm( + func=mdp.joint_mirror, + weight=0.0, + params={ + "asset_cfg": SceneEntityCfg("robot"), + "mirror_joints": [["FR.*", "RL.*"], ["FL.*", "RR.*"]], + }, + ) + + action_mirror = RewTerm( + func=mdp.action_mirror, + weight=0.0, + params={ + "asset_cfg": SceneEntityCfg("robot"), + "mirror_joints": [["FR.*", "RL.*"], ["FL.*", "RR.*"]], + }, + ) + + action_sync = RewTerm( + func=mdp.action_sync, + weight=0.0, + params={ + "asset_cfg": SceneEntityCfg("robot"), + "joint_groups": [ + ["FR_hip_joint", "FL_hip_joint", "RL_hip_joint", "RR_hip_joint"], + ["FR_thigh_joint", "FL_thigh_joint", "RL_thigh_joint", "RR_thigh_joint"], + ["FR_calf_joint", "FL_calf_joint", "RL_calf_joint", "RR_calf_joint"], + ], + }, + ) + + # Action penalties + applied_torque_limits = RewTerm( + func=mdp.applied_torque_limits, + weight=0.0, + params={"asset_cfg": SceneEntityCfg("robot", joint_names=".*")}, + ) + action_rate_l2 = RewTerm(func=mdp.action_rate_l2, weight=0.0) + # smoothness_1 = RewTerm(func=mdp.smoothness_1, weight=0.0) # Same as action_rate_l2 + # smoothness_2 = RewTerm(func=mdp.smoothness_2, weight=0.0) # Unvaliable now + + # Contact sensor + undesired_contacts = RewTerm( + func=mdp.undesired_contacts, + weight=0.0, + params={ + "sensor_cfg": SceneEntityCfg("contact_forces", body_names=""), + "threshold": 1.0, + }, + ) + contact_forces = RewTerm( + func=mdp.contact_forces, + weight=0.0, + params={"sensor_cfg": SceneEntityCfg("contact_forces", body_names=""), "threshold": 100.0}, + ) + + # Velocity-tracking rewards + track_lin_vel_xy_exp = RewTerm( + func=mdp.track_lin_vel_xy_exp, weight=0.0, params={"command_name": "base_velocity", "std": math.sqrt(0.5)} + ) + track_ang_vel_z_exp = RewTerm( + func=mdp.track_ang_vel_z_exp, weight=0.0, params={"command_name": "base_velocity", "std": math.sqrt(0.5)} + ) + + # Others + # feet_air_time = RewTerm( + # func=mdp.feet_air_time, + # weight=0.0, + # params={ + # "command_name": "base_velocity", + # "threshold": 0.5, + # "sensor_cfg": SceneEntityCfg("contact_forces", body_names=""), + # }, + # ) + + feet_air_time = RewTerm( + func=mdp.feet_air_time_including_ang_z, + weight=0.0, + params={ + "command_name": "base_velocity", + "threshold": 0.5, + "sensor_cfg": SceneEntityCfg("contact_forces", body_names=""), + }, + ) + + feet_air_time_variance = RewTerm( + func=mdp.feet_air_time_variance_penalty, + weight=0, + params={"sensor_cfg": SceneEntityCfg("contact_forces", body_names="")}, + ) + + feet_gait = RewTerm( + func=mdp.GaitReward, + weight=0.0, + params={ + "std": math.sqrt(0.5), + "command_name": "base_velocity", + "max_err": 0.2, + "velocity_threshold": 0.5, + "command_threshold": 0.1, + "synced_feet_pair_names": (("", ""), ("", "")), + "asset_cfg": SceneEntityCfg("robot"), + "sensor_cfg": SceneEntityCfg("contact_forces"), + }, + ) + + feet_contact = RewTerm( + func=mdp.feet_contact, + weight=0.0, + params={ + "sensor_cfg": SceneEntityCfg("contact_forces", body_names=""), + "command_name": "base_velocity", + "expect_contact_num": 2, + }, + ) + + feet_contact_without_cmd = RewTerm( + func=mdp.feet_contact_without_cmd, + weight=0.0, + params={ + "sensor_cfg": SceneEntityCfg("contact_forces", body_names=""), + "command_name": "base_velocity", + }, + ) + + feet_stumble = RewTerm( + func=mdp.feet_stumble, + weight=0.0, + params={ + "sensor_cfg": SceneEntityCfg("contact_forces", body_names=""), + }, + ) + + feet_slide = RewTerm( + func=mdp.feet_slide, + weight=0.0, + params={ + "sensor_cfg": SceneEntityCfg("contact_forces", body_names=""), + "asset_cfg": SceneEntityCfg("robot", body_names=""), + }, + ) + + stand_still = RewTerm( + func=mdp.stand_still_joint_deviation_l1, + weight=0, + params={ + "command_name": "base_velocity", + "asset_cfg": SceneEntityCfg("robot", joint_names=""), + }, + ) # negetive + + feet_height = RewTerm( + func=mdp.feet_height, + weight=0.0, + params={ + "asset_cfg": SceneEntityCfg("robot", body_names=""), + "tanh_mult": 2.0, + "target_height": 0.05, + "command_name": "base_velocity", + }, + ) + + feet_height_body = RewTerm( + func=mdp.feet_height_body, + weight=0.0, + params={ + "asset_cfg": SceneEntityCfg("robot", body_names=""), + "tanh_mult": 2.0, + "target_height": -0.3, + "command_name": "base_velocity", + }, + ) + + feet_distance_y_exp = RewTerm( + func=mdp.feet_distance_y_exp, + weight=0.0, + params={ + "std": math.sqrt(0.25), + "asset_cfg": SceneEntityCfg("robot", body_names=""), + "stance_width": float, + }, + ) + + # feet_distance_xy_exp = RewTerm( + # func=mdp.feet_distance_xy_exp, + # weight=0.0, + # params={ + # "std": math.sqrt(0.25), + # "asset_cfg": SceneEntityCfg("robot", body_names=""), + # "stance_length": float, + # "stance_width": float, + # }, + # ) + + upward = RewTerm(func=mdp.upward, weight=0.0) + + # lin_vel_xy_l2_with_ang_z_command = RewTerm( + # func=mdp.lin_vel_xy_l2_with_ang_z_command, + # weight=0, + # params={ + # "command_name": "base_velocity", + # "command_threshold": 0.1, + # }, + # ) # negetive + + +@configclass +class TerminationsCfg: + """Termination terms for the MDP.""" + + # MDP terminations + time_out = DoneTerm(func=mdp.time_out, time_out=True) + # command_resample + terrain_out_of_bounds = DoneTerm( + func=mdp.terrain_out_of_bounds, + params={"asset_cfg": SceneEntityCfg("robot"), "distance_buffer": 3.0}, + time_out=True, + ) + + # Contact sensor + illegal_contact = DoneTerm( + func=mdp.illegal_contact, + params={"sensor_cfg": SceneEntityCfg("contact_forces", body_names=""), "threshold": 1.0}, + ) + + bad_orientation_2 = DoneTerm(func=mdp.bad_orientation_2) + + + + + +@configclass +class CurriculumCfg: + """Curriculum terms for the MDP.""" + + terrain_levels = CurrTerm(func=mdp.terrain_levels_vel) + + command_levels = CurrTerm( + func=mdp.command_levels_vel, + params={ + "reward_term_name": "track_lin_vel_xy_exp", + "range_multiplier": (0.1, 1.0), + }, + ) + + +## +# Environment configuration +## + + +@configclass +class LocomotionVelocityRoughEnvCfg(ManagerBasedRLEnvCfg): + """Configuration for the locomotion velocity-tracking environment.""" + + # Scene settings + scene: MySceneCfg = MySceneCfg(num_envs=4096, env_spacing=2.5) + # Basic settings + observations: ObservationsCfg = ObservationsCfg() + actions: ActionsCfg = ActionsCfg() + commands: CommandsCfg = CommandsCfg() + # MDP settings + rewards: RewardsCfg = RewardsCfg() + terminations: TerminationsCfg = TerminationsCfg() + events: EventCfg = EventCfg() + curriculum: CurriculumCfg = CurriculumCfg() + + def __post_init__(self): + """Post initialization.""" + # general settings + self.decimation = 4 + self.episode_length_s = 20.0 + # simulation settings + self.sim.dt = 0.005 + self.sim.render_interval = self.decimation + self.sim.physics_material = self.scene.terrain.physics_material + self.sim.physx.gpu_max_rigid_patch_count = 10 * 2**15 + self.sim.physx.max_position_iteration_count = 4 + self.sim.physx.max_velocity_iteration_count = 1 + # update sensor update periods + # we tick all the sensors based on the smallest update period (physics update period) + if self.scene.height_scanner is not None: + self.scene.height_scanner.update_period = self.decimation * self.sim.dt + if self.scene.contact_forces is not None: + self.scene.contact_forces.update_period = self.sim.dt + + # check if terrain levels curriculum is enabled - if so, enable curriculum for terrain generator + # this generates terrains with increasing difficulty and is useful for training + if getattr(self.curriculum, "terrain_levels", None) is not None: + if self.scene.terrain.terrain_generator is not None: + self.scene.terrain.terrain_generator.curriculum = True + else: + if self.scene.terrain.terrain_generator is not None: + self.scene.terrain.terrain_generator.curriculum = False + + def disable_zero_weight_rewards(self): + """If the weight of rewards is 0, set rewards to None""" + for attr in dir(self.rewards): + if not attr.startswith("__"): + reward_attr = getattr(self.rewards, attr) + if not callable(reward_attr) and reward_attr.weight == 0: + setattr(self.rewards, attr, None) + + +def create_obsgroup_class(class_name, terms, enable_corruption=False, concatenate_terms=True): + """ + Dynamically create and register a ObsGroup class based on the given configuration terms. + + :param class_name: Name of the configuration class. + :param terms: Configuration terms, a dictionary where keys are term names and values are term content. + :param enable_corruption: Whether to enable corruption for the observation group. Defaults to False. + :param concatenate_terms: Whether to concatenate the observation terms in the group. Defaults to True. + :return: The dynamically created class. + """ + # Dynamically determine the module name + module_name = inspect.getmodule(inspect.currentframe()).__name__ + + # Define the post-init function + def post_init_wrapper(self): + setattr(self, "enable_corruption", enable_corruption) + setattr(self, "concatenate_terms", concatenate_terms) + + # Dynamically create the class using ObsGroup as the base class + terms["__post_init__"] = post_init_wrapper + dynamic_class = configclass(type(class_name, (ObsGroup,), terms)) + + # Custom serialization and deserialization + def __getstate__(self): + state = self.__dict__.copy() + return state + + def __setstate__(self, state): + self.__dict__.update(state) + + # Add custom serialization methods to the class + dynamic_class.__getstate__ = __getstate__ + dynamic_class.__setstate__ = __setstate__ + + # Place the class in the global namespace for accessibility + globals()[class_name] = dynamic_class + + # Register the dynamic class in the module's dictionary + if module_name in sys.modules: + sys.modules[module_name].__dict__[class_name] = dynamic_class + else: + raise ImportError(f"Module {module_name} not found.") + + # Return the class for external instantiation + return dynamic_class diff --git a/source/rl_training/rl_training/ui_extension_example.py b/source/rl_training/rl_training/ui_extension_example.py new file mode 100644 index 0000000..90dd2cb --- /dev/null +++ b/source/rl_training/rl_training/ui_extension_example.py @@ -0,0 +1,52 @@ +# Copyright (c) 2025 Deep Robotics +# SPDX-License-Identifier: BSD 3-Clause +# +# # Copyright (c) 2024-2025 Ziqi Fan +# SPDX-License-Identifier: Apache-2.0 + +# Copyright (c) 2022-2025, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import omni.ext + + +# Functions and vars are available to other extension as usual in python: `example.python_ext.some_public_function(x)` +def some_public_function(x: int): + print("[rl_training] some_public_function was called with x: ", x) + return x**x + + +# Any class derived from `omni.ext.IExt` in top level module (defined in `python.modules` of `extension.toml`) will be +# instantiated when extension gets enabled and `on_startup(ext_id)` will be called. Later when extension gets disabled +# on_shutdown() is called. +class ExampleExtension(omni.ext.IExt): + # ext_id is current extension id. It can be used with extension manager to query additional information, like where + # this extension is located on filesystem. + def on_startup(self, ext_id): + print("[rl_training] startup") + + self._count = 0 + + self._window = omni.ui.Window("My Window", width=300, height=300) + with self._window.frame: + with omni.ui.VStack(): + label = omni.ui.Label("") + + def on_click(): + self._count += 1 + label.text = f"count: {self._count}" + + def on_reset(): + self._count = 0 + label.text = "empty" + + on_reset() + + with omni.ui.HStack(): + omni.ui.Button("Add", clicked_fn=on_click) + omni.ui.Button("Reset", clicked_fn=on_reset) + + def on_shutdown(self): + print("[rl_training] shutdown") diff --git a/source/rl_training/setup.py b/source/rl_training/setup.py new file mode 100644 index 0000000..1060a6c --- /dev/null +++ b/source/rl_training/setup.py @@ -0,0 +1,55 @@ +# Copyright (c) 2025 Deep Robotics +# SPDX-License-Identifier: BSD 3-Clause +# +# # Copyright (c) 2024-2025 Ziqi Fan +# SPDX-License-Identifier: Apache-2.0 + +"""Installation script for the 'rl_training' python package.""" + +import os +import toml + +from setuptools import setup + +# Obtain the extension data from the extension.toml file +EXTENSION_PATH = os.path.dirname(os.path.realpath(__file__)) +# Read the extension.toml file +EXTENSION_TOML_DATA = toml.load(os.path.join(EXTENSION_PATH, "config", "extension.toml")) + +# Minimum dependencies required prior to installation +INSTALL_REQUIRES = [ + # base + "psutil", + "colorama", + "xacrodoc", + # amp + "numpy", + "pandas", + "pinocchio", + # rl + "cusrl[all]", +] + +# Installation operation +setup( + name="rl_training", + packages=["rl_training"], + author=EXTENSION_TOML_DATA["package"]["author"], + maintainer=EXTENSION_TOML_DATA["package"]["maintainer"], + url=EXTENSION_TOML_DATA["package"]["repository"], + version=EXTENSION_TOML_DATA["package"]["version"], + description=EXTENSION_TOML_DATA["package"]["description"], + keywords=EXTENSION_TOML_DATA["package"]["keywords"], + install_requires=INSTALL_REQUIRES, + license="Apache License 2.0", + include_package_data=True, + python_requires=">=3.10", + classifiers=[ + "Natural Language :: English", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Isaac Sim :: 4.5.0", + "Isaac Sim :: 5.0.0", + ], + zip_safe=False, +)