commit 07027e8fde9b19c467698fc918fa9ef5db530bab Author: syfffffjj Date: Mon Nov 2 16:41:07 2026 +0800 feat: huoyi M20 smooth/bridge training and stability docs Current code snapshot: smooth/bridge task configs, smooth random terrain, roll/pitch stability rewards, M20 URDF updates, eval profile and stability training documentation. 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..1240b9b --- /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..752ad5a --- /dev/null +++ b/docs/fire_ant_prototype_specs.md @@ -0,0 +1,201 @@ +# 火蚁样机参数简表 + +本文档依据当前火蚁仿真模型与训练配置整理： + +- 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 +``` + +### 膝关节向内折叠限制 + +膝关节角为 0 时膝连杆竖直向下，因此膝关节角就是腿相对竖直方向的偏角。前腿向内折叠对应关节角减小，后腿向内折叠对应关节角增大。训练要求向内最多折叠到竖直向下，不能更向内： + +```text +前膝 q >= 0 +后膝 q <= 0 +``` + +超出部分以二次惩罚计入奖励： + +```text +penalty = sum(max(-q_front, 0)^2) + sum(max(q_hind, 0)^2) +``` + +实现位置： + +- 惩罚函数：`mdp/rewards.py::knee_inward_fold_penalty` +- 奖励项与权重：`config/wheeled/deeprobotics_m20/rough_env_cfg.py` 中的 `knee_inward_fold_penalty` +- 当前权重为 `-1.0`，`max_inward_tilt` 为 `0.0`（即允许的最大向内倾角为 0，腿不能越过竖直向下姿态） + +向外折叠（前膝 q > 0、后膝 q < 0）不受该惩罚约束，仍由常规关节位置限制约束。若训练中仍出现越界，可增大惩罚权重；若限制过于保守影响越障，可将 `max_inward_tilt` 调大（例如 `0.1` rad）或在阶段 2 适当减小权重。 + +## 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/fire_ant_smooth_training.md b/docs/fire_ant_smooth_training.md new file mode 100644 index 0000000..2e5999a --- /dev/null +++ b/docs/fire_ant_smooth_training.md @@ -0,0 +1,561 @@ +# 火蚁平稳性训练任务文档 + +> **任务 ID**：`Smooth-Deeprobotics-M20-v0` +> **机器人**：火蚁 / Solenopsis（8 自由度轮腿） +> **训练目标**：在低频随机起伏地形上保持机身 pitch / roll 水平，同时稳定跟踪线速度与偏航速度 +> **文档依据**：当前工作区代码 + 2026-10-28 两次实际运行的 `params/env.yaml`、`params/agent.yaml`、`train_smooth*.log` +> **文档更新时间**：2026-11-02 + +## 1. 任务定位 + +平稳性训练是 Stage-1 粗糙地形策略之上的微调任务。动作空间、观测维度、机器人模型、质量/接触随机化都保持不变，因此可以从 Stage-1 checkpoint 直接 resume；变化集中在“地形”和“姿态平稳性奖励”两处。 + +### 1.1 与 Stage-1（Rough）的关键差异 + +下表按当前工作区代码对比；实际运行的参数以各运行目录下的 `params/env.yaml` 为准。 + +| 类别 | Stage-1 Rough | 平稳性 Smooth | +| --- | --- | --- | +| 地形 | 台阶、箱子、随机粗糙、斜坡混合 | 100% 低频正弦平滑起伏，无台阶、无尖边 | +| 地形初始化等级 | `max_init_terrain_level = 2` | `max_init_terrain_level = 9` | +| 重置位姿 | x/y ±1.0 m，z 偏移 0，roll/pitch ±0.3 rad | x/y ±2.0 m，z 偏移 +0.30 m，roll/pitch ±0.05 rad | +| 姿态惩罚 | `flat_orientation_l2 = -0.35`、`roll_l2 = -0.45` | `flat_orientation_l2` 置 0；`roll_l2 = -0.4`、`pitch_l2 = -0.4` | +| 机身角速度惩罚 | `ang_vel_xy_l2 = -0.05` | `ang_vel_xy_l2 = -0.1` | + +速度指令分布（`x ∈ [-2, 2] m/s`、`y = 0`、`yaw ∈ [-1.5, 1.5] rad/s`、heading 指令开启、2% 站立环境）在两种地形下保持一致。 + +### 1.2 平稳性由哪些设置共同保证 + +| 机制 | 具体设置 | 作用 | +| --- | --- | --- | +| 俯仰惩罚 | `pitch_l2 = -0.4`，弧度平方 | 直接压低机头/机尾俯仰角 | +| 横滚惩罚 | `roll_l2 = -0.4`，弧度平方 | 直接压低左右横滚角 | +| 角速度阻尼 | `ang_vel_xy_l2 = -0.1` | 抑制 roll/pitch 方向的晃动速度 | +| 竖直速度惩罚 | `lin_vel_z_l2 = -2.0` | 抑制上下弹跳 | +| 竖直姿态奖励 | `upward = +0.08`，`(1 - g_z)²` | 机身越接近水平奖励越高（直立时 g_z = -1，项值最大） | +| 机身高度约束 | `base_height_l2 = -0.5`，目标 0.45 m（相对机身下方地形均值） | 防止趴下或过度抬升 | +| 膝内折约束 | `knee_inward_fold_penalty = -1.0`，`max_inward_tilt = 0.0` | 膝不能向内侧折过竖直位置 | +| 倾倒终止 | `bad_orientation_2`，倾角超过 0.9 rad 终止 | 直接淘汰失稳策略 | +| 地形 | 平滑正弦起伏，无台阶冲击 | 只施加连续的 pitch/roll 扰动，专练平稳性 | + +## 2. 相关文件 + +| 内容 | 路径 | +| --- | --- | +| 任务注册 | `source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/config/wheeled/deeprobotics_m20/__init__.py` | +| 平稳性环境配置 | `.../deeprobotics_m20/smooth_env_cfg.py` | +| 父类环境配置 | `.../deeprobotics_m20/rough_env_cfg.py` | +| PPO 配置 | `.../deeprobotics_m20/agents/rsl_rl_ppo_cfg.py` | +| 平地/基础环境 | `source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/velocity_env_cfg.py` | +| 平稳地形生成 | `.../locomotion/velocity/mdp/terrain.py` | +| 自定义奖励/观测 | `.../locomotion/velocity/mdp/rewards.py`、`observations.py` | +| 机器人模型与执行器 | `source/rl_training/rl_training/assets/deeprobotics.py` | +| URDF | `deep_robotics_model/M20/M20_urdf/urdf/Solenopsis.urdf` | + +## 3. 实际训练运行记录 + +本次平稳性训练由两次运行接力完成，日志与 checkpoint 都保存在 `logs/rsl_rl/huoyi_smooth/`。 + +### 3.1 运行 A：小地形试跑（2026-10-28_16-50-47） + +| 项目 | 值 | +| --- | --- | +| 起点 checkpoint | `/home/suyufei/syf/huoyi_training/logs/rsl_rl/deeprobotics_m20_rough/2026-10-28_10-23-40_stage1_readapt_realmass/model_24800.pt` | +| 并行环境数 | 2048 | +| 迭代数 | 50 | +| GPU | `cuda:2` | +| 地形参数 | `height_range=(0.02, 0.05)`、`wavelength_range=(1.5, 3.0)`、`max_init_terrain_level=1` | +| 姿态奖励 | `flat_orientation_l2 = -0.8`（未拆分 roll/pitch） | +| 结束状态 | 迭代 24849，mean reward 60.34，mean episode length 919.83 | +| 产出 | `model_24800.pt`、`model_24849.pt` | + +运行 A 的目的：先在小起伏、低地形等级上确认从粗糙地形策略切换过来不会崩，同时保留一份可回退的中间策略。 + +### 3.2 运行 B：正式平稳性训练（2026-10-28_18-16-44） + +| 项目 | 值 | +| --- | --- | +| 起点 checkpoint | `logs/rsl_rl/huoyi_smooth/2026-10-28_16-50-47/model_24849.pt` | +| 并行环境数 | 4096 | +| 迭代数 | 2500（从 24849 续到 27348） | +| GPU | `cuda:0` | +| 地形参数 | `height_range=(0.10, 0.20)`、`wavelength_range=(1.0, 2.2)`、`max_init_terrain_level=9` | +| 姿态奖励 | `flat_orientation_l2 = 0`；`roll_l2 = -0.4`、`pitch_l2 = -0.4`；`ang_vel_xy_l2 = -0.1` | +| 总样本量 | 245,760,000 steps（2500 × 24 × 4096） | +| 训练耗时 | 约 2 小时 57 分（Iteration time 约 3.09–3.20 s） | +| 最终产出 | `model_27348.pt`，以及 `exported/policy.pt`、`exported/policy.onnx` | + +运行 B 才是本任务的有效训练：地形起伏幅度、波长密度和初始地形等级都明显提高，逼迫策略在连续起伏上主动调平机身。 + +### 3.3 复现命令 + +运行 A（小地形试跑）： + +```bash +python scripts/reinforcement_learning/rsl_rl/train.py \ + --task Smooth-Deeprobotics-M20-v0 \ + --num_envs 2048 \ + --max_iterations 50 \ + --headless \ + --resume \ + --load_run /home/suyufei/syf/huoyi_training/logs/rsl_rl/deeprobotics_m20_rough/2026-10-28_10-23-40_stage1_readapt_realmass \ + --checkpoint /home/suyufei/syf/huoyi_training/logs/rsl_rl/deeprobotics_m20_rough/2026-10-28_10-23-40_stage1_readapt_realmass/model_24800.pt +``` + +运行 B（从运行 A 续训）： + +```bash +python scripts/reinforcement_learning/rsl_rl/train.py \ + --task Smooth-Deeprobotics-M20-v0 \ + --num_envs 4096 \ + --max_iterations 2500 \ + --headless \ + --resume \ + --load_run /home/suyufei/syf/huoyi_training_level/logs/rsl_rl/huoyi_smooth/2026-10-28_16-50-47 \ + --checkpoint /home/suyufei/syf/huoyi_training_level/logs/rsl_rl/huoyi_smooth/2026-10-28_16-50-47/model_24849.pt +``` + +说明： + +- 恢复训练时策略/价值网络输入维度必须一致（本任务 policy 33 维、critic 223 维、动作 8 维）。 +- `--checkpoint` 是脚本实际使用的参数名；日志 `params/agent.yaml` 中对应字段名为 `load_checkpoint`。 +- 不传 `--max_iterations` 时使用配置默认值 5000；上述两次运行都用命令行覆盖了默认值。 +- 环境种子默认 42，实际运行未覆盖。 + +### 3.4 播放与导出 + +```bash +python scripts/reinforcement_learning/rsl_rl/play.py \ + --task Smooth-Deeprobotics-M20-v0 \ + --num_envs 16 \ + --checkpoint /home/suyufei/syf/huoyi_training_level/logs/rsl_rl/huoyi_smooth/2026-10-28_18-16-44/model_27348.pt +``` + +## 4. 仿真与场景设置 + +| 项目 | 值 | +| --- | --- | +| 物理步长 `sim.dt` | 0.005 s（200 Hz） | +| 控制抽帧 `decimation` | 4 | +| 控制步长 | 0.02 s（50 Hz） | +| 回合长度 `episode_length_s` | 20 s（1000 个控制步） | +| 并行环境数 `num_envs` | 4096（正式训练） | +| 环境间距 `env_spacing` | 2.5 m | +| 随机种子 | 42 | +| 设备 | `cuda:0` | +| 重力 | `(0, 0, -9.81)` | +| PhysX 位置/速度迭代 | 4 / 1 | +| CCD、稳定性增强 | 均关闭 | +| 地形物理材质 | 静摩擦 1.0、动摩擦 1.0、恢复系数 0.0，摩擦/恢复均为 multiply | +| 渲染间隔 | 4 个物理步 | +| 接触传感器更新周期 | 0.005 s，历史长度 3，`track_air_time=True` | +| 高度扫描更新周期 | 0.02 s | + +实测运行环境：IsaacLab 2.3.2 / Isaac Sim 5.1，8 × RTX 4090（正式训练使用 0 号卡），72 核 Xeon Platinum 8352V，约 512 GB 内存。 + +## 5. 机器人模型 + +| 项目 | 值 | +| --- | --- | +| URDF | `deep_robotics_model/M20/M20_urdf/urdf/Solenopsis.urdf` | +| 自由度 | 8（4 个膝关节 + 4 个轮关节） | +| 整机仿真质量 | 约 23.69 kg | +| 轮半径 | 0.11 m（奖励计算同样使用 0.11） | +| 轮距 | 0.50 m（奖励计算同样使用 0.50） | +| 根初始高度 | (0, 0, 0.32) m | +| 初始关节角 | 前膝 `+0.5 rad`，后膝 `-0.5 rad`，轮 `0` | +| 软关节限位系数 | 0.9 | +| 自碰撞 | 关闭 | + +### 5.1 关节顺序 + +```text +膝关节：fl_knee_joint, fr_knee_joint, hl_knee_joint, hr_knee_joint +轮关节：fl_wheel_joint, fr_wheel_joint, hl_wheel_joint, hr_wheel_joint +``` + +### 5.2 执行器（DelayedPDActuator） + +| 参数 | 膝关节 | 轮关节 | +| --- | ---: | ---: | +| `effort_limit` | 75.0 N·m | 28.68 N·m | +| `velocity_limit` | 17.8024 rad/s | 109.956 rad/s | +| `stiffness` | 80.0 | 0.0 | +| `damping` | 2.0 | 0.6 | +| `friction` | 0.0 | 1.0 | +| `armature` | 0.0 | 0.00243216 | +| 动作延迟 | 0–5 个物理步 | 0–5 个物理步 | + +膝关节是位置控制，轮关节是纯速度控制（stiffness 为 0）。执行器的刚度、阻尼会在每次 reset 时按 0.85–1.15 缩放。 + +## 6. 动作空间 + +动作为 8 维，四膝位置 + 四轮速度： + +| 动作项 | 关节 | 缩放 | 默认偏置 | 裁剪 | 说明 | +| --- | --- | ---: | --- | --- | --- | +| `joint_pos` | 4 个膝关节 | 0.25 | `use_default_offset=True` | ±100 | 目标角 = 默认角 + 0.25 × 动作 | +| `joint_vel` | 4 个轮关节 | 5.0 | 默认 0 | ±100 | 目标轮速 = 5.0 × 动作 (rad/s) | + +动作顺序被 `preserve_order=True` 固定，分别对应配置中的膝关节顺序与轮关节顺序，避免按正则匹配后顺序漂移导致 resume 时动作错位。 + +## 7. 地形设置 + +平稳性地形由 `mdp/terrain.py` 中的 `smooth_random_terrain` 生成，核心是一个“多个低频正弦波叠加”的连续高度场。 + +### 7.1 生成公式 + +对每个地形块： + +1. 在 `8 m × 8 m` 的网格上生成坐标（水平分辨率 0.05 m）； +2. 叠加 `num_waves = 3` 个正弦波，每个波随机采样： + - 方向 `direction ~ U(0, 2π)` + - 波长 `wavelength ~ U(1.0, 2.2) m` + - 相位 `phase ~ U(0, 2π)` + - 权重 `weight ~ U(0.5, 1.0)` +3. 减去高度均值； +4. 按 `max(|h|)` 归一化到当前课程幅度 `amplitude`； +5. 按 `vertical_scale = 0.005 m` 量化为 int16 高度场。 + +课程幅度： + +```text +amplitude = 0.10 + difficulty × (0.20 - 0.10) # difficulty ∈ [0, 1] +``` + +即课程起点起伏峰值 0.10 m，终点 0.20 m。地形没有台阶和尖锐边缘，机器人受到的是连续的 pitch/roll 扰动。 + +### 7.2 地形参数 + +| 参数 | 值 | +| --- | --- | +| 地形块尺寸 | 8.0 m × 8.0 m | +| 行数 × 列数 | 10 × 20 | +| 边界宽度 / 高度 | 2.0 m / 1.0 m | +| 水平分辨率 | 0.05 m | +| 垂直分辨率 | 0.005 m | +| 是否缓存 | 否（`use_cache=False`） | +| 课程 | 开启（`curriculum=True`），难度 0 → 1 | +| 子地形比例 | `smooth_random` = 1.0 | +| 起伏峰值范围 | 0.10 → 0.20 m（课程起止） | +| 波长范围 | 1.0 → 2.2 m | +| 正弦波数量 | 3 | +| 子地形边界宽度 | 0.25 m | +| 初始地形等级上限 | 9（10 行中的最高等级） | + +### 7.3 地形课程规则 + +使用 IsaacLab 标准的 `terrain_levels_vel`： + +- **升级**：本回合行走距离 > 地形块尺寸的一半（4 m）； +- **降级**：行走距离 < 指令速度 × 20 s × 0.5； +- 每个环境初始等级从 `[0, max_init_terrain_level]` 均匀采样。 + +正式训练开始时平均地形等级约 4.95，结束时升至 5.96（满级 9），说明策略在持续向更难地形推进。 + +### 7.4 地形材质 + +| 参数 | 值 | +| --- | --- | +| 静摩擦 | 1.0 | +| 动摩擦 | 1.0 | +| 恢复系数 | 0.0 | +| 摩擦组合方式 | multiply | +| 恢复组合方式 | multiply | + +轮子材质另有随机化（见第 11 节），因此轮地有效摩擦主要由轮子材质范围决定。 + +## 8. 观测空间 + +采用非对称 Actor-Critic：策略只看本体感知，价值网络额外看地形高度扫描。 + +### 8.1 Policy 观测（33 维） + +| 序号 | 观测项 | 维度 | 噪声 | 缩放 | 说明 | +| ---: | --- | ---: | --- | ---: | --- | +| 0 | `base_ang_vel` | 3 | U(-0.2, 0.2) | 0.25 | 机身角速度 | +| 1 | `projected_gravity` | 3 | U(-0.05, 0.05) | 1.0 | 重力在机体系的投影，直接反映 pitch/roll | +| 2 | `velocity_commands` | 3 | 无 | 1.0 | 线速度 x、y 与偏航速度指令 | +| 3 | `joint_pos` | 8 | U(-0.01, 0.01) | 1.0 | 4 个膝相对角 + 4 个轮关节占位（轮位置恒置 0） | +| 4 | `joint_vel` | 8 | U(-1.5, 1.5) | 0.05 | 4 膝 + 4 轮关节速度 | +| 5 | `actions` | 8 | 无 | 1.0 | 上一步动作 | + +合计 3 + 3 + 3 + 8 + 8 + 8 = 33 维。 + +**策略是“盲”的**：`base_lin_vel` 和 `height_scan` 在 policy 组中被显式置空，机器人只能依靠 IMU、关节编码器和指令来保持平稳，不能直接“看到”地形。 + +### 8.2 Critic 观测（223 维） + +| 序号 | 观测项 | 维度 | 说明 | +| ---: | --- | ---: | --- | +| 0 | `base_lin_vel` | 3 | 机身线速度（特权信息） | +| 1 | `base_ang_vel` | 3 | 机身角速度 | +| 2 | `projected_gravity` | 3 | 重力投影 | +| 3 | `velocity_commands` | 3 | 速度指令 | +| 4 | `joint_pos` | 8 | 关节位置 | +| 5 | `joint_vel` | 8 | 关节速度 | +| 6 | `actions` | 8 | 上一步动作 | +| 7 | `height_scan` | 187 | 机身下方 1.6 m × 1.0 m 地形高度网格（17 × 11，分辨率 0.1 m），裁剪到 ±1.0 | + +critic 组不注入噪声。合计 3 + 3 + 3 + 3 + 8 + 8 + 8 + 187 = 223 维。 + +### 8.3 传感器 + +| 传感器 | 配置 | +| --- | --- | +| `contact_forces` | 覆盖 `/World/envs/env_.*/Robot/.*`，更新周期 0.005 s，历史长度 3，记录 air time | +| `height_scanner` | 挂在 `base_link`，从上方 20 m 向下投射，1.6 m × 1.0 m，分辨率 0.1 m，yaw 对齐，187 个点 | +| `height_scanner_base` | 挂在 `base_link`，0.1 m × 0.1 m，分辨率 0.05 m，仅用于 `base_height_l2` 估计脚下地形高度 | + +## 9. 奖励设置（22 项） + +奖励项权重与最终训练日志值如下。`Episode_Reward` 是 IsaacLab 奖励管理器的日志量：回合累积奖励（已按 `dt = 0.02` 积分）除以 20 s，因此它与权重和回合长度都相关，只适合看趋势和相对大小。 + +### 9.1 姿态平稳性 + +| 奖励项 | 权重 | 关键参数 | 末次日志值 | 作用 | +| --- | ---: | --- | ---: | --- | +| `lin_vel_z_l2` | -2.0 | 无 | -0.1122 | 抑制机身竖直方向弹跳 | +| `ang_vel_xy_l2` | -0.1 | 无 | -0.0815 | 抑制 roll/pitch 角速度 | +| `roll_l2` | -0.4 | 欧拉角 roll，弧度平方 | -0.0043 | 直接惩罚横滚角 | +| `pitch_l2` | -0.4 | 欧拉角 pitch，弧度平方 | -0.0049 | 直接惩罚俯仰角 | +| `base_height_l2` | -0.5 | 目标高度 0.45 m，`height_scanner_base` 估计脚下地形 | -0.0208 | 维持机身高度 | +| `upward` | +0.08 | `(1 - g_z)²`，直立时 g_z = -1 | +0.2872 | 机身越水平奖励越高 | + +`flat_orientation_l2` 在平稳性任务中被置 0（不参与训练），由 `roll_l2` + `pitch_l2` 取代。运行 A 及更早的设置中 `flat_orientation_l2` 为 -0.8；正式训练把它拆成 -0.4 / -0.4，便于分别调节俯仰和横滚。 + +### 9.2 速度与偏航跟踪 + +| 奖励项 | 权重 | 关键参数 | 末次日志值 | 作用 | +| --- | ---: | --- | ---: | --- | +| `track_lin_vel_xy_exp` | +2.0 | std = 0.7071（即 √0.5） | +1.6189 | 跟踪 xy 线速度 | +| `track_ang_vel_z_exp` | +1.5 | std = 0.5 | +1.0915 | 跟踪偏航角速度 | +| `track_diff_wheel_yaw_exp` | +1.0 | std = 0.6，轮半径 0.11，轮距 0.50，左轮 (0,2)，右轮 (1,3)，指令阈值 0.05 | +0.5449 | 用左右轮速差估计偏航，鼓励差速转向 | +| `lin_vel_xy_yaw_command_l2` | -0.4 | 偏航阈值 0.25、线速度阈值 0.15 | -0.0012 | 原地转向时惩罚线速度漂移 | + +### 9.3 关节、动作与功率 + +| 奖励项 | 权重 | 作用关节 | 末次日志值 | 作用 | +| --- | ---: | --- | ---: | --- | +| `joint_torques_l2` | -2.5e-5 | 4 膝 | -0.0356 | 抑制膝力矩 | +| `joint_acc_l2` | -3e-6 | 4 膝 | -0.1570 | 抑制膝加速度 | +| `joint_acc_wheel_l2` | -1e-7 | 4 轮 | -0.0413 | 轻微抑制轮加速度 | +| `joint_pos_limits` | -5.0 | 4 膝 | -0.0000 | 关节限位惩罚 | +| `joint_power` | -2e-5 | 4 膝 | -0.0014 | 抑制膝功率 | +| `knee_joint_pos_penalty` | -0.1 | 4 膝 | -0.0751 | 站立或低速时把膝拉回默认姿态（`stand_still_scale=5`、速度阈值 0.5、指令阈值 0.1） | +| `knee_inward_fold_penalty` | -1.0 | 前膝 / 后膝 | -0.0002 | 禁止膝向内折过竖直位置（`max_inward_tilt = 0.0`） | +| `action_rate_l2` | -0.015 | 全部动作 | -0.1226 | 动作平滑 | +| `stand_still` | -2.0 | 4 膝 | -0.0495 | 站立指令下关节偏离默认姿态的惩罚 | + +`knee_inward_fold_penalty` 的方向约定：膝关节角为 0 时膝连杆竖直向下；前膝角减小为内折，后膝角增大为内折。`max_inward_tilt = 0.0` 表示不允许越过竖直位置，向外折叠仍由关节限位约束。 + +### 9.4 接触与轮地 + +| 奖励项 | 权重 | 关键参数 | 末次日志值 | 作用 | +| --- | ---: | --- | ---: | --- | +| `undesired_contacts` | -1.0 | 除轮子外的所有 body，力阈值 1.0 N | -0.0002 | 禁止机身、膝等部位触地拖行 | +| `contact_forces` | -1.5e-4 | 4 个轮子，阈值 100 N | -0.0166 | 抑制轮地过大冲击力 | +| `feet_contact_without_cmd` | +0.1 | 4 个轮子 | +0.0069 | 无速度指令时保持轮子着地 | + +### 9.5 被关闭的奖励项 + +以下奖励项权重为 0，训练时被移除或置空：`is_terminated`、`flat_orientation_l2`、`body_lin_acc_l2`、`joint_vel_l2`、`joint_vel_wheel_l2`、`joint_torques_wheel_l2`、`joint_vel_limits`、`feet_air_time`、`feet_air_time_variance`、`feet_contact`、`feet_stumble`、`feet_slide`、`feet_height`、`feet_height_body`、`feet_distance_y_exp`、`feet_gait`、`wheel_vel_penalty`、`joint_mirror`、`action_mirror`、`action_sync`、`applied_torque_limits`、`joint_deviation_l1`、`stand_still_without_cmd`、`joint_pos_penalty`、`hipx_joint_pos_penalty`、`hipy_joint_pos_penalty`、`command_progress`、`support_contact_progress`、`support_contact_excess`。 + +注意：平稳性任务没有 `illegal_contact` 终止，机身/膝触地不会立即结束回合，而是通过 `undesired_contacts` 持续扣分；这与 Stage-2 允许短时支撑的设计不同。 + +## 10. 终止条件 + +| 终止项 | 参数 | `time_out` | 说明 | +| --- | --- | --- | --- | +| `time_out` | 20 s | 是 | 正常回合结束 | +| `terrain_out_of_bounds` | 距地形块原点 3.0 m 缓冲 | 是 | 跑出当前地形块，按超时处理，不计为摔倒 | +| `bad_orientation_2` | 倾角上限 0.9 rad（约 51.6°） | 否 | 机身过度倾斜，判定失稳 | +| `illegal_contact` | — | — | 已禁用，机身/膝触地不终止 | + +`bad_orientation_2` 使用 `mdp.bad_orientation`，比较机身 z 轴与重力方向的夹角。它和 `roll_l2`、`pitch_l2` 形成“软惩罚 + 硬终止”的双层约束。 + +## 11. 速度指令 + +| 参数 | 值 | +| --- | --- | +| 指令类 | `UniformThresholdVelocityCommand` | +| 重采样周期 | 10 s | +| `lin_vel_x` | -2.0 ~ 2.0 m/s | +| `lin_vel_y` | 0.0（不允许横移） | +| `ang_vel_z` | -1.5 ~ 1.5 rad/s | +| heading 指令 | 开启，`heading_control_stiffness = 0.5` | +| heading 范围 | -π ~ π | +| 站立环境比例 | 2% | +| heading 环境比例 | 100% | + +`UniformThresholdVelocityCommand` 会在采样后把平面线速度模长 ≤ 0.2 m/s 的指令清零，因此低速线速度指令会被视为 0；偏航速度不受此阈值影响。这保证了“要么有明显前进/后退，要么原地转向”，减少模棱两可的小速度指令。 + +## 12. 域随机化设置 + +所有事件均由父类 `EventCfg` 继承，平稳性任务未额外修改（除重置位姿在 `smooth_env_cfg.py` 中覆盖）。 + +| 事件 | 模式 | 对象 | 参数 | +| --- | --- | --- | --- | +| `randomize_rigid_body_mass` | startup | 除 `base_link` 外的所有 body | 质量 scale 0.85–1.15，重算惯量 | +| `randomize_rigid_body_mass_base` | startup | `base_link` | 质量 add -1.0 ~ +3.0 kg，重算惯量 | +| `randomize_rigid_body_inertia` | startup | 所有 body | 惯量 scale 0.85–1.15 | +| `randomize_com_positions` | startup | `base_link` | COM x/y ±0.03 m，z ±0.02 m | +| `randomize_apply_external_force_torque` | reset | `base_link` | 力 ±10 N，力矩 ±10 N·m | +| `randomize_reset_joints` | reset | 全部关节 | 位置 scale 1.0–1.0（等同不随机），速度 0 | +| `randomize_actuator_gains` | reset | 4 膝 + 4 轮 | 刚度、阻尼均 scale 0.85–1.15，均匀分布 | +| `randomize_reset_base` | reset | 机器人根 | 见下方重置位姿 | +| `randomize_push_robot` | interval 10–15 s | 机器人根 | x/y 速度各 ±0.5 m/s | +| `randomize_wheel_material` | startup | `.*_wheel` | 静摩擦 0.9–1.4，动摩擦 0.6–0.9，恢复 0–0.05，1024 桶 | +| `randomize_support_material` | startup | `base_link`、`.*_knee` | 静摩擦 0.4–0.8，动摩擦 0.25–0.4，恢复 0–0.05，1024 桶 | +| `randomize_rigid_body_material` | — | 全部 body | 已禁用，由上面两个材质事件替代 | + +### 12.1 重置位姿 + +| 轴 | 范围 | +| --- | --- | +| x | -2.0 ~ 2.0 m（相对默认根位置偏移） | +| y | -2.0 ~ 2.0 m（相对默认根位置偏移） | +| z | +0.30 m（相对默认根高度的偏移） | +| roll | -0.05 ~ 0.05 rad | +| pitch | -0.05 ~ 0.05 rad | +| yaw | -π ~ π | +| 线速度 x | -0.1 ~ 0.1 m/s | +| 其余速度 | 0 | + +注意：IsaacLab 的 `reset_root_state_uniform` 把 `z` 当作**相对默认根状态的偏移**。机器人 `init_state` 的根高度是 0.32 m，因此重置时实际高度约为 0.32 + 0.30 = 0.62 m（相对地形原点），会先经历一次落地再进入正常控制。若希望原地生成，应把 z 改为 0.0。 + +## 13. PPO 超参数 + +| 参数 | 值 | +| --- | --- | +| Runner | `OnPolicyRunner`（RSL-RL） | +| 策略结构 | ActorCritic，MLP 512-256-128，ELU | +| Actor 输入 / 输出 | 33 / 8 | +| Critic 输入 / 输出 | 223 / 1 | +| 初始噪声标准差 | 1.0，`noise_std_type="log"` | +| `num_steps_per_env` | 24（单环境 rollout 0.48 s） | +| 每轮样本数 | 24 × 4096 = 98,304 | +| `max_iterations` | 2500（命令行覆盖；配置默认 5000） | +| `save_interval` | 100 | +| `num_learning_epochs` | 5 | +| `num_mini_batches` | 4（mini-batch 24,576 样本） | +| 学习率 | 1e-3，adaptive schedule，`desired_kl = 0.01` | +| 折扣 / GAE | `gamma = 0.99`，`lam = 0.95` | +| PPO clip | 0.2 | +| 熵系数 | 0.01 | +| 价值损失系数 | 1.0，使用 clipped value loss | +| 梯度裁剪 | 1.0 | +| 观测归一化 | 关闭（`empirical_normalization=False`） | +| 动作裁剪 | ±100 | +| 日志 | TensorBoard | +| 实验名 | `huoyi_smooth` | +| 随机种子 | 42 | + +单轮耗时约 3.1–3.2 s，采集速度约 8.8k steps/s，2500 轮合计约 2 小时 57 分。 + +## 14. 训练结果（运行 B，迭代 27348） + +| 指标 | 值 | +| --- | --- | +| 有效迭代 | 2500（24849 → 27348） | +| 总样本量 | 245,760,000 steps | +| 训练耗时 | 2 小时 57 分 08 秒 | +| 平均回合奖励 | 55.95 | +| 平均回合长度 | 886.40 步（约 17.73 s / 20 s） | +| 平均地形等级 | 5.96 / 9 | +| 线速度跟踪误差 `error_vel_xy` | 0.3823 m/s | +| 偏航速度跟踪误差 `error_vel_yaw` | 0.4167 rad/s | +| `time_out` 终止比例 | 80.10% | +| `terrain_out_of_bounds` 比例 | 18.89% | +| `bad_orientation_2` 比例 | 1.04% | +| 价值损失 | 0.0929 | +| 策略替代损失 | 0.0028 | +| 熵损失 | 7.4557 | +| 平均动作噪声标准差 | 0.68 | + +### 14.1 奖励分解（末次日志） + +正向项合计约 +3.55，负向项合计约 -0.72，主要贡献来自： + +| 奖励项 | 末次日志值 | +| --- | ---: | +| `track_lin_vel_xy_exp` | +1.6189 | +| `track_ang_vel_z_exp` | +1.0915 | +| `track_diff_wheel_yaw_exp` | +0.5449 | +| `upward` | +0.2872 | +| `feet_contact_without_cmd` | +0.0069 | +| `joint_acc_l2` | -0.1570 | +| `action_rate_l2` | -0.1226 | +| `lin_vel_z_l2` | -0.1122 | +| `ang_vel_xy_l2` | -0.0815 | +| `knee_joint_pos_penalty` | -0.0751 | +| `stand_still` | -0.0495 | +| `joint_acc_wheel_l2` | -0.0413 | +| `joint_torques_l2` | -0.0356 | +| `base_height_l2` | -0.0208 | +| `contact_forces` | -0.0166 | +| `pitch_l2` | -0.0049 | +| `roll_l2` | -0.0043 | +| 其他 | 均接近 0 | + +### 14.2 关于 resume 首轮的“异常低值” + +运行 B 的第一条日志（迭代 24849）显示 mean reward -0.94、mean episode length 12.99。这是 resume 时 `init_at_random_ep_len=True` 造成的统计瞬态：所有环境被随机放到回合计时器的任意位置，第一轮 rollout 内瞬间完成的“残余回合”被计入均值，并不代表策略崩溃。后续迭代迅速恢复到正常运行区间。 + +### 14.3 与运行 A 的对比 + +| 指标 | 运行 A（小地形） | 运行 B（正式） | +| --- | ---: | ---: | +| 并行环境 | 2048 | 4096 | +| 地形峰值 | 0.02–0.05 m | 0.10–0.20 m | +| 初始地形等级上限 | 1 | 9 | +| 平均奖励 | 60.34 | 55.95 | +| 平均回合长度 | 919.83 | 886.40 | + +运行 A 的奖励更高，但地形明显更容易，两者不能直接比较；运行 B 是在更难的起伏地形和更高初始等级下取得的 55.95。 + +## 15. 调参要点与注意事项 + +1. **地形难度要分级推进。** 运行 A（0.02–0.05 m、初始等级 1）先确认策略不崩，运行 B 再切到 0.10–0.20 m、初始等级 9。直接从大起伏开始容易出现早期大量 `bad_orientation`。 +2. **roll / pitch 分开调。** 平稳性任务把原来的 `flat_orientation_l2` 拆成 `roll_l2` 与 `pitch_l2`，当前各 -0.4；如果某一轴仍晃动，可单独加大该轴权重，而不会同时把另一轴压得过死。 +3. **角速度惩罚不要太小。** `ang_vel_xy_l2` 从 Stage-1 的 -0.05 提到 -0.1，是抑制“角度不大但持续晃”的关键项。 +4. **resume 时不要改观测/动作维度或顺序。** 当前契约为 policy 33 维、critic 223 维、动作 8 维，且膝关节顺序被 `preserve_order=True` 固定；改顺序会导致动作错位但训练不报错。 +5. **`base_height_l2` 的 0.45 m 是相对 `height_scanner_base` 所测脚下地形均值的高度。** 修改该扫描器的尺寸或分辨率会改变高度目标的实际含义。 +6. **重置 z 是偏移量。** 当前 z = +0.30 m 是叠加在默认根高度 0.32 m 上的偏移，实际约 0.62 m；若不需要落地过程，应改为 0.0。 +7. **奖励里的轮半径 0.11 m、轮距 0.50 m 要与 URDF 同步。** `track_diff_wheel_yaw_exp` 用这两个常数把轮速差换算成偏航角速度，URDF 改轮径后必须同步修改。 +8. **`undesired_contacts` 阈值只有 1.0 N。** 机身、膝几乎一碰地就会扣分，这是“不允许趴地爬行”的强约束；若发现策略过度保守，可先检查这一项。 +9. **`terrain_out_of_bounds` 占 18.89%。** 它被标记为 `time_out=True`，不算摔倒，但会截断回合。8 m × 8 m 地形块配合最高 2 m/s 的指令，高速直行时容易跑出边界；若希望提高有效回合长度，可增大块尺寸或 `distance_buffer`。 +10. **策略是盲的。** `base_lin_vel` 和 `height_scan` 只在 critic 中；部署时不需要地形感知，但也意味着策略完全依赖 IMU 与关节反馈来保持平稳。 +11. **控制频率与延迟。** 50 Hz 控制、执行器 0–5 个物理步延迟（0–25 ms），仿真里的平稳性结论不能直接等同于实机，实机部署前应确认延迟与电机响应。 +12. **不要用运行 A 的奖励值判断平稳性优劣。** 两次运行的地形分布不同，平均奖励不可直接横比。 + +## 16. 产物与日志 + +| 内容 | 路径 | +| --- | --- | +| 运行 A 目录 | `logs/rsl_rl/huoyi_smooth/2026-10-28_16-50-47/` | +| 运行 B 目录 | `logs/rsl_rl/huoyi_smooth/2026-10-28_18-16-44/` | +| 最终 checkpoint | `logs/rsl_rl/huoyi_smooth/2026-10-28_18-16-44/model_27348.pt` | +| 导出策略 | `.../2026-10-28_18-16-44/exported/policy.pt`、`policy.onnx` | +| 运行参数 | `.../2026-10-28_18-16-44/params/env.yaml`、`agent.yaml` | +| TensorBoard | `.../2026-10-28_18-16-44/events.out.tfevents.*` | +| 代码快照 | `.../2026-10-28_18-16-44/git/IsaacLab-2.3.2.diff`、`huoyi_training_level.diff` | +| 运行 A 日志 | `train_smooth50.log` | +| 运行 B 日志 | `train_smooth4096_1028_1816.log` | +| 未进入训练的启动日志 | `train_smooth4096_1028_1812.log`（启动阶段退出，未产生迭代） | + +checkpoint 命名与迭代号对应：`model_24900.pt` … `model_27348.pt`，每 100 轮保存一次，最后一次为 27348。 + +## 17. 可选的后续实验 + +以下仅为建议，未在本次训练中执行： + +1. 把地形峰值从 0.20 m 提高到 0.25 m，或把波长下限继续降到 0.8 m，观察 `bad_orientation` 与 `error_vel_xy` 的变化。 +2. 针对 `terrain_out_of_bounds` 偏高的问题，尝试 10 m × 10 m 地形块或更大的 `distance_buffer`，减少高速直行被截断。 +3. 在保持 policy 观测不变的前提下，尝试把 `ang_vel_xy_l2` 从 -0.1 调到 -0.15，观察横滚/俯仰速度是否进一步下降。 +4. 单独统计 roll、pitch 的 RMS 与峰值，作为“平稳性”的量化指标，而不是只看总奖励。 +5. 若后续 URDF 轮径或轮距变化，同步更新 `track_diff_wheel_yaw_exp` 的 `wheel_radius`、`track_width`，并重新校准 `base_height_l2` 的目标高度。 + +--- + +**文档维护**：修改 `smooth_env_cfg.py`、`rough_env_cfg.py`、`terrain.py`、`rewards.py` 或 PPO 配置后，请同步更新本文对应章节，并在“实际训练运行记录”中追加新的运行目录与关键指标。 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/eval_kit/README.md b/eval_kit/README.md new file mode 100644 index 0000000..f598c4f --- /dev/null +++ b/eval_kit/README.md @@ -0,0 +1,561 @@ +# eval_kit —— 可搬运的机器人性能评价套件 + +**这份套件与 `scripts/reinforcement_learning/rsl_rl/eval.py` 并存，不导入也不修改它。** +两者的指标名完全一致、取点一致、公式一致，所以输出可以直接互相对照。 + +## 为什么要另写一套 + +`eval.py` 是一个能跑通的评估脚本，但它把三件事混在一起：读仿真环境、算指标、格式化输出。 +其中「读环境」必须知道具体机器人长什么样——它写死了 `.*_wheel` 和 `["base_link", ".*_knee"]`。 +于是它只能评估火蚁；而且因为 IsaacLab 的 `find_bodies` 对匹配不到的正则会直接抛 +`ValueError`，换一个没有轮子的机器人跑会在进循环之前就崩，不是给出一个错的分数。 + +本套件把三层切开，让不依赖框架的部分可以单独测试、单独搬运： + +``` +eval_kit/ +├── core/ ★ 零框架依赖（只用 numpy + 标准库），可整块拷到别的项目 +│ ├── schema.py │ RobotProfile / EvalProtocol / Trace 三个数据契约 +│ ├── metrics.py │ 指标纯函数（唯一算数的地方） +│ ├── aggregation.py │ 上限网格、score、跨种子 mean±std +│ ├── terrain_spec.py │ 地形等级规格（纯数据） +│ ├── report.py │ 格式化与指标筛选 +│ └── naming.py │ 容错版的名字/正则解析 +├── isaaclab/ × 依赖 isaaclab，唯一接触仿真内部的地方 +│ ├── env_setup.py │ 把协议与地形规格翻译成 env_cfg 改动 +│ ├── recorder.py │ 跑 rollout 并采集成 Trace +│ └── policy.py │ checkpoint 解析与策略加载 +├── profiles.py ★ 机器人形态描述：加机器人 = 加一条记录 +├── protocols.py ★ 评估协议预设：改口径 = 改这里 +└── tests/ + ├── test_core.py ★ 合成轨迹测试指标（不需要仿真） + └── test_cli.py ★ 命令行参数与口径覆盖自检（不需要仿真） +``` + +打 ★ 的是「可套用」的部分：`core/`、`profiles.py`、`protocols.py` 与两个测试文件 +可以在没有 Isaac Sim、没有 GPU 的机器上运行。将来要接别的框架（legged_gym、mujoco）， +只需要再写一个适配层产出 `Trace`，`core/` 一行都不用改。 + +## 用法 + +### 1. 激活火蚁项目的虚拟环境 + +```bash +cd ~/syf/IsaacLab-2.3.2 +source env_isaaclab_py311/bin/activate +cd ~/syf/huoyi_training # 评估脚本在仓库根目录跑 +``` + +验证环境是否正确：`python -c "import isaaclab"` 应当无报错，且 `isaaclab.__file__` +指向 `IsaacLab-2.3.2/source/isaaclab/isaaclab/__init__.py`（可编辑安装）。 + +注意有两个容易踩空的近邻环境：系统 `python3`（3.10）和 conda 的 +`env_rl_training` 都**不能**跑——前者没有 isaaclab/isaacsim，后者有 isaacsim 5.1 +但没有 isaaclab。只有 `env_isaaclab_py311` 是完整的。仓库根目录是必须的， +因为 `logs/` 查找用的是相对 cwd 的路径。 + +### 2. 让 IsaacLab 的日志目录可写 + +`/tmp/isaaclab` 属于机器上的其它用户，默认会 `PermissionError`。 +`configure_logging` 用的是 `tempfile.gettempdir()/isaaclab/logs`，而 `gettempdir()` +尊重 `TMPDIR`，所以设一下即可（不必去动别人的目录）： + +```bash +export TMPDIR=/home/suyufei/tmp +``` + +### 3. 跑评估 + +两条硬性约束，先看： + +- **必须传 `--checkpoint`**（绝对路径最省事）。编排进程不启动仿真，因此不加载 + hydra，无法按 `--load-run` 用实验名查找权重。 +- **每个种子一个子进程**。IsaacLab 在本配置下**一个进程只能创建一个环境**： + 第一个环境建好后，`ROUGH_TERRAINS_CFG` 这类模块级单例会被运行时长出自引用， + 之后任何 cfg 构造（新实例、开场存的 pristine 深拷贝、重新深拷贝）都会在 + `cfg.validate()` 里 RecursionError —— 三种补救方式都实测失败。所以 + `eval_suite.py` 的编排进程自己不建环境，而是为每个（阶段, 种子）fork 一个子进程。 + 代价是每个子进程要重新启动一次 Isaac Sim（约 15 s），所以 + `--protocol standard_v1_multiseed`（5 个种子）大约要 5~10 分钟。 + +```bash +# 标准评估（默认用 solenopsis 描述 + standard_v1 协议） +python scripts/reinforcement_learning/rsl_rl/eval_suite.py \ + --task Rough-Deeprobotics-M20-v0 \ + --checkpoint /绝对路径/model_19999.pt \ + --headless + +# 上限评估 + 5 个随机种子 + 存轨迹 +python scripts/reinforcement_learning/rsl_rl/eval_suite.py \ + --task Rough-Deeprobotics-M20-v0 \ + --checkpoint /绝对路径/model_19999.pt \ + --protocol standard_v1_multiseed --save-trace /tmp/traces --headless + +# 抗扰评估（每 2.01 s 施加一次水平扰动） +python scripts/reinforcement_learning/rsl_rl/eval_suite.py \ + --task Rough-Deeprobotics-M20-v0 \ + --checkpoint /绝对路径/model_19999.pt \ + --eval-type robustness --headless + +# 只看某一种地形的上限 +python scripts/reinforcement_learning/rsl_rl/eval_suite.py \ + --task Rough-Deeprobotics-M20-v0 \ + --checkpoint /绝对路径/model_19999.pt \ + --eval-type upper_limits --terrain-type stairs_up --headless + +# 推荐用于机器人对比：放大地砖，消除越界终止项的误判（见下一节） +python scripts/reinforcement_learning/rsl_rl/eval_suite.py \ + --task Rough-Deeprobotics-M20-v0 \ + --checkpoint /绝对路径/model_19999.pt \ + --protocol standard_v1_clean --headless + +# 或者直接关掉那个终止项 +python scripts/reinforcement_learning/rsl_rl/eval_suite.py \ + --task Rough-Deeprobotics-M20-v0 \ + --checkpoint /绝对路径/model_19999.pt \ + --no-terrain-bounds --headless + +# 冒烟测试（小规模，只为确认流水线通） +python scripts/reinforcement_learning/rsl_rl/eval_suite.py \ + --task Rough-Deeprobotics-M20-v0 \ + --checkpoint /绝对路径/model_19999.pt \ + --protocol smoke --headless +``` + +输出默认写到 `/eval_kit/_<时间戳>.json`——刻意与 +`eval.py` 的 `/eval/` 分开，避免两套结果混在一个目录里。 + +快速自检（不需要仿真）： + +```bash +python3 eval_kit/tests/test_core.py # 24 项：指标、聚合、名字解析、轨迹存取 +python3 eval_kit/tests/test_cli.py # 19 项：参数解析、口径覆盖、路径规则、推流开关 +``` + +### 4. 用 WebRTC 看画面（可视化） + +机器上没有显示器时，用 IsaacLab 的推流模式看评估过程：加 `--livestream` 即可 +（这个参数由 `AppLauncher` 提供，`eval_suite.py` 自己没有定义）。 + +```bash +# 局域网直连（推荐）：客户端和服务器在同一网段 +python scripts/reinforcement_learning/rsl_rl/eval_suite.py \ + --task Rough-Deeprobotics-M20-v0 \ + --checkpoint /绝对路径/model_18200.pt \ + --protocol smoke --terrain-type stairs_up \ + --livestream 2 --device cuda:6 +``` + +- `--livestream 2`：局域网/私有网络 WebRTC，客户端直连 `<本机 IP>:49100`。 +- `--livestream 1`：公网 WebRTC，走 NVIDIA NVCF 信令，需要 `PUBLIC_IP` 且能出公网。 +- 两个值都会让 `AppLauncher` 自动进入 headless（推流由服务端渲染）， + **不要再传 `--headless`**（传了也无害，只是多余）。 + +客户端用 NVIDIA 的 **Isaac Sim WebRTC Streaming Client**（随 Isaac Sim 一起提供， +Windows/Linux 都有），填入服务器地址即可；端口默认 49100，脚本启动时会把实际 +地址打印出来： + +``` +[eval_kit] WebRTC 推流已启用（--livestream 2，局域网直连） +[eval_kit] 客户端连接地址：10.62.25.137:49100 192.168.77.26:49100 +``` + +脚本为推流额外做三件事： + +1. **视口相机默认跟着被评估的机器人**。`--view-env N` 选第 N 个环境（默认 env 0）， + 上限评估时会打印这个编号对应哪种地形、哪个难度等级。不设的话推流画面停在 + 默认相机上，看不到机器人。 + 想自己控制镜头就加 **`--free-camera`**：改为不跟随任何物体，起始摆在俯视全场的 + 位置（按实际地形范围算中心与半径），之后在客户端里用鼠标/键盘自由移动 —— + Isaac Lab 的 `ViewportCameraController` 只在 `origin_type="asset_root"` 时每帧 + 改写相机，设成 `world` 就只在开头写一次，不会把你的操作抢回去。同时给了 + `--view-env` 时以 `--free-camera` 为准。 +2. **开始 rollout 前留出连接时间**：`--pause-s`（秒），推流时默认 30 s， + 期间**物理被冻结**、画面照常刷新，在终端按回车可提前开始。不推流时默认 0。 + 已验证：停顿与否跑出来的指标逐位相同（同一 seed）。 + 注意这一窗口里机器人是**悬空静止**的：它停在协议规定的起始位姿上，而 + `standard_v1_clean` 的 `start_z_offset_m=0.5` 会把它抬高 0.5 m 再落下， + 所以「离地没落下、也不动」正是正常现象；`sim.play()` 之后才开始动。 + 启动时会打印倒计时，别把它当成卡死。 +3. **打印客户端连接地址**（从 carb 设置里读实际端口，不是写死的）。 + +**想看清机器人走路**：`smoke`/`standard_v1` 一类协议一个 episode 只有 2~10 s 仿真时间， +扣掉 30 s 等待后留给你看的时间很短。只为观察的话把 episode 拉长（这会改变 `protocol` +指纹，**别拿这些数字做对比**）： + +```bash +# 只看画面：60 s 的 episode、60 个机器人，混合地形上一直走 +python scripts/reinforcement_learning/rsl_rl/eval_suite.py \ + --task Rough-Deeprobotics-M20-v0 --checkpoint /绝对路径/model_18200.pt \ + --protocol standard_v1_clean --eval-type velocity_error \ + --episode-length-s 60 --num-envs 60 \ + --livestream 2 --view-env 0 --pause-s 30 --device cuda:0 +``` + +**每个（阶段, 种子）是一个独立子进程**，各自开一次推流；子进程退出时画面中断， +下一个阶段会自动重新推流，客户端重连即可。所以 `--protocol standard_v1_multiseed` +这种 8 个种子的协议会看到 8 次「断线—重连」。想连续观察就用单种子的协议 +（`standard_v1` / `smoke`），或者按地形单独跑。 + +**这块机器上跑推流要先挑一张空闲卡**。渲染 + 编码要占显存，而渲染路径（MGPU/ +流编码器）会在 GPU 0 上申请资源，GPU 0 正被别人的训练占满时会刷大量 +`Out of GPU memory` / `Buffer creation failed for the device: 0`，严重时 +`CUDA error: an illegal memory access` 直接把评估打断。处置： + +```bash +nvidia-smi --query-gpu=index,memory.used --format=csv # 挑一张空卡 +--device cuda:N # 渲染与物理都钉到 N +--rendering_mode performance # 可选：低画质换显存与延迟 +--kit_args=--/renderer/multiGpu/maxGpuCount=1 # 可选：限制渲染只用一张卡 +``` + +`--device` 现在同时决定**策略网络**放在哪张卡：以前策略只认 agent 配置里的 +`device`（默认 `cuda:0`），在别人占着 0 号卡时会莫名其妙 OOM。 + +**读数提醒（未解）**：同一 checkpoint、同一协议、同一 seed，**启动方式不同会得到不同的 +数值**。实测（`--protocol smoke --eval-type velocity_error --num-envs 64`，同一个 +`model_18200.pt`）： + +| 启动方式 | mean_lin_vel_error | mean_yaw_vel_error | 中途早退 | +| --- | --- | --- | --- | +| `--headless --device cuda:6`（三次独立重跑） | 0.8159 | 0.1480 | 10 | +| `CUDA_VISIBLE_DEVICES=6 --headless --device cuda:0` | 0.8507 | 0.3334 | 17 | +| `--headless --enable_cameras`（渲染开，无窗口） | 0.8507 | 0.3334 | 17 | +| `--livestream 2 --device cuda:6`（推流，暂停 0~30 s 都试过） | 0.8507 | 0.3334 | 17 | + +每一组内部都是**逐位可复现**的（同组三次重跑完全一致），跨组则有稳定差异；`--pause-s` +已验证与差异无关（headless 加暂停与不加完全相同；推流下暂停 0 s 与 30 s 也相同）。 +两档之间的成因尚未定位（与「是否渲染」和「`--device` 写法/可见卡数」都相关，未能归因到单一因素）。 +**在这个问题查清之前：拿数字做对比时固定一种启动方式**，不要把可视化那次的数值与历史 +headless 结果混在一个表里报。看画面用推流没有影响，只是它的数字属于另一档。 + +**画面卡住时先看终端，不要猜**。rollout 期间每 10% 打印一行进度： + +``` +[eval_kit] 停顿 5.0 s 结束，开始 rollout（画面从这里才动） +[eval_kit] rollout 100/1000 步，用时 7.0 s +[eval_kit] rollout 200/1000 步，用时 13.6 s +... +[eval_kit] 结果已保存：.../upper_limits_2026-10-19_17-16-19.json +``` + +- **有 `rollout` 打印在涨** → 仿真在正常推进，卡的是画面（渲染/编码侧）。按优先级试： + 换到空闲卡（`--device cuda:N`，N 用空闲卡的编号）、 + `--rendering_mode performance`、确认客户端没断线、机器人数量少一点（`--num-envs`）。 + 已知一种具体成因已经默认修掉：客户端改窗口大小/全屏会发 `onCustomResizeMessage`， + 服务端默认拒绝并让推流的 `shared swapchain buffer` 进入坏状态，此后画面不再更新 + （仿真还在跑）。脚本现在会给推流自动追加 + `--/app/livestream/allowResize=true --/app/livestream/allowDynamicResize=true`。 + 如果日志里仍出现 `onCustomResizeMessage ... must both be positive`，**别缩放客户端窗口**， + 并把那段日志发出来。 +- **`rollout` 打印不再涨** → 是仿真停了，通常是在客户端 UI 里动了时间轴/仿真控件： + 视口下方的时间轴 Play/Pause/Stop 按钮、IsaacLab 面板里的 `Record Animation` + （START 之后默认 10 s 结束会**直接关掉应用**）、`Reset Simulation on Stop`。 + 看画面时别点它们；真被误触了就 Ctrl-C 重跑。 +- **进程已经结束**（终端出现 `结果已保存`，或被 Ctrl-C）→ 客户端只会把最后一帧留在屏幕上， + 看起来就是「画面定住」，这既不是卡住也没有在算，重跑即可。 +- 一轮跑完的标志只有两个：终端出现 `结果已保存：<路径>.json`，以及 + `/eval_kit/` 下多出一个 JSON。 + +### 已知问题：WebRTC 画面在客户端接上后停止更新 + +实测到过这个现象：仿真、渲染、相机都正常（`--record-frames` 抓出的帧显示机器人在动、 +相机跟着走），**只有推流给客户端的画面不再刷新**。已排除的原因：物理冻结（终端 +`rollout` 进度一直在涨）、暂停窗口（停顿结束后仍不动）、客户端改画幅被拒（已默认打开 +`allowResize`/`allowDynamicResize`）、时间轴被暂停（那样进度会停）。 + +推流器默认把日志静音（`app.livestream.logLevel = "silence"`），所以要先让它说话： + +```bash +--kit_args=--/app/livestream/logLevel=debug +# 另外可以试：--kit_args=--/renderer/multiGpu/maxGpuCount=1 +# （IsaacLab 自己的 isaaclab.python.rendering.kit 就是这么设的，多卡机器上更保险） +``` + +再不行就用 `--record-frames` 那条路径看评估过程（功能上等价，只是不能实时移动相机）。 + +**推流不可用时的替代路径：抓帧成视频。** `--record-frames DIR` 会在 rollout 期间按 +间隔把视口画面存成 PNG，跑完打印合成命令： + +```bash +CUDA_VISIBLE_DEVICES=7 python scripts/reinforcement_learning/rsl_rl/eval_suite.py \ + --task Rough-Deeprobotics-M20-v0 --checkpoint /绝对路径/model_18200.pt \ + --protocol standard_v1_clean --eval-type velocity_error \ + --episode-length-s 30 --num-envs 16 \ + --livestream 2 --pause-s 0 --view-env 0 --device cuda:0 \ + --record-frames /tmp/eval_frames --record-every 5 +``` + +抓帧依赖视口渲染，所以**必须开着渲染**：推流（`--livestream 2`）或本地 GUI 都行， +纯 headless（`--headless` 且不推流）没有视口，抓不到帧。跑完合成 +（`--record-every 5` 时每帧 0.1 s 仿真时间，10 fps 即按仿真速度播放）： + +```bash +cd /tmp/eval_frames && ffmpeg -y -framerate 10 -pattern_type glob -i '*.png' \ + -pix_fmt yuv420p -c:v libx264 -vf 'pad=ceil(iw/2)*2:ceil(ih/2)*2' /tmp/eval.mp4 +``` + +已实测：600 步抓到 102 张 1280×720 PNG，合出的 mp4 能正常播放，画面里机器人从 +起始高度落地并在地形上移动、相机跟着它走 —— 这条路径绕开了推流传输层， +在推流卡住时仍然能看评估过程。 + +## 三个对象 + +- **`RobotProfile`**（`profiles.py`）：这个机器人谁是足端、谁是支撑体、哪些关节耗电、多重。 + 评价代码里没有任何 link/joint 名字或质量常数。 +- **`EvalProtocol`**（`protocols.py`）：一次评估的全部口径——指令、时长、成功距离、 + 倾角阈值、加速度阈值、能耗约定、地形规格、种子。决定数值的字段会被哈希成 + `protocol_fingerprint` 写进输出 JSON。 +- **`Trace`**（`core/schema.py`）：一次 rollout 的逐帧记录（机身位姿/速度、关节力矩与速度、 + 足端位置、支撑接触力与时长）。第 0 帧是动作前的起始状态，之后每步一帧。 + +流水线是单向的：适配器产 `Trace` → 内核算指标 → 报告层格式化。 +`Trace` 可以存成 `.npz`，事后加新指标不用重跑仿真。 + +## 加一个机器人 + +在 `profiles.py` 加一条记录即可，评估代码不动： + +```python +MY_ROBOT = RobotProfile( + name="my_robot", + mass_kg=None, # None = 运行时从资产读真实质量 + feet_body_patterns=(".*_foot",), + support_body_patterns=(), # 没有支撑体语义 → 相关指标返回 nan + actuated_joint_patterns=(".*_joint",), + support_force_threshold_n=15.0, + body_length_m=0.5, +) +``` + +`test_core.py::test_naming_does_not_crash_when_morphology_is_missing` 就是在验证这件事： +用火蚁的描述去解析一个没有轮子的机器人，得到的是「未命中」记录 + 跳过对应指标，而不是崩溃。 + +## 相对 `eval.py` 的三处有意差异 + +指标名、取点、公式全部保持一致，只有三处刻意不同： + +| 项 | `eval.py` | 本套件 | 原因 | +|---|---|---|---| +| COT 的质量 | `--robot_mass` 默认 25.482 | 默认从 `root_physx_view.get_masses()` 读（`Solenopsis.urdf` 实为 26.581 kg），命令行可覆盖 | 旧默认值是更早的 M20 简化模型名义质量，会让 COT 偏低约 4% | +| 加速度阈值 | `--acc_threshold` 2.0 m/s²（绝对值） | `acc_threshold_g` 重力倍数（默认仍是 2.0/9.81） | 绝对值对不同量级的机器人不公平 | +| 缺失形态的指标 | 返回 0.0 | 返回 `nan`，并在 JSON 里记录 `stability_margin_method="unavailable"` | 0.0 看起来像测出来的值，实际是没测 | + +另外新增了两项检查：输出 JSON 里带 `protocol_fingerprint`（防止不同口径的结果被悄悄比较）， +以及 `early_resets`（环境在 episode 中途被重置时会警告，因为那会让路程/进度类指标混入两次 rollout）。 + +## 口径与 `criterion.py` 的对照（备忘） + +`criterion.py`（来自 legged_gym 的山猫项目）与 `eval.py` 在 7 处口径上不一致， +移植/对照时要注意。已经核对过的结论： + +| 项 | `criterion.py` | `eval.py` / 本套件 | +|---|---|---| +| 能耗积分 | `Σ(τ·q̇)·dt`（有符号） | `Σ\|τ·q̇\|·dt`（绝对值，本套件可切换但默认绝对值） | +| COT 质量 | 写死 27.78822 | 自动读取（26.581 kg） | +| 倾角阈值 | π/6 ≈ 0.524 rad | 0.70 rad | +| 稳定裕度 | ZMP 到足端凸包距离（scipy + shapely） | 机身到足端包围盒距离 | +| 成功判定 | env 维护目标区域 | 沿指令方向前进 ≥ success_distance | +| 时间效率基准 | 走过半块地形的时间 | 走完 success_distance 的时间 | +| 上限 score | 累积成功率加权 | **代数等价，已验算一致** ✓ | + +其中「能耗积分」这一条最要紧：有符号口径下负功会抵消正功，同一段运动算出的能耗可以 +差一倍以上，所以跨机器人比较时必须在协议里钉死一种。 + +## ⚠️ 发现的问题：越界终止项让约 20% 的环境从第 0 步就被重置 + +这是实跑冒烟测试时发现的，**是 `eval.py` 评估设置里既有的问题，不是本套件引入的**。 + +IsaacLab 的 `terrain_out_of_bounds` 终止项（`isaaclab_tasks/.../mdp/terminations.py`） +比较的是机器人的**世界坐标**与**整张地图**的范围，而不是机器人自己那块地砖： + +```python +map_width = n_rows * tile + 2 * border_width +x_out_of_bounds = |root_pos_w[:, 0]| > 0.5 * map_width - distance_buffer +``` + +在 `eval.py` 的默认评估配置下（地砖 2.0 m、`border_width` 2.0、`distance_buffer` 3.0、 +起始位姿 x = −0.9），网格最外侧格子中心在 ±(rows−1)/2 × 2.0 处，机器人起始位置因此 +越过阈值。实测数据（8 难度 × 2 地形 × 每格 2 个 = 32 个环境）： + +- 9 行格子中有 2 个整格越界（origin x = −7.0 → 起始 x = −7.9 > 阈值 7.0） +- **首步就有 6/32 = 19% 的环境被判越界并重置** +- 5×5 网格下同理：最左一整列（20%）从第 0 步起每步都被重置 + +还有第二个机制：`|y|` 恰好**等于**阈值的那些列（y = ±4.0 对阈值 4.0）处于浮点边界上， +取决于最后一个 bit 是否触发。 + +两个后果，都影响可比性： + +1. 这些环境永远没有机会尝试地形，却全部被记成失败 → 成功率和 score 系统性偏低。 +2. 地形行序是**打乱**的（哪一行落在最外侧每次都不同）→ **被清零的是哪个难度等级 + 跨种子都不稳定**，所以逐等级成功率和 score 连自己跟自己比都不可靠。 + +### 三种处置（按推荐顺序） + +| 方案 | 做法 | 效果 | +|---|---|---| +| 1 | `--terrain-size 4.0` 或 `--protocol standard_v1_clean` | 保留终止项，消除误判。阈值由 `minimum_terrain_size_for_bounds()` 给出：`tile >= 2*(0.9 + 3.0 − 2.0) = 3.8 m` | +| 2 | `--no-terrain-bounds` 或 `--protocol standard_v1_no_bounds` | 直接去掉该终止项。成功判定本就按前进距离算，走出地砖不构成失败 | +| 3 | 保持默认（2.0 m 地砖 + 启用终止项） | 与旧 `eval.py` 数值可比，但带着上述偏差；只建议用于对照 | + +### 因果验证 + +同一配置（5×5 网格、10 个环境、8 步、零动作策略），只切换该终止项： + +| `enforce_terrain_bounds` | 各环境中途重置次数 | 结论 | +|---|---|---| +| `True`（默认，与 eval.py 一致） | `[7, 1, 0, 0, 0, 0, 0, 7, 7, 0]` | 4/10 环境受污染 | +| `False` | `[0, 0, 0, 0, 0, 0, 0, 0, 0, 0]` | 无早退 | + +本套件的 `early_resets` 检查就是为发现这类问题加的，输出 JSON 里一直带着它。 + +## ⚠️ 发现的问题二、三：地形参数与生成位置 + +下面两处同样是实跑时发现的，也都继承自 ``eval.py`` 的评估设置。 + +### 问题二：机器人在倒金字塔台阶地形里被生成到几何体内部 + +旧写法把重置高度钉在瓦片中心的 ``env_origin.z`` 上（``pose_range.z = (0, 0)``）， +这只对「瓦片中心是最高点」的地形成立。 + +``stairs_down`` 用的 ``MeshInvertedPyramidStairsTerrainCfg`` **中心是最低点**， +再叠加 ``platform_width`` 被从 3.0 m 压到 0.4 m（``min(0.4, tile*0.2)``）， +起始点 x = −0.9 m 就落在平台之外且低于该处地表 —— 机器人被生成在几何体内部、卡住。 + +结果是 ``stairs_down`` **在全部 8 个难度上都是 0.00**，连只有 4 cm 台阶的 L1 也是， +看上去像「这机器人不会下台阶」，实际是 harness 的问题。 + +**修复**：``EvalProtocol.start_z_offset_m``（CLI ``--start-z-offset``）。 +实测对照（同一 checkpoint、只跑 ``--terrain-type stairs_down``）： + +| 起始高度 | score | 成功率 | +|---|---|---| +| 0.0（旧 eval.py） | 0.00 / 8 | 0.05 | +| 0.3 m | 3.64 / 8 | 0.47 | +| **0.5 m** | **4.53 / 8** | — | + +取 0.5 m 是因为嵌入深度随难度递增（L1 约 0.11 m、L8 约 0.50 m），0.5 刚好覆盖最深一级。 +``standard_v1_clean`` 已内置该值；``standard_v1`` 保持 0.0 以维持与旧 eval.py 的可比性。 +代价是机器人在平台上方落下会弹一下，``stairs_up`` 略降，但整体是净收益。 + +### 问题三：粗糙地形的噪声参数被当成起伏幅度 + +``rough_slope`` 用的 ``HfRandomUniformTerrainCfg`` 有两个容易误解的点，IsaacLab 源码里 +写得很清楚但很容易忽略： + +1. **``difficulty`` 参数被忽略**（docstring 原文：*"The difficulty parameter is ignored + for this terrain"*）。所以那一列**没有难度梯度**，8 行是同一分布的重复采样 —— + 这也正是它能被看出来的原因：它的成功率不单调（0.12/0.00/0.38/0.25/0.62/0.75/0.12/0.00）。 +2. **``noise_range`` 是绝对高度范围，不是起伏幅度**。旧 eval.py 设成 ``(0.05, 0.40)``， + 而 ``noise_step`` 又被改成 0.01 —— 在 1 cm 水平网格上放 40 cm 高的随机起伏， + 等于一堵堵随机尖刺墙，轮半径只有 0.096 m 的机器人无法通过。 + +**修复**：``noise_step`` 回到上游的 0.02，``noise_range`` 的上限改为 0.10（与训练地形 +``noise_range=(0.02, 0.10)`` 同量级）。实测： + +| ``noise_range`` | score | 成功率 | +|---|---|---| +| (0.05, 0.40)（旧 eval.py）+ noise_step 0.01 | 0.12 / 8 | 0.05 | +| (0.02, 0.10) + noise_step 0.02 | **6.38 / 8** | 0.86 | + +同时给规格加了 ``isaaclab_extra_params``（写与等级无关的固定字段），并把裸 ``setattr`` +换成带 ``hasattr`` 校验的 ``_set_sub_terrain_field``：字段名写错现在会直接报错， +不会再静默造一个假属性、让地形保持默认难度（那看起来只是「很难」）。 + +``rough_slope`` 的 ``unit`` 里也写明了「此地形忽略难度参数，等级排名无意义」， +所以这条限制会随输出 JSON 一起保存。 + +### 三处修正的合计效果 + +同一 checkpoint（``model_19999.pt``、seed 0），只切换协议预设： + +| 地形 | 旧 eval.py 口径 | ``standard_v1_clean`` | +|---|---|---| +| stairs_up | 2.88 | 3.03 | +| stairs_down | 0.00 | **4.21** | +| smooth_slope | 6.50 | 6.50 | +| rough_slope | 0.12 | **6.03** | +| boxes | 4.28 | 4.60 | +| **总成功率** | 0.434 | **0.659** | +| **平均前进距离** | 3.59 m | **5.26 m** | +| **COT** | 4.50 | **2.45** | +| **score 总和** | 13.78 | **24.38** | + +## ⚠️ 尚存的未解异常：所有地形的 L8 都是 0.00 + +修正后在完整评估里观察到：**五个地形在最高难度 L8 上的成功率全部是 0.00**。 + +对 ``stairs_up`` / ``boxes`` 这类台阶高度确实随难度递增的地形，L8 失败是合理的； +但 ``rough_slope`` 忽略难度参数，L8 的统计分布与 L1–L7 完全相同（那 7 行是 0.75–1.00）， +却精确地得到 0.00 —— 这在统计上不可能是能力问题，指向**最后一行的地形原点或生成高度有问题** +（最可能是 ``terrain_origins`` 最后一行的 z 或布局异常，使机器人被放在地形之外/之下）。 + +影响：每个地形的 ``score`` 末项 ``acc[-1] * num_levels`` 因 ``acc[-1] = 0`` 而丢失， +score 被低估约 1 分（以 smooth_slope 为例，6.50 → 若 L8 正常约 7.40）。 + +**这一条尚未定位，也没有修。** 在有明确结论之前，请把矩阵最后一行的数据当作不可用， +score 当作保守下界。 + +## 验证状态（请如实看待） + +**已在真机（`IsaacLab-2.3.2/env_isaaclab_py311` + RTX 4090）上实跑验证**： + +- `eval_kit/isaaclab/` 三个模块已端到端跑通，三步冒烟全部 PASS： + - `names` —— 形态解析在真实 body 名上命中：足端 4 个 + `['fl_wheel','fr_wheel','hl_wheel','hr_wheel']`、支撑体 5 个 + `['base_link','fl_knee',...,'hr_knee']`、驱动关节 8 个，全部与描述文件一致 + - `plain` —— 配置 → 采集 → 指标全链路：11 帧 × 4 环境，指标齐出 + （`margin=0.1667 m`、`max_tilt=0.051 rad`、`energy=8.0572 J`、`COT=2.1570`） + - `upper` —— 上限网格链路：8 个环境按 2 类型 × 2 等级分格，矩阵与 score 正确折叠 +- **整机质量从资产读出为 26.581 kg**，与 `docs/solenopsis_urdf_comparison.md` 的记录一致； + 这确认了 `eval.py` 的 `--robot_mass` 默认值 25.482 确实偏低约 4%，本套件改为自动读取 + 是有必要的。 +- 越界终止项的偏差已定位、量化并做过因果验证（见上一节）。 +- `eval_kit/core/` 的指标用合成轨迹验证，期望值手算：22 项测试通过。 +- `eval_suite.py` 的参数解析与口径覆盖用桩替换 `isaaclab.app` 后测试：11 项通过。 + +**仍未验证**： + +- 没有用**训练到收敛的** checkpoint 跑过，所以指标数值本身没有和 `eval.py` 做过逐项对照。 + 现有的 `model_0.pt` 是 iteration 0（基本等于随机权重），跑出来的 `success_rate` 恒为 0、 + `time_consumption` 恒为 nan，拿这种退化值对照没有信息量——**这时"通过"是弱证据， + "报错"才是强证据**，而报错部分已经全部排除。 +- 因此**第一次拿到训练好的 checkpoint 时仍建议做等价性验证**：同一 checkpoint、 + 同一组参数分别跑 `eval.py` 与 `eval_suite.py --protocol standard_v1`，除「相对 eval.py + 的三处有意差异」一节列出的三项外应逐位相同。注意 `eval.py` 目前**无法启动**（见下）。 +- `profiles.py` 里的 `lite3` 记录没有在仿真里跑过。 + +## ⚠️ `eval.py` 目前无法启动 + +实跑时顺带发现的：`eval.py` 自己也加了一个 `--verbose`，而 +`AppLauncher.add_app_launcher_args()` 会无条件添加同名参数，argparse 因此在**构造 +parser 阶段**就抛错，脚本连 `--help` 都跑不出来： + +``` +argparse.ArgumentError: argument --verbose: conflicting option string: --verbose +``` + +`--device` 和 `--verbose` 是两次踩到的坑，但只有 `device / headless / enable_cameras / +experience / livestream / rendering_mode / xr` 这 7 个会被 AppLauncher 自己的显式检查 +拦下并给出可读报错；`--verbose / --info / --cpu` 等其余 7 个由 argparse 拦。 +本套件把 14 个名字都记在 `tests/test_cli.py::APPLAUNCHER_ARG_NAMES` 里，并在 +`test_parser_compatible_with_app_launcher` 中复刻检查，所以这类冲突不必启动仿真就能发现。 + +修 `eval.py` 只需把那行 `--verbose` 改名（例如 `--verbose-terrain`）并同步 +`args_cli.verbose` 的用法。**我没有改动 `eval.py`**——按你的要求新写一套、不碰原文件。 + +## 已知限制 + +- **一个进程只能建一个环境**（IsaacLab 在本配置下的限制，实测确认）。因此 + `eval_suite.py` 用「编排进程 + 每（阶段, 种子）一个子进程」的结构，每个子进程 + 重新启动一次 Isaac Sim。这不是本套件的设计偏好，是绕不过去的约束。 +- **地形规格是唯一真正不可移植的部分。** 规格数值（等级 × 类型 × 绝对高度）是纯数据， + 但「把规格翻译成地形生成器」的代码是 IsaacLab 专用（`isaaclab/env_setup.py`）。 + 换框架需要另写这个 builder。 +- **稳定裕度口径较粗。** 当前是「机身到足端轴对齐包围盒的距离」，比 ZMP/凸包版本粗糙， + 但不会在退化多边形上失败、也不需要 scipy。算法名记录在输出 JSON 里。 +- **重复试验只覆盖随机种子。** `aggregate_runs` 给的是种子间的 mean±std， + 不含 checkpoint 之间的方差（同一实验不同迭代的权重差异）——那个要靠多取几个 checkpoint 评估。 +- **单次 rollout 的噪声很大。** 比较两个策略/机器人时请用 `--seeds` 给多个种子， + 看 mean±std，不要看单次数值。 \ No newline at end of file diff --git a/eval_kit/__init__.py b/eval_kit/__init__.py new file mode 100644 index 0000000..06fc066 --- /dev/null +++ b/eval_kit/__init__.py @@ -0,0 +1,35 @@ +# Copyright (c) 2025 Deep Robotics +# SPDX-License-Identifier: BSD 3-Clause + +"""可搬运的机器人性能评价套件。 + +分层： + +- :mod:`eval_kit.core` —— 零框架依赖（只用 numpy + 标准库），可整体拷到别的项目 +- :mod:`eval_kit.isaaclab` —— IsaacLab 适配层，唯一接触仿真环境内部的地方 +- :mod:`eval_kit.profiles` —— 机器人形态描述（加机器人 = 加一条记录） +- :mod:`eval_kit.protocols` —— 评估协议预设 + +命令行入口：``scripts/reinforcement_learning/rsl_rl/eval_suite.py``。 +""" + +from eval_kit.core.aggregation import aggregate_runs, build_upper_limit_matrix +from eval_kit.core.metrics import compute_metrics +from eval_kit.core.schema import EvalProtocol, RobotProfile, Trace +from eval_kit.core.terrain_spec import get_terrain_spec +from eval_kit.profiles import PROFILES, get_profile +from eval_kit.protocols import PROTOCOLS, get_protocol + +__all__ = [ + "EvalProtocol", + "PROFILES", + "PROTOCOLS", + "RobotProfile", + "Trace", + "aggregate_runs", + "build_upper_limit_matrix", + "compute_metrics", + "get_profile", + "get_protocol", + "get_terrain_spec", +] diff --git a/eval_kit/core/__init__.py b/eval_kit/core/__init__.py new file mode 100644 index 0000000..709f27f --- /dev/null +++ b/eval_kit/core/__init__.py @@ -0,0 +1,26 @@ +# Copyright (c) 2025 Deep Robotics +# SPDX-License-Identifier: BSD 3-Clause + +"""评价内核：不依赖 isaaclab / isaacsim / torch / GPU，只依赖 numpy 与标准库。 + +这里的模块可以在任何机器上 import 和测试，包括没有仿真环境的机器。 +""" + +from eval_kit.core.aggregation import aggregate_runs, build_upper_limit_matrix +from eval_kit.core.metrics import compute_metrics +from eval_kit.core.naming import NameResolution, resolve_matching_indices +from eval_kit.core.schema import EvalProtocol, RobotProfile, Trace +from eval_kit.core.terrain_spec import TERRAIN_SPECS, get_terrain_spec + +__all__ = [ + "TERRAIN_SPECS", + "EvalProtocol", + "NameResolution", + "RobotProfile", + "Trace", + "aggregate_runs", + "build_upper_limit_matrix", + "compute_metrics", + "get_terrain_spec", + "resolve_matching_indices", +] diff --git a/eval_kit/core/aggregation.py b/eval_kit/core/aggregation.py new file mode 100644 index 0000000..84f1e37 --- /dev/null +++ b/eval_kit/core/aggregation.py @@ -0,0 +1,212 @@ +# Copyright (c) 2025 Deep Robotics +# SPDX-License-Identifier: BSD 3-Clause + +"""上限评估的网格汇总：把逐环境结果折叠成「难度等级 × 地形类型」的成功率矩阵。 + +纯数据变换，只依赖 numpy，与仿真框架无关。 + +``score`` 的公式与 ``criterion.py`` / 旧 ``eval.py`` 完全一致（三者代数等价）： +对每个地形类型，先算逐级累积成功率 ``acc[i] = Π_{j<=i} rate[j]``， +再算 ``score = Σ_i acc[i] * (i+1) * (1 - rate[i+1])``，最后一级的 ``(1 - rate)`` 取 1。 +含义是「在爬到第 i+1 级之前失败」的累积期望，值越大说明越能爬到高难度。 +""" + +from __future__ import annotations + +from typing import Any, Sequence + +import numpy as np + +from eval_kit.core.terrain_spec import TerrainSpec + + +def accumulated_success(rates: Sequence[float]) -> list[float]: + """逐级累积成功率。""" + accumulated: list[float] = [] + running = 1.0 + for rate in rates: + running *= float(rate) + accumulated.append(running) + return accumulated + + +def upper_limit_score(rates: Sequence[float]) -> float: + """由某一地形类型各难度等级的成功率算 score。""" + rates = [float(rate) for rate in rates] + accumulated = accumulated_success(rates) + score = 0.0 + for index, acc_rate in enumerate(accumulated): + next_failure = 1.0 - rates[index + 1] if index + 1 < len(rates) else 1.0 + score += acc_rate * (index + 1) * next_failure + return score + + +def aggregate_runs(runs: Sequence[tuple[int, dict[str, Any]]]) -> dict[str, Any]: + """把多次（不同随机种子）的指标聚合成 mean/std。 + + 单次 rollout 的噪声足以淹没两个策略之间的真实差异，所以做对比时应该跑多个 + 种子并看分布，而不是看单个数。 + + Args: + runs: ``(seed, metrics)`` 列表，``metrics`` 来自 + :func:`eval_kit.core.metrics.compute_metrics`。 + """ + if not runs: + raise ValueError("没有任何一次 run 可以聚合。") + + # 逐环境的私有字段不参与聚合 + metric_keys: list[str] = [] + for _, metrics in runs: + for key, value in metrics.items(): + if key.startswith("_"): + continue + if isinstance(value, (int, float)) and key not in metric_keys: + metric_keys.append(key) + + summary: dict[str, Any] = {} + for key in metric_keys: + values = np.array([float(metrics.get(key, float("nan"))) for _, metrics in runs], dtype=np.float64) + finite = values[np.isfinite(values)] + if finite.size == 0: + summary[key] = { + "mean": float("nan"), + "std": float("nan"), + "min": float("nan"), + "max": float("nan"), + "values": values.tolist(), + "finite_runs": 0, + } + continue + summary[key] = { + "mean": float(np.mean(finite)), + "std": float(np.std(finite, ddof=1)) if finite.size > 1 else 0.0, + "min": float(np.min(finite)), + "max": float(np.max(finite)), + "values": values.tolist(), + "finite_runs": int(finite.size), + } + + consistent: dict[str, Any] = {} + for key in runs[0][1]: + if not key.startswith("_") and isinstance(runs[0][1][key], str): + values = {metrics.get(key) for _, metrics in runs} + consistent[key] = runs[0][1][key] + if len(values) > 1: + consistent[f"{key}_INCONSISTENT"] = sorted(str(value) for value in values) + + return { + "num_runs": len(runs), + "seeds": [int(seed) for seed, _ in runs], + "summary": summary, + "constant": consistent, + "per_run": [{"seed": int(seed), "metrics": metrics} for seed, metrics in runs], + } + + +def build_upper_limit_matrix( + per_env: dict[str, list[float]], + level_ids: np.ndarray, + type_ids: np.ndarray, + terrain_types: Sequence[str], + spec: TerrainSpec, + num_levels: int, +) -> dict[str, Any]: + """把逐环境结果汇总成网格。 + + Args: + per_env: 含 ``success`` / ``max_progress`` / ``max_tilt`` 的逐环境列表， + 由 :func:`eval_kit.core.metrics.compute_metrics` 在 + ``include_per_env=True`` 时给出。 + level_ids: ``(N,)`` 每个环境所属的难度等级（0 起）。 + type_ids: ``(N,)`` 每个环境所属的地形类型序号（对应 ``terrain_types``）。 + terrain_types: 地形类型名，同时决定矩阵列顺序。 + spec: 地形规格，用于取每级的难度值。 + num_levels: 难度等级数（矩阵行数）。 + """ + success = np.asarray(per_env["success"], dtype=np.float64) + max_progress = np.asarray(per_env["max_progress"], dtype=np.float64) + max_tilt = np.asarray(per_env["max_tilt"], dtype=np.float64) + level_ids = np.asarray(level_ids, dtype=np.int64) + type_ids = np.asarray(type_ids, dtype=np.int64) + + if not (success.shape == max_progress.shape == max_tilt.shape == level_ids.shape == type_ids.shape): + raise ValueError("逐环境数组与网格编号数组长度不一致。") + + per_type: dict[str, Any] = {} + for type_index, terrain_type in enumerate(terrain_types): + level_values = spec.levels_for(terrain_type, num_levels) + rates: list[float] = [] + level_metrics: list[dict[str, Any]] = [] + for level_index, level_value in enumerate(level_values): + mask = (type_ids == type_index) & (level_ids == level_index) + total = int(np.count_nonzero(mask)) + if total == 0: + rates.append(float("nan")) + level_metrics.append( + { + "level": level_index + 1, + "difficulty": float(level_value), + "success_rate": float("nan"), + "success_count": 0.0, + "num_envs": 0, + "avg_progress_m": float("nan"), + "avg_max_tilt_rad": float("nan"), + } + ) + continue + count = float(np.count_nonzero(success[mask] > 0.5)) + rates.append(count / total) + level_metrics.append( + { + "level": level_index + 1, + "difficulty": float(level_value), + "success_rate": count / total, + "success_count": count, + "num_envs": total, + "avg_progress_m": float(np.mean(max_progress[mask])), + "avg_max_tilt_rad": float(np.mean(max_tilt[mask])), + } + ) + + per_type[terrain_type] = { + "levels": level_metrics, + "success_rate": rates, + "success_rate_accumulated": accumulated_success(rates), + "score": upper_limit_score(rates), + "max_level": len(level_values), + "difficulty_unit": spec.unit.get(terrain_type, ""), + } + + # 打成便于打印与比较的矩阵：行 = 难度等级，列 = 地形类型 + success_matrix: list[list[float]] = [] + accumulated_matrix: list[list[float]] = [] + difficulty_matrix: list[list[float]] = [] + for level_index in range(num_levels): + success_row: list[float] = [] + accumulated_row: list[float] = [] + difficulty_row: list[float] = [] + for terrain_type in terrain_types: + entry = per_type[terrain_type] + rates = entry["success_rate"] + accumulated = entry["success_rate_accumulated"] + levels = entry["levels"] + success_row.append(rates[level_index] if level_index < len(rates) else float("nan")) + accumulated_row.append(accumulated[level_index] if level_index < len(accumulated) else float("nan")) + difficulty_row.append(levels[level_index]["difficulty"] if level_index < len(levels) else float("nan")) + success_matrix.append(success_row) + accumulated_matrix.append(accumulated_row) + difficulty_matrix.append(difficulty_row) + + return { + "benchmark": spec.name, + "terrain_types": list(terrain_types), + "num_levels": num_levels, + "difficulty_unit": {terrain_type: spec.unit.get(terrain_type, "") for terrain_type in terrain_types}, + "per_type": per_type, + "scores": {terrain_type: per_type[terrain_type]["score"] for terrain_type in terrain_types}, + "matrix": { + "success_rate": success_matrix, + "success_rate_accumulated": accumulated_matrix, + "difficulty": difficulty_matrix, + }, + } diff --git a/eval_kit/core/metrics.py b/eval_kit/core/metrics.py new file mode 100644 index 0000000..695fa75 --- /dev/null +++ b/eval_kit/core/metrics.py @@ -0,0 +1,262 @@ +# Copyright (c) 2025 Deep Robotics +# SPDX-License-Identifier: BSD 3-Clause + +"""指标计算：纯函数，只依赖 numpy。 + +输入是 :class:`~eval_kit.core.schema.Trace`、:class:`~eval_kit.core.schema.RobotProfile` +和 :class:`~eval_kit.core.schema.EvalProtocol`，输出是一个扁平的指标字典。 + +指标口径与 ``scripts/reinforcement_learning/rsl_rl/eval.py`` 保持一致（同样的取点、 +同样的公式），因此默认协议下两套评估的数值可以直接对照；差别只在三处，且都是 +有意为之： + +1. **质量**不再写死在命令行默认值里，而是从 ``Trace.mass_kg`` 取（适配层从资产读）， + 传入的 ``profile.mass_kg`` 只作兜底。 +2. **加速度阈值**以重力倍数表达，换算成 m/s² 后参与计算，使阈值可跨量级机器人复用。 +3. **缺失的形态**返回 ``nan`` 而不是 0。旧实现里没有足端的机器人会得到 + 「支撑裕度 = 0.00」，看起来像测出来的值，实际是没测。 +""" + +from __future__ import annotations + +from typing import Any + +import numpy as np + +from eval_kit.core.schema import EvalProtocol, RobotProfile, Trace + +#: 全部帧，用于成功判定与最大进度（含起始帧，与旧实现取点一致）。 +_ALL_FRAMES = slice(None) +#: 动作执行之后的帧，用于逐帧累积类指标。 +_STEP_FRAMES = slice(1, None) + +_GRAVITY = 9.81 + + +def tilt_from_projected_gravity(projected_gravity_b: np.ndarray) -> np.ndarray: + """由机体系投影重力算倾角 [rad]。 + + ``projected_gravity_b`` 为 ``(T, N, 3)``，返回 ``(T, N)``。 + """ + z_axis = np.clip(-projected_gravity_b[..., 2], -1.0, 1.0) + return np.arccos(z_axis) + + +def support_margin_xy(root_xy: np.ndarray, foot_xy: np.ndarray) -> np.ndarray: + """机身投影点到足端轴对齐包围盒的距离 [m]，在盒内为正、盒外为负。 + + ``root_xy`` 为 ``(N, 2)``，``foot_xy`` 为 ``(N, F, 2)``，返回 ``(N,)``。 + + 这是当前实现采用的口径（与旧 ``eval.py`` 相同），记录在 + ``EvalProtocol.stability_margin_method`` 里。它比凸包/ZMP 版本粗糙， + 但不会在退化多边形上失败，也不需要 scipy。 + """ + min_xy = np.min(foot_xy, axis=1) + max_xy = np.max(foot_xy, axis=1) + + 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 = np.minimum( + np.minimum(root_xy[:, 0] - min_xy[:, 0], max_xy[:, 0] - root_xy[:, 0]), + np.minimum(root_xy[:, 1] - min_xy[:, 1], max_xy[:, 1] - root_xy[:, 1]), + ) + outside_dx = np.clip(np.maximum(min_xy[:, 0] - root_xy[:, 0], root_xy[:, 0] - max_xy[:, 0]), 0.0, None) + outside_dy = np.clip(np.maximum(min_xy[:, 1] - root_xy[:, 1], root_xy[:, 1] - max_xy[:, 1]), 0.0, None) + outside_dist = np.sqrt(outside_dx**2 + outside_dy**2) + + return np.where(inside, inside_dist, -outside_dist) + + +def _nanmean_or_nan(values: np.ndarray) -> float: + """均值；空数组返回 nan（而不是 0，也不是 numpy 的警告）。""" + if values.size == 0: + return float("nan") + return float(np.mean(values)) + + +def compute_metrics( + trace: Trace, + profile: RobotProfile, + protocol: EvalProtocol, + *, + include_per_env: bool = False, +) -> dict[str, Any]: + """由一次 rollout 的轨迹算出全部指标。 + + 指标名的命名与旧 ``eval.py`` 完全一致，便于新旧结果直接对照。 + """ + num_steps = trace.num_steps + num_envs = trace.num_envs + dt = float(trace.dt) + + if abs(protocol.dt - dt) > 0.01 * dt: + raise ValueError( + f"协议 dt={protocol.dt} 与轨迹 dt={dt} 不一致。" + "通常意味着协议里的步长与任务的 decimation 不匹配。" + ) + + command = np.asarray(trace.command, dtype=np.float64) + command_xy = command[:2] + command_yaw = float(command[2]) + + command_dir = command_xy.copy() + norm = float(np.linalg.norm(command_dir)) + command_dir = command_dir / norm if norm > 1e-6 else np.array([1.0, 0.0]) + + # ------------------------------------------------------------------ 成功判定 + base_xy = np.asarray(trace.base_pos_w, dtype=np.float64)[..., :2] + delta_xy = base_xy - base_xy[0][None, :, :] + progress = np.einsum("tnc,c->tn", delta_xy[_ALL_FRAMES], command_dir) + max_progress = progress.max(axis=0) + reached = max_progress >= protocol.success_distance_m + + tilt = tilt_from_projected_gravity(np.asarray(trace.projected_gravity_b, dtype=np.float64)) + tilt_steps = tilt[_STEP_FRAMES] + max_tilt = tilt_steps.max(axis=0) + ever_fallen = max_tilt > protocol.tilt_threshold_rad + + success = reached & ~ever_fallen + success_count = int(np.count_nonzero(success)) + + # 首次达成成功的时刻：第 i 帧（i>=1）达成则记 i*dt。 + crossed = progress[_STEP_FRAMES] >= protocol.success_distance_m + complete_time = np.full((num_envs,), np.nan, dtype=np.float64) + for env_index in range(num_envs): + hit = np.flatnonzero(crossed[:, env_index]) + if hit.size: + complete_time[env_index] = float(hit[0] + 1) * dt + + successful_time = complete_time[success & np.isfinite(complete_time)] + avg_time = _nanmean_or_nan(successful_time) + + # ------------------------------------------------------------- 速度跟踪误差 + lin_vel_b = np.asarray(trace.base_lin_vel_b, dtype=np.float64)[_STEP_FRAMES] + ang_vel_b = np.asarray(trace.base_ang_vel_b, dtype=np.float64)[_STEP_FRAMES] + lin_error = np.linalg.norm(lin_vel_b[..., :2] - command_xy[None, None, :], axis=-1) + yaw_error = np.abs(ang_vel_b[..., 2] - command_yaw) + + # -------------------------------------------------------------- 能耗与路程 + joint_torque = np.asarray(trace.joint_torque, dtype=np.float64)[_STEP_FRAMES] + joint_vel = np.asarray(trace.joint_vel, dtype=np.float64)[_STEP_FRAMES] + mechanical_power = joint_torque * joint_vel + if protocol.energy_convention == "absolute": + integrand = np.abs(mechanical_power) + elif protocol.energy_convention == "signed": + integrand = mechanical_power + else: + raise ValueError( + f"未知能耗约定 {protocol.energy_convention!r}；可选 'absolute' 或 'signed'。" + ) + energy = integrand.sum(axis=(0, 2)) * dt + + lin_vel_w = np.asarray(trace.base_lin_vel_w, dtype=np.float64) + route = np.linalg.norm(lin_vel_w[_STEP_FRAMES][..., :2], axis=-1).sum(axis=0) * dt + + # ------------------------------------------------------------ 运动不稳定性 + accel = np.linalg.norm(np.diff(lin_vel_w, axis=0), axis=-1) / dt + instability_count = (accel > protocol.acc_threshold_mps2).sum(axis=0) + + stable_time = (tilt_steps < protocol.tilt_threshold_rad).sum(axis=0) * dt + + # ------------------------------------------------------------- 支撑多边形 + feet_pos_w = np.asarray(trace.feet_pos_w, dtype=np.float64)[_STEP_FRAMES] + num_feet = int(feet_pos_w.shape[2]) + if num_feet > 0: + margins = support_margin_xy(base_xy[_STEP_FRAMES][:, :, :].reshape(-1, 2), feet_pos_w.reshape(-1, num_feet, 2)) + margins = margins.reshape(feet_pos_w.shape[0], num_envs) + avg_margin = _nanmean_or_nan(margins) + margin_method = protocol.stability_margin_method + else: + avg_margin = float("nan") + margin_method = "unavailable" + + # --------------------------------------------------------------- 支撑接触 + support_force = np.asarray(trace.support_force_n, dtype=np.float64) + support_contact_time = np.asarray(trace.support_contact_time_s, dtype=np.float64) + num_support_bodies = int(support_force.shape[2]) + support_threshold = ( + profile.support_force_threshold_n + if protocol.support_force_threshold_n is None + else protocol.support_force_threshold_n + ) + if num_support_bodies > 0: + active = np.any(support_force[_STEP_FRAMES] > support_threshold, axis=2) # (T-1, N) + support_steps = active.sum(axis=0) + # 起始帧之前视为「未接触」，因此首帧接触也算一次事件（与旧实现一致）。 + active_padded = np.concatenate([np.zeros((1, num_envs), dtype=bool), active], axis=0) + support_events = (active_padded[1:] & ~active_padded[:-1]).sum(axis=0) + max_support_time = support_contact_time[_STEP_FRAMES].max(axis=(0, 2)) + support_contact_ratio = _nanmean_or_nan(support_steps / max(num_steps, 1)) + avg_support_events = float(np.mean(support_events)) + avg_max_support_time = float(np.mean(max_support_time)) + if success_count > 0: + success_support_ratio = float(np.mean((support_events[success] > 0).astype(np.float64))) + else: + success_support_ratio = float("nan") + else: + support_contact_ratio = float("nan") + avg_support_events = float("nan") + avg_max_support_time = float("nan") + success_support_ratio = float("nan") + + # --------------------------------------------------------- 质量与 COT + mass = trace.mass_kg if trace.mass_kg is not None else profile.mass_kg + if mass is None or mass <= 0.0: + cost_of_transport = np.full((num_envs,), np.nan) + mass_used = float("nan") + else: + cost_of_transport = energy / ((route + 1e-6) * mass * _GRAVITY) + mass_used = float(mass) + + metrics: dict[str, Any] = { + "profile": profile.name, + "protocol_fingerprint": protocol.fingerprint(), + "num_envs": float(num_envs), + "num_steps": float(num_steps), + "episode_length_s": float(num_steps * dt), + "terrain_size_m": float(protocol.terrain_size_m), + "success_distance_m": float(protocol.success_distance_m), + "robot_mass_kg": mass_used, + "energy_convention": protocol.energy_convention, + "stability_margin_method": margin_method, + "tilt_threshold_rad": float(protocol.tilt_threshold_rad), + "acc_threshold_mps2": float(protocol.acc_threshold_mps2), + "support_force_threshold_n": float(support_threshold) if num_support_bodies > 0 else float("nan"), + # --- 成功率与进度 + "success_count": float(success_count), + "success_rate": float(success_count / max(num_envs, 1)), + "avg_final_progress_m": float(np.mean(max_progress)), + # --- 时间 + "avg_time_consumption_s": avg_time, + "time_efficiency": float( + avg_time / (protocol.success_distance_m / max(abs(float(command[0])), 1e-6)) + ) + if np.isfinite(avg_time) + else float("nan"), + # --- 跟踪误差 + "mean_lin_vel_error": float(np.mean(lin_error)), + "mean_yaw_vel_error": float(np.mean(yaw_error)), + # --- 能耗 + "avg_energy_j": float(np.mean(energy)), + "avg_cost_of_transport": _nanmean_or_nan(cost_of_transport), + # --- 稳定性 + "unstability_rate": float(np.mean(instability_count / max(num_steps, 1))), + "avg_stability_margin_m": avg_margin, + "avg_max_tilt_rad": float(np.mean(max_tilt)), + "avg_stable_time_s": float(np.mean(stable_time)), + # --- 支撑接触 + "avg_support_contact_ratio": support_contact_ratio, + "avg_support_events": avg_support_events, + "avg_max_support_contact_time_s": avg_max_support_time, + "successful_runs_with_support_ratio": success_support_ratio, + } + + if include_per_env: + metrics["_per_env"] = { + "success": success.astype(float).tolist(), + "max_progress": max_progress.tolist(), + "max_tilt": max_tilt.tolist(), + } + return metrics diff --git a/eval_kit/core/naming.py b/eval_kit/core/naming.py new file mode 100644 index 0000000..07eac9b --- /dev/null +++ b/eval_kit/core/naming.py @@ -0,0 +1,94 @@ +# Copyright (c) 2025 Deep Robotics +# SPDX-License-Identifier: BSD 3-Clause + +"""按正则列表解析名字，与 IsaacLab 的 ``resolve_matching_names`` 语义一致， +但**匹配不到不抛异常**。 + +IsaacLab 的实现用 ``re.fullmatch`` 逐个比对，并且：某个名字同时命中多个正则会报错， +某个正则一个都没命中也报错。对环境配置来说这种严格是有用的（写错了应该立刻发现）， +但评价代码不该因为「这个机器人没有轮子」就崩掉——那时正确行为是跳过依赖轮子的指标。 + +这个模块是纯逻辑，不依赖 isaaclab，因此可以单独测试。 +""" + +from __future__ import annotations + +import re +from typing import Sequence + + +class NameResolution: + """一次名字解析的结果。""" + + __slots__ = ("indices", "names", "unmatched_patterns", "ambiguous_names") + + def __init__( + self, + indices: list[int], + names: list[str], + unmatched_patterns: list[str], + ambiguous_names: list[tuple[str, list[str]]], + ) -> None: + self.indices = indices + self.names = names + #: 一个都没命中的正则。 + self.unmatched_patterns = unmatched_patterns + #: 同时命中多个正则的名字，以及命中的正则。 + self.ambiguous_names = ambiguous_names + + def __bool__(self) -> bool: + return bool(self.indices) + + def __len__(self) -> int: + return len(self.indices) + + def __repr__(self) -> str: + return ( + f"NameResolution({len(self.indices)} 项, " + f"未命中={self.unmatched_patterns}, 多重命中={self.ambiguous_names})" + ) + + +def resolve_matching_indices(patterns: Sequence[str], names: Sequence[str]) -> NameResolution: + """把正则列表解析成 ``names`` 中的下标。 + + 行为与 IsaacLab 的 ``resolve_matching_names`` 一致之处：用 ``re.fullmatch``； + 返回顺序遵循 ``patterns`` 的顺序。 + 不同之处：不抛异常，而是把「未命中」和「多重命中」记在返回值里，由调用方决定 + 是警告还是忽略。 + + Args: + patterns: 正则列表，例如 ``("base_link", ".*_knee")``。 + names: 候选名字，例如 ``articulation.body_names``。 + """ + hits: dict[int, list[str]] = {} + pattern_hits: dict[str, int] = {pattern: 0 for pattern in patterns} + + for index, name in enumerate(names): + for pattern in patterns: + try: + matched = re.fullmatch(pattern, name) is not None + except re.error as exc: # 正则本身写错，属于配置错误，应当立刻暴露 + raise ValueError(f"机器人描述里的正则 {pattern!r} 无法编译：{exc}") from exc + if matched: + hits.setdefault(index, []).append(pattern) + pattern_hits[pattern] += 1 + + indices: list[int] = [] + resolved_names: list[str] = [] + ambiguous: list[tuple[str, list[str]]] = [] + taken: set[int] = set() + + # 按 patterns 顺序输出，跳过已被更早的正则取走的下标 + for pattern in patterns: + for index, matched_patterns in hits.items(): + if pattern not in matched_patterns or index in taken: + continue + taken.add(index) + indices.append(index) + resolved_names.append(names[index]) + if len(matched_patterns) > 1: + ambiguous.append((names[index], list(matched_patterns))) + + unmatched = [pattern for pattern, count in pattern_hits.items() if count == 0] + return NameResolution(indices, resolved_names, unmatched, ambiguous) diff --git a/eval_kit/core/report.py b/eval_kit/core/report.py new file mode 100644 index 0000000..c5e98fa --- /dev/null +++ b/eval_kit/core/report.py @@ -0,0 +1,191 @@ +# Copyright (c) 2025 Deep Robotics +# SPDX-License-Identifier: BSD 3-Clause + +"""结果的格式化与筛选：纯函数，只依赖 numpy 与标准库。 + +输出格式与旧 ``eval.py`` 保持一致，便于对照。``filter_metrics`` 决定每种 +``eval_type`` 关心哪些指标。 +""" + +from __future__ import annotations + +from typing import Any, Sequence + +import numpy as np + +#: 支持的评估类型。 +EVAL_TYPES: tuple[str, ...] = ( + "all", + "upper_limits", + "velocity_error", + "access_rate", + "time_consumption", + "energy_consumption", + "motion_stability", + "stability_margin", + "robustness", +) + +#: 每种评估类型关心哪些指标。 +METRIC_SELECTION: dict[str, tuple[str, ...]] = { + "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"), +} + + +def filter_metrics(metrics: dict[str, Any], eval_type: str) -> dict[str, Any]: + """按评估类型裁剪指标字典；``all`` 返回全部。""" + if eval_type == "all": + return metrics + if eval_type not in METRIC_SELECTION: + raise KeyError(f"未知评估类型 {eval_type!r}；可选：{sorted(EVAL_TYPES)}") + return {key: metrics[key] for key in METRIC_SELECTION[eval_type] if key in metrics} + + +def _format_cell(value: float, width: int = 14) -> str: + if value is None or not np.isfinite(value): + return " -".rjust(width) + return f"{value:.2f}".rjust(width) + + +def format_matrix(matrix: Sequence[Sequence[float]], terrain_types: Sequence[str], *, title: str) -> str: + """把「行=等级、列=地形」的矩阵格式化成字符串。""" + width = 14 + header = "level".ljust(8) + "".join(terrain_type[: width - 1].rjust(width) for terrain_type in terrain_types) + lines = [title, header] + for level_index, row in enumerate(matrix, start=1): + cells = "".join(_format_cell(float(value), width) for value in row) + lines.append(f"{level_index:<8}{cells}") + return "\n".join(lines) + + +def format_numpy_matrix(matrix: Sequence[Sequence[float]], *, title: str) -> str: + """紧凑的 numpy 风格矩阵，便于贴进报告；nan 显示成 ``-``。 + + 用 ``formatter`` 而不是 ``nanstr``——后者是 ``np.set_printoptions`` 的参数， + 不是 ``np.array2string`` 的关键字（numpy 1.21 与 1.26 都会对前者报 TypeError）。 + """ + array = np.round(np.array(matrix, dtype=np.float32), 2) + + def _cell(value: float) -> str: + return " -" if not np.isfinite(value) else f"{value:.2f}" + + body = np.array2string(array, formatter={"float_kind": _cell}) + return f"{title}:\n{body}" + + +def format_upper_limit_summary(result: dict[str, Any]) -> str: + """上限评估的汇总输出：成功率矩阵、累积成功率矩阵、各类型 score。""" + terrain_types = result["terrain_types"] + matrix = result["matrix"] + scores = np.array([result["scores"][terrain_type] for terrain_type in terrain_types], dtype=np.float32) + lines = [ + format_numpy_matrix(matrix["success_rate"], title="success rate"), + format_numpy_matrix(matrix["success_rate_accumulated"], title="success rate accumulated"), + f"score: {np.array2string(scores, precision=2, suppress_small=True)} of {result['num_levels']}", + ] + units = {unit for unit in result.get("difficulty_unit", {}).values() if unit} + if units: + lines.append(f"difficulty units: {'; '.join(sorted(units))}") + return "\n".join(lines) + + +def format_rollout_summary(metrics: dict[str, Any]) -> str: + """混合地形汇总跑的全指标输出。""" + + def fmt(key: str, spec: str = "{:.2f}") -> str: + value = float(metrics.get(key, float("nan"))) + return "nan" if not np.isfinite(value) else spec.format(value) + + success_ratio = float(metrics.get("successful_runs_with_support_ratio", float("nan"))) + success_ratio_text = "nan" if not np.isfinite(success_ratio) else f"{success_ratio * 100:.1f}%" + support_ratio = float(metrics.get("avg_support_contact_ratio", float("nan"))) + support_ratio_text = "nan" if not np.isfinite(support_ratio) else f"{support_ratio * 100:.1f}%" + mass = float(metrics.get("robot_mass_kg", float("nan"))) + mass_text = "nan" if not np.isfinite(mass) else f"{mass:.3f}" + + return "\n".join( + [ + f"success rate: {fmt('success_rate')}", + f"success count: {int(round(float(metrics.get('success_count', 0.0))))} " + f"time consumption: {fmt('avg_time_consumption_s')}", + f"time efficiency: {fmt('time_efficiency')}", + f"mean lin vel error: {fmt('mean_lin_vel_error')} " + f"mean yaw vel error: {fmt('mean_yaw_vel_error')}", + f"average energy: {fmt('avg_energy_j')} J " + f"average cost of transport: {fmt('avg_cost_of_transport')}", + f"unstability rate: {fmt('unstability_rate', '{:.3f}')} " + f"average stability margin: {fmt('avg_stability_margin_m')}", + f"average max tilt: {fmt('avg_max_tilt_rad')} rad " + f"average stable time: {fmt('avg_stable_time_s')} s", + f"support contact ratio: {support_ratio_text} " + f"events/episode: {fmt('avg_support_events')} " + f"max duration: {fmt('avg_max_support_contact_time_s')} s " + f"successful runs using support: {success_ratio_text}", + f"[protocol] mass={mass_text} kg energy={metrics.get('energy_convention')} " + f"margin={metrics.get('stability_margin_method')} " + f"tilt>={fmt('tilt_threshold_rad')} acc>={fmt('acc_threshold_mps2')} " + f"fingerprint={metrics.get('protocol_fingerprint')}", + ] + ) + + +def format_aggregate_summary(aggregate: dict[str, Any], keys: Sequence[str] | None = None) -> str: + """多次试验的 ``mean ± std`` 汇总。 + + 做机器人对比时应看这个而不是单次数值。 + """ + summary = aggregate.get("summary", {}) + selected = list(keys) if keys is not None else list(summary.keys()) + lines = [f"runs = {aggregate.get('num_runs')} seeds = {aggregate.get('seeds')}"] + for key in selected: + entry = summary.get(key) + if entry is None: + continue + mean = float(entry["mean"]) + std = float(entry["std"]) + if not np.isfinite(mean): + lines.append(f" {key:<38} nan（{entry.get('finite_runs', 0)}/{aggregate.get('num_runs')} 次有效）") + else: + lines.append(f" {key:<38} {mean:>10.4f} ± {std:<8.4f} [{entry['min']:.4f}, {entry['max']:.4f}]") + return "\n".join(lines) + + +def format_metric_result(eval_type: str, metrics: dict[str, Any]) -> str: + """单一评估类型的精简输出。""" + + def fmt(key: str) -> str: + value = float(metrics.get(key, float("nan"))) + return "nan" if not np.isfinite(value) else f"{value:.2f}" + + if eval_type == "velocity_error": + return f"mean lin vel error: {fmt('mean_lin_vel_error')}\nmean yaw vel error: {fmt('mean_yaw_vel_error')}" + if eval_type == "access_rate": + return f"success rate: {fmt('success_rate')}\navg final progress: {fmt('avg_final_progress_m')} m" + if eval_type == "time_consumption": + return ( + f"success count: {int(round(float(metrics.get('success_count', 0.0))))} " + f"time consumption: {fmt('avg_time_consumption_s')}\n" + f"time efficiency: {fmt('time_efficiency')}" + ) + if eval_type == "energy_consumption": + return f"average energy: {fmt('avg_energy_j')} J\naverage cost of transport: {fmt('avg_cost_of_transport')}" + if eval_type == "motion_stability": + return f"unstability rate: {fmt('unstability_rate')}" + if eval_type == "stability_margin": + return ( + f"average stability margin: {fmt('avg_stability_margin_m')} m " + f"(method: {metrics.get('stability_margin_method')})" + ) + if eval_type == "robustness": + return ( + f"success rate: {fmt('success_rate')}\n" + f"average max tilt: {fmt('avg_max_tilt_rad')} rad\n" + f"average stable time: {fmt('avg_stable_time_s')} s" + ) + return format_rollout_summary(metrics) diff --git a/eval_kit/core/schema.py b/eval_kit/core/schema.py new file mode 100644 index 0000000..9afbda4 --- /dev/null +++ b/eval_kit/core/schema.py @@ -0,0 +1,392 @@ +# Copyright (c) 2025 Deep Robotics +# SPDX-License-Identifier: BSD 3-Clause + +"""评价套件的数据契约。 + +三个对象构成整条流水线： + +``RobotProfile`` 机器人形态（谁是足端、谁是支撑体、质量多少） +``EvalProtocol`` 评估协议（指令、时长、各种阈值、能耗约定） +``Trace`` 一次 rollout 的标准化逐帧记录 + +上游（任何框架的适配器）负责产出 ``Trace``，下游（:mod:`eval_kit.core.metrics`） +只从 ``Trace`` 计算指标。指标代码因此不接触任何仿真对象。 + +这个模块只依赖 numpy 与标准库。 +""" + +from __future__ import annotations + +import hashlib +import json +from dataclasses import asdict, dataclass +from typing import Any + +import numpy as np + +#: 帧索引约定：``Trace`` 的第 0 帧是执行动作之前的起始状态， +#: 第 ``i`` 帧（``i >= 1``）是第 ``i`` 步动作之后的状态。 +#: 因此一次 ``steps`` 步的 rollout 会产生 ``steps + 1`` 帧。 +START_FRAME = 0 + +#: 成功判定与最大进度使用全部帧（含起始帧），与 legged_gym 侧 +#: ``criterion.py`` 及现有 ``eval.py`` 的取点一致。 +PROGRESS_FRAMES = slice(None) + +#: 逐帧累积类指标只使用动作之后的帧。 +STEP_FRAMES = slice(1, None) + + +@dataclass(frozen=True) +class RobotProfile: + """机器人形态描述。 + + 评价代码里不应出现具体的 link/joint 名字或质量常数，全部集中在这里。 + 新增一个机器人 = 新增一条 ``RobotProfile`` 记录，不改评价代码。 + """ + + name: str + """机器人标识，同时也是输出文件名的一部分。""" + + mass_kg: float | None + """整机质量 [kg]；``None`` 表示运行时从仿真资产读取（推荐）。""" + + feet_body_patterns: tuple[str, ...] = () + """支撑多边形的接触点，例如火蚁的四个轮。空元组表示该机器人不计算支撑裕度。""" + + support_body_patterns: tuple[str, ...] = () + """允许作为「短暂支撑」的机体，例如火蚁的机身与膝部。空元组表示不做支撑统计。""" + + actuated_joint_patterns: tuple[str, ...] = () + """参与能耗统计的关节。""" + + support_force_threshold_n: float = 15.0 + """判定「支撑接触」的法向力阈值 [N]。""" + + body_length_m: float = 0.0 + """机身长度 [m]，仅用于校验地形瓦片尺寸是否够大。""" + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +@dataclass(frozen=True) +class EvalProtocol: + """评估协议：一次评估的全部「口径」。 + + 这些字段决定数值大小，因此都会被写进输出 JSON 的 ``fingerprint``， + 避免不同设置下的结果被悄悄拿来比较。 + + 默认值与现有 ``scripts/reinforcement_learning/rsl_rl/eval.py`` 的默认 + 命令行参数一致（能耗用绝对值口径、倾角阈值 0.70 rad、加速度阈值 2.0 m/s² + 折算成重力倍数），因此默认协议下的新结果可以直接和旧评估结果对照。 + """ + + steps: int = 500 + """每个 episode 的控制步数。dt=0.02 时 500 步 = 10 s。""" + + dt: float = 0.02 + """控制步长 [s]，即 ``decimation * sim.dt``。""" + + command: tuple[float, float, float] = (1.0, 0.0, 0.0) + """固定速度指令 ``(lin_vel_x, lin_vel_y, ang_vel_z)``。""" + + start_pose: tuple[float, float, float] = (-0.9, 0.0, 0.0) + """起始位姿 ``(x, y, yaw)``，相对地形瓦片中心。""" + + start_z_offset_m: float = 0.0 + """起始高度相对地形瓦片中心高度的额外抬升 [m]。 + + 默认 0.0 表示把机器人直接放在 ``env_origin.z`` 上，与旧 ``eval.py`` 的写法一致。 + 但那样只对「瓦片中心是最高点」的地形成立：像倒金字塔台阶（``pyramid_stairs_inv``， + 中心是最低点）这类地形，机器人在 ``|x| > platform`` 处会被直接生成到地形几何体 + 内部，卡住无法移动，导致该地形在所有难度上都得到 0 成功率。 + + 给一个正值（例如 0.3 m）让机器人从地形表面上方落下，可以适配任意地形形状。 + 改这个值会改变 ``protocol_fingerprint``。 + """ + + success_distance_m: float = 2.0 + """沿指令方向前进多少米算成功。""" + + tilt_threshold_rad: float = 0.70 + """倾角超过该值 [rad] 视为摔倒。0.70 rad ≈ 40°。""" + + acc_threshold_g: float = 2.0 / 9.81 + """机身加速度超过该值（单位 g）计为一次不稳定。用重力倍数而非绝对 m/s²， + 使该阈值对不同量级的机器人保持可比。默认值 2.0 m/s² 折算而来。""" + + energy_convention: str = "absolute" + """能耗积分约定，``"absolute"`` 表示 ``Σ|τ·q̇|·dt``， + ``"signed"`` 表示 ``Σ(τ·q̇)·dt``。两者数值差异可达一倍以上， + 全局必须一致；该值会写进输出 JSON。""" + + stability_margin_method: str = "foot_bbox" + """支撑裕度算法。``"foot_bbox"`` = 机体到足端轴对齐包围盒的距离（当前实现）。""" + + support_force_threshold_n: float | None = 15.0 + """判定支撑接触的法向力阈值 [N]；``None`` 时改用机器人描述里的值。""" + + enable_interval_pushes: bool = False + """是否启用周期性外力扰动（``robustness`` 评估用）。""" + + enforce_terrain_bounds: bool = True + """是否保留 IsaacLab 的 ``terrain_out_of_bounds`` 终止项。 + + **默认 True 是为了与旧 ``eval.py`` 的结果可比，但已知它有系统性偏差**： + 该终止项比较的是机器人的世界坐标与整张地图的范围，而不是它自己那块地砖。 + 在地砖 2 m、网格 5x5 或 8xN 的评估配置下，最外侧的格子会被误判为越界， + 机器人从第 0 步起每步都被重置，永远没机会尝试地形；而 ``|y|`` 恰好等于阈值 + 的那两列则处于浮点边界上。实测：8x2 网格下首步就有 6/32 个环境被误判。 + + 三种处置方式（建议前两种，按顺序）： + + 1. 把 :attr:`terrain_size_m` 提到 **≥ 3.8**，即可让所有起始位置落进阈值内， + 同时保留这个终止项； + 2. 置为 ``False``，评估时去掉这个终止项（成功判定本来就是按前进距离算的， + 走出地砖不构成失败）； + 3. 保持 ``True`` 且用小地砖 —— 数值会偏低，且被影响的难度等级因行序打乱而 + 每次不同，不建议用来做对比。 + """ + + push_interval_s: tuple[float, float] = (2.0, 2.01) + """扰动间隔 [s]，取一个极窄区间让扰动近似等周期但可比较。""" + + push_velocity_range: tuple[float, float] = (-0.5, 0.5) + """每个水平方向上的扰动速度范围 [m/s]。""" + + terrain_spec: str = "standard_v1" + """地形等级规格名，见 :mod:`eval_kit.core.terrain_spec`。""" + + terrain_size_m: float = 2.0 + """单块地形瓦片边长 [m]。""" + + num_levels: int = 8 + """上限评估的难度等级数（网格行数）。""" + + envs_per_cell: int = 1 + """上限评估每个「地形类型 × 难度」格子里放几个机器人。 + + 默认 1。**注意：由于起始位姿被钉死，同一格的多个机器人起始坐标完全相同**， + 它们不是独立样本，差异只来自物理求解器的非确定性。所以靠「一格放多个」来提高 + 样本量是假精度；正确做法是 ``envs_per_cell=1`` + 多个 :attr:`seeds` + （不同种子会生成不同的地形实例，因此跨种子才是独立样本）。 + + 代价：``envs_per_cell=1`` 时单次评估里每格成功率只能是 0 或 1， + 必须配合足够多的种子才能得到有意义的分辨率。经验值：种子数 ≥ 8 时 + 每格分辨率达到 1/8 ≈ 12.5%。 + """ + + seeds: tuple[int, ...] = (0,) + """重复试验的随机种子。单次 rollout 噪声很大，比较机器人时应跑多个种子。 + + 每个种子会重新生成地形实例（``terrain_generator.seed = seed``），所以跨种子 + 的样本是独立的；这是提高统计可靠性的正确途径，而不是加大 ``envs_per_cell``。 + """ + + @property + def acc_threshold_mps2(self) -> float: + """加速度阈值的 m/s² 形式，供指标计算使用。""" + return self.acc_threshold_g * 9.81 + + def minimum_terrain_size_for_bounds(self, *, border_width: float = 2.0, distance_buffer: float = 3.0) -> float: + """保留 ``terrain_out_of_bounds`` 终止项且不误判所需的最小地砖边长 [m]。 + + 这个终止项判定 ``|origin + start| <= 0.5 * 地图边长 - distance_buffer``， + 而最外侧格子中心在 ``(rows-1)/2 * tile``，地图边长 ``= rows * tile + 2 * border``。 + 两边的 ``rows`` 项相消，于是结论与网格行数无关： + + ``tile >= 2 * (|start_x| + distance_buffer - border_width)`` + + 默认的 ``border_width=2.0`` / ``distance_buffer=3.0`` 与 + :func:`eval_kit.isaaclab.env_setup.apply_eval_terrain_baseline` 及 IsaacLab 的 + 默认值一致；取 ``start_x = -0.9`` 得 ``tile >= 3.8 m``。 + + 放在内核里（而不是适配层）是为了让编排进程也能用——它不启动仿真， + 无法 import 任何依赖 isaaclab 的模块。 + """ + return max(0.0, 2.0 * (abs(self.start_pose[0]) + distance_buffer - border_width)) + + def fingerprint(self) -> str: + """协议指纹：决定数值的字段的稳定哈希，写进输出 JSON。""" + payload = { + "steps": self.steps, + "dt": self.dt, + "command": list(self.command), + "start_pose": list(self.start_pose), + "start_z_offset_m": self.start_z_offset_m, + "success_distance_m": self.success_distance_m, + "tilt_threshold_rad": self.tilt_threshold_rad, + "acc_threshold_g": self.acc_threshold_g, + "energy_convention": self.energy_convention, + "stability_margin_method": self.stability_margin_method, + "support_force_threshold_n": self.support_force_threshold_n, + "enable_interval_pushes": self.enable_interval_pushes, + "enforce_terrain_bounds": self.enforce_terrain_bounds, + "push_interval_s": list(self.push_interval_s), + "push_velocity_range": list(self.push_velocity_range), + "terrain_spec": self.terrain_spec, + "terrain_size_m": self.terrain_size_m, + "num_levels": self.num_levels, + "envs_per_cell": self.envs_per_cell, + } + blob = json.dumps(payload, sort_keys=True).encode("utf-8") + return hashlib.sha256(blob).hexdigest()[:12] + + def to_dict(self) -> dict[str, Any]: + data = asdict(self) + data["acc_threshold_mps2"] = self.acc_threshold_mps2 + data["fingerprint"] = self.fingerprint() + return data + + +@dataclass +class Trace: + """一次 rollout 的标准化逐帧记录。 + + 所有数组的第一维都是帧、第二维都是环境。第 0 帧是动作执行之前的起始状态。 + + ``feet_pos_w`` / ``support_force_n`` / ``support_contact_time_s`` 的第二维 + 可能是 0（该机器人没有足端或没有支撑体语义），指标层会跳过对应指标。 + """ + + dt: float + command: np.ndarray + """``(3,)`` 本次 rollout 使用的固定指令。""" + + base_pos_w: np.ndarray + """``(T, N, 3)`` 机身世界位置。""" + + base_lin_vel_w: np.ndarray + """``(T, N, 3)`` 机身世界系线速度，用于路程与加速度（不稳定性）。""" + + base_lin_vel_b: np.ndarray + """``(T, N, 3)`` 机体系线速度，用于速度跟踪误差。""" + + base_ang_vel_b: np.ndarray + """``(T, N, 3)`` 机体系角速度，用第 2 个分量做 yaw 跟踪误差。""" + + projected_gravity_b: np.ndarray + """``(T, N, 3)`` 机体系投影重力，用于倾角。""" + + joint_torque: np.ndarray + """``(T, N, J)`` 实际施加的关节力矩。""" + + joint_vel: np.ndarray + """``(T, N, J)`` 关节速度。""" + + feet_pos_w: np.ndarray + """``(T, N, F, 2)`` 足端/轮心的世界 xy，用于支撑多边形。""" + + support_force_n: np.ndarray + """``(T, N, S)`` 支撑体的法向合力。""" + + support_contact_time_s: np.ndarray + """``(T, N, S)`` 支撑体的连续接触时长。""" + + done: np.ndarray + """``(T, N)`` 每步之后的终止标志（1.0 表示该步结束后被重置）。""" + + mass_kg: float | None = None + """实际使用的整机质量 [kg]。``None`` 表示上游未能从资产读到。""" + + def __post_init__(self) -> None: + self.validate() + + @property + def num_frames(self) -> int: + return int(self.base_pos_w.shape[0]) + + @property + def num_steps(self) -> int: + """实际执行的控制步数（帧数减一）。""" + return self.num_frames - 1 + + @property + def num_envs(self) -> int: + return int(self.base_pos_w.shape[1]) + + @property + def num_joints(self) -> int: + return int(self.joint_torque.shape[2]) + + def validate(self) -> None: + """形状一致性检查。早失败好过算出一堆形状广播出来的错数字。""" + frames, envs = self.num_frames, self.num_envs + if frames < 2: + raise ValueError(f"Trace 至少需要 2 帧（起始帧 + 至少一步），收到 {frames} 帧。") + + expected: dict[str, tuple[int, ...]] = { + "base_pos_w": (frames, envs, 3), + "base_lin_vel_w": (frames, envs, 3), + "base_lin_vel_b": (frames, envs, 3), + "base_ang_vel_b": (frames, envs, 3), + "projected_gravity_b": (frames, envs, 3), + "done": (frames, envs), + } + for name, shape in expected.items(): + array = np.asarray(getattr(self, name)) + if array.shape != shape: + raise ValueError(f"{name} 形状应为 {shape}，实际 {array.shape}。") + + joints = self.num_joints + for name in ("joint_torque", "joint_vel"): + if np.asarray(getattr(self, name)).shape != (frames, envs, joints): + raise ValueError( + f"{name} 形状应为 {(frames, envs, joints)}，实际 {np.asarray(getattr(self, name)).shape}。" + ) + + feet = np.asarray(self.feet_pos_w) + if feet.shape[:2] != (frames, envs) or feet.ndim != 4 or feet.shape[3] != 2: + raise ValueError(f"feet_pos_w 形状应为 {(frames, envs, feet.shape[2] if feet.ndim > 2 else 0, 2)}，" + f"实际 {feet.shape}。") + + for name in ("support_force_n", "support_contact_time_s"): + array = np.asarray(getattr(self, name)) + if array.ndim != 3 or array.shape[:2] != (frames, envs): + raise ValueError(f"{name} 形状应为 {(frames, envs, array.shape[2] if array.ndim > 2 else 0)}，" + f"实际 {array.shape}。") + + if np.asarray(self.command).shape != (3,): + raise ValueError(f"command 应为 (3,)，实际 {np.asarray(self.command).shape}。") + + def save(self, path: str) -> None: + """把轨迹存成 ``.npz``，便于事后用新指标重新分析，无需重跑仿真。""" + np.savez_compressed( + path, + dt=np.asarray(self.dt), + command=self.command, + base_pos_w=self.base_pos_w, + base_lin_vel_w=self.base_lin_vel_w, + base_lin_vel_b=self.base_lin_vel_b, + base_ang_vel_b=self.base_ang_vel_b, + projected_gravity_b=self.projected_gravity_b, + joint_torque=self.joint_torque, + joint_vel=self.joint_vel, + feet_pos_w=self.feet_pos_w, + support_force_n=self.support_force_n, + support_contact_time_s=self.support_contact_time_s, + done=self.done, + mass_kg=np.asarray(-1.0 if self.mass_kg is None else self.mass_kg), + ) + + @classmethod + def load(cls, path: str) -> Trace: + with np.load(path) as data: + mass = float(data["mass_kg"]) + return cls( + dt=float(data["dt"]), + command=data["command"], + base_pos_w=data["base_pos_w"], + base_lin_vel_w=data["base_lin_vel_w"], + base_lin_vel_b=data["base_lin_vel_b"], + base_ang_vel_b=data["base_ang_vel_b"], + projected_gravity_b=data["projected_gravity_b"], + joint_torque=data["joint_torque"], + joint_vel=data["joint_vel"], + feet_pos_w=data["feet_pos_w"], + support_force_n=data["support_force_n"], + support_contact_time_s=data["support_contact_time_s"], + done=data["done"], + mass_kg=None if mass < 0.0 else mass, + ) diff --git a/eval_kit/core/terrain_spec.py b/eval_kit/core/terrain_spec.py new file mode 100644 index 0000000..d94164c --- /dev/null +++ b/eval_kit/core/terrain_spec.py @@ -0,0 +1,157 @@ +# Copyright (c) 2025 Deep Robotics +# SPDX-License-Identifier: BSD 3-Clause + +"""地形等级规格：纯数据，与仿真框架无关。 + +规格描述「评估用哪几种地形、每种地形分几级、每级多高」，用绝对高度表示， +因为机器人越障能力本来就是由绝对高度决定的。 + +规格本身不含任何框架信息；「把规格翻译成本框架的地形生成器参数」这件事 +由各框架的 builder 负责（当前是 :mod:`eval_kit.isaaclab.env_setup`）。 +换框架时改 builder，不改这里的数值。 + +数值与 ``scripts/reinforcement_learning/rsl_rl/eval.py`` 的 ``TERRAIN_LEVELS`` +保持一致，因此两套评估所用地形严格相同。 +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Mapping + +#: 无量纲坡度：高 / 水平距离。0.40 相当于 21.8°。 +_UNIT_SLOPE = "slope (rise/run, 0.40 ≈ 21.8°)" +#: 台阶高度，米。 +_UNIT_METER = "m" + + +@dataclass(frozen=True) +class TerrainSpec: + """一套地形评估规格。""" + + name: str + """规格名，写进输出 JSON。""" + + types: tuple[str, ...] + """参与评估的地形类型，同时也是上限矩阵的列顺序。""" + + levels: Mapping[str, tuple[float, ...]] + """每种地形的难度等级值（绝对量，单位见 :attr:`unit`），由易到难。""" + + unit: Mapping[str, str] + """每种地形等级值的单位说明。""" + + isaaclab_subterrain: Mapping[str, str] + """地形类型 → IsaacLab ``ROUGH_TERRAINS_CFG`` 里的子地形名。""" + + isaaclab_param: Mapping[str, str] + """地形类型 → 需要在子地形配置里写入的字段名。""" + + feature_scaling: Mapping[str, tuple[float, float]] = field(default_factory=dict) + """瓦片尺寸缩小时的特征尺寸上限：``字段名 -> (相对瓦片的比例, 绝对上限)``。 + + 例如 ``{"step_width": (0.125, 0.25)}`` 表示 ``step_width = min(0.25, tile * 0.125)``。 + 这样同一套规格可以用在不同大小的地形瓦片上。 + """ + + grid_rows: int = 5 + grid_cols: int = 5 + """非上限评估时地形生成器的网格规模。""" + + isaaclab_extra_params: Mapping[str, Mapping[str, float]] = field(default_factory=dict) + """地形类型 → 需要额外写入的固定字段，例如 ``{"rough_slope": {"noise_step": 0.02}}``。 + + 与 :attr:`isaaclab_param` 的区别：后者写的是随等级变化的首尾值，这里写的是与等级 + 无关的固定量。写入时会校验字段是否存在，写错会直接报错而不是静默造一个假属性。 + """ + + def levels_for(self, terrain_type: str, num_levels: int | None = None) -> tuple[float, ...]: + """取某种地形的等级值，``num_levels`` 为 None 时取全部。""" + if terrain_type not in self.levels: + raise KeyError(f"地形规格 {self.name!r} 里没有 {terrain_type!r}；可选：{sorted(self.levels)}") + values = tuple(self.levels[terrain_type]) + if num_levels is None: + return values + if num_levels > len(values): + raise ValueError( + f"地形 {terrain_type!r} 只定义了 {len(values)} 级难度，请求了 {num_levels} 级。" + ) + return values[:num_levels] + + def scaling_for(self, field_name: str, tile_size_m: float) -> float | None: + """按瓦片尺寸算某个特征字段的上限值，未定义则返回 None。""" + rule = self.feature_scaling.get(field_name) + if rule is None: + return None + ratio, cap = rule + return min(cap, tile_size_m * ratio) + + def to_dict(self) -> dict[str, object]: + return { + "name": self.name, + "types": list(self.types), + "levels": {key: list(value) for key, value in self.levels.items()}, + "unit": dict(self.unit), + } + + +STANDARD_V1 = TerrainSpec( + name="standard_v1", + types=("stairs_up", "stairs_down", "smooth_slope", "rough_slope", "boxes"), + 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), + # 注意：random_rough 用的是 HfRandomUniformTerrainCfg，IsaacLab 明确写了 + # 「difficulty 参数被忽略」，所以这里只有首尾值真正起作用 —— 写成 + # noise_range=(levels[0], levels[-1])，是**绝对高度范围**而非起伏幅度。 + # 旧 eval.py 把它设成 (0.05, 0.40)，在 2 cm 水平网格上等于 40 cm 高的随机尖刺， + # 轮半径 0.096 m 的机器人无法通过，那一列几乎恒为 0（实测只从 0.12 升到 0.45）。 + # 这里回到与训练地形同量级的范围（上游训练用 noise_range=(0.02, 0.10)）。 + "rough_slope": (0.02, 0.03, 0.04, 0.05, 0.06, 0.08, 0.09, 0.10), + "boxes": (0.04, 0.06, 0.08, 0.10, 0.12, 0.14, 0.16, 0.18), + }, + unit={ + "stairs_up": _UNIT_METER, + "stairs_down": _UNIT_METER, + "smooth_slope": _UNIT_SLOPE, + "rough_slope": _UNIT_METER + "（高度范围；此地形忽略难度参数，8 行同分布，等级排名无意义）", + "boxes": _UNIT_METER, + }, + isaaclab_subterrain={ + "stairs_up": "pyramid_stairs", + "stairs_down": "pyramid_stairs_inv", + "smooth_slope": "hf_pyramid_slope", + "rough_slope": "random_rough", + "boxes": "boxes", + }, + isaaclab_param={ + "stairs_up": "step_height_range", + "stairs_down": "step_height_range", + "smooth_slope": "slope_range", + "rough_slope": "noise_range", + "boxes": "grid_height_range", + }, + feature_scaling={ + "platform_width": (0.2, 0.4), + "border_width": (0.05, 0.1), + "step_width": (0.125, 0.25), + "grid_width": (0.15, 0.3), + }, + # 上游 ROUGH_TERRAINS_CFG 用 noise_step=0.02；旧 eval.py 把它改成 0.01 后又叠加 + # noise_range 上限 0.40 m，相当于 1 cm 水平网格上放 40 cm 起伏 —— 极端尖锐的尖刺， + # 轮半径只有 0.096 m 的机器人基本无法通过。这里回到上游的 0.02。 + isaaclab_extra_params={ + "rough_slope": {"noise_step": 0.02}, + }, + grid_rows=5, + grid_cols=5, +) + +TERRAIN_SPECS: dict[str, TerrainSpec] = {STANDARD_V1.name: STANDARD_V1} + + +def get_terrain_spec(name: str) -> TerrainSpec: + if name not in TERRAIN_SPECS: + raise KeyError(f"未知地形规格 {name!r}；可选：{sorted(TERRAIN_SPECS)}") + return TERRAIN_SPECS[name] diff --git a/eval_kit/isaaclab/__init__.py b/eval_kit/isaaclab/__init__.py new file mode 100644 index 0000000..6604867 --- /dev/null +++ b/eval_kit/isaaclab/__init__.py @@ -0,0 +1,18 @@ +# Copyright (c) 2025 Deep Robotics +# SPDX-License-Identifier: BSD 3-Clause + +"""IsaacLab 适配层：唯一接触仿真环境内部的地方。 + +这一层**依赖** isaaclab / isaacsim，因此必须在 ``AppLauncher`` 启动之后才能 import。 +上面的 :mod:`eval_kit.core` 与 :mod:`eval_kit.profiles` 不依赖它，可以单独使用。 + +职责边界： + +- :mod:`eval_kit.isaaclab.env_setup` —— 把协议与地形规格翻译成 env_cfg 的改动 +- :mod:`eval_kit.isaaclab.recorder` —— 跑 rollout 并采集成 :class:`~eval_kit.core.schema.Trace` +- :mod:`eval_kit.isaaclab.policy` —— checkpoint 解析与策略加载 + +这里没有任何指标计算逻辑，指标全在内核里。 +""" + +__all__ = ["env_setup", "policy", "recorder"] diff --git a/eval_kit/isaaclab/env_setup.py b/eval_kit/isaaclab/env_setup.py new file mode 100644 index 0000000..26fee2c --- /dev/null +++ b/eval_kit/isaaclab/env_setup.py @@ -0,0 +1,307 @@ +# Copyright (c) 2025 Deep Robotics +# SPDX-License-Identifier: BSD 3-Clause + +"""把评估协议与地形规格翻译成 IsaacLab 的 ``env_cfg`` 改动。 + +这些函数对应旧 ``eval.py`` 里的 ``apply_eval_terrain_baseline`` / +``configure_eval_env`` / ``configure_upper_limit_env`` / ``assign_upper_limit_origins``， +区别是： + +- 数值全部来自 :class:`~eval_kit.core.terrain_spec.TerrainSpec` 与 + :class:`~eval_kit.core.schema.EvalProtocol`，没有散落的常数 +- 需要关掉的随机化事件分「通用」与「机器人相关」两组，机器人相关的那组用 + ``hasattr`` 守卫，因此换机器人不会因为字段不存在而报错 +""" + +from __future__ import annotations + +import copy +from typing import Any, Sequence + +import torch + +from isaaclab.terrains.config.rough import ROUGH_TERRAINS_CFG + +from eval_kit.core.schema import EvalProtocol +from eval_kit.core.terrain_spec import TerrainSpec + +#: 所有任务都需要的通用随机化项。 +COMMON_RANDOMIZATION_TERMS: tuple[str, ...] = ( + "randomize_rigid_body_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", +) + +#: 火蚁专属的随机化项（机身/膝部与轮胎分开做材质随机化）。 +ROBOT_SPECIFIC_RANDOMIZATION_TERMS: tuple[str, ...] = ( + "randomize_wheel_material", + "randomize_support_material", +) + + +def set_fixed_command(env_cfg: Any, protocol: EvalProtocol) -> None: + """把速度指令钉成一个恒定值，并在整个 episode 内不重采样。""" + command_cfg = env_cfg.commands.base_velocity + command_x, command_y, command_yaw = protocol.command + command_cfg.heading_command = False + command_cfg.resampling_time_range = ( + protocol.steps * protocol.dt + 0.01, + protocol.steps * protocol.dt + 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) + command_cfg.debug_vis = False + + +def apply_eval_terrain_baseline(env_cfg: Any, spec: TerrainSpec, protocol: EvalProtocol) -> None: + """所有任务统一用同一套评估地形，保证不同机器人/不同实验面对的是同一块地。""" + 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) + + generator = env_cfg.scene.terrain.terrain_generator + generator.size = (protocol.terrain_size_m, protocol.terrain_size_m) + generator.border_width = 2.0 + generator.curriculum = False + generator.num_rows = spec.grid_rows + generator.num_cols = spec.grid_cols + generator.use_cache = False + env_cfg.scene.terrain.max_init_terrain_level = None + + # 瓦片变小的时候，台阶宽度/平台宽度等特征也要跟着收，否则一块地里塞不下 + for sub_cfg in generator.sub_terrains.values(): + for field_name in spec.feature_scaling: + if hasattr(sub_cfg, field_name): + setattr(sub_cfg, field_name, spec.scaling_for(field_name, protocol.terrain_size_m)) + + +def _build_push_event(protocol: EvalProtocol): + """构造周期性扰动事件。""" + from isaaclab.managers import EventTermCfg as EventTerm + + try: + # 与旧 eval.py 同源，保证扰动行为完全一致 + import rl_training.tasks.manager_based.locomotion.velocity.mdp as mdp + except ImportError: + from isaaclab_tasks.manager_based.locomotion.velocity import mdp + + return EventTerm( + func=mdp.push_by_setting_velocity, + mode="interval", + interval_range_s=(protocol.push_interval_s[0], protocol.push_interval_s[1]), + params={ + "velocity_range": { + "x": protocol.push_velocity_range, + "y": protocol.push_velocity_range, + } + }, + ) + + +def _configure_viewer(env_cfg: Any, viewer_env_index: int | None, free_camera: bool) -> None: + """设置视口相机：跟随某个机器人，或交给客户端自己控制。 + + ``ViewportCameraController`` 只在 ``origin_type`` 是 ``asset_root`` / + ``asset_body`` 时每帧改写相机位姿；设成 ``world`` 只在初始化时写一次， + 之后客户端的操作不会被抢回去 —— 这就是「自由视角」。起始视角由子进程按 + 实际地形范围另设（见 ``eval_suite._set_free_camera_overview``）。 + """ + viewer = getattr(env_cfg, "viewer", None) + if viewer is None: + return + if free_camera: + viewer.origin_type = "world" + # 仅作初始化时的临时值，正式视角在环境建好后按地形范围重设 + viewer.eye = (0.0, -20.0, 15.0) + viewer.lookat = (0.0, 0.0, 0.0) + return + if viewer_env_index is None: + return + viewer.origin_type = "asset_root" + viewer.asset_name = "robot" + viewer.env_index = int(viewer_env_index) + # 相对机器人的偏移：斜前方 2.5 m、抬高 1.6 m，看向机身略前方 + viewer.eye = (2.5, -2.5, 1.6) + viewer.lookat = (0.4, 0.0, 0.1) + + +def configure_eval_env( + env_cfg: Any, + spec: TerrainSpec, + protocol: EvalProtocol, + *, + num_envs: int | None = None, + device: str | None = None, + seed: int | None = None, + viewer_env_index: int | None = None, + free_camera: bool = False, +) -> None: + """把训练用配置改造成单次评估用配置。 + + 关掉观测噪声、课程、随机化，把重置位姿钉死，指令固定，地形统一。 + + ``viewer_env_index`` 不为 None 时，把视口相机改成跟随该环境的机器人； + ``free_camera`` 为 True 时改成不跟随任何物体、由客户端自己移动（两者互斥， + 前者优先）。这两种都只在有画面时（本地 GUI 或推流）才有意义 —— headless + 且不推流时 IsaacLab 根本不创建相机控制器，设了也无害。 + """ + env_cfg.scene.num_envs = num_envs if num_envs is not None else env_cfg.scene.num_envs + env_cfg.episode_length_s = protocol.steps * protocol.dt + if seed is not None: + env_cfg.seed = seed + if device is not None: + env_cfg.sim.device = device + + _configure_viewer(env_cfg, viewer_env_index, free_camera) + + env_cfg.observations.policy.enable_corruption = False + set_fixed_command(env_cfg, protocol) + apply_eval_terrain_baseline(env_cfg, spec, protocol) + + # 课程必须关掉，否则难度/指令会在评估过程中自己变 + for term in ("terrain_levels", "command_levels"): + if getattr(env_cfg.curriculum, term, None) is not None: + setattr(env_cfg.curriculum, term, None) + + for term in (*COMMON_RANDOMIZATION_TERMS, *ROBOT_SPECIFIC_RANDOMIZATION_TERMS): + if hasattr(env_cfg.events, term): + setattr(env_cfg.events, term, None) + + # 重置位姿钉死：同一批机器人从同一个位置、同一姿态、零速度出发 + reset_base = getattr(env_cfg.events, "randomize_reset_base", None) + if reset_base is not None: + start_x, start_y, start_yaw = protocol.start_pose + # z 是相对地形瓦片中心高度的抬升。默认 0.0 与旧 eval.py 一致，但只对 + # 「瓦片中心最高」的地形成立；倒金字塔台阶这类中心最低的地形需要正值， + # 否则机器人在 |x|>platform 处会被生成到地形几何体内部而卡住。 + reset_base.params = { + "pose_range": { + "x": (start_x, start_x), + "y": (start_y, start_y), + "z": (protocol.start_z_offset_m, protocol.start_z_offset_m), + "roll": (0.0, 0.0), + "pitch": (0.0, 0.0), + "yaw": (start_yaw, start_yaw), + }, + "velocity_range": {axis: (0.0, 0.0) for axis in ("x", "y", "z", "roll", "pitch", "yaw")}, + } + + if hasattr(env_cfg.events, "randomize_push_robot"): + env_cfg.events.randomize_push_robot = _build_push_event(protocol) if protocol.enable_interval_pushes else None + + # 越界终止项：它比较的是世界坐标与整张地图的范围，不是机器人自己那块地砖， + # 小地砖 + 大网格时会把最外侧格子判成越界（详见 EvalProtocol 的说明）。 + if not protocol.enforce_terrain_bounds and getattr(env_cfg.terminations, "terrain_out_of_bounds", None) is not None: + env_cfg.terminations.terrain_out_of_bounds = None + + +# 说明：判定地砖是否够大的逻辑已移到 eval_kit.core.schema.EvalProtocol. +# minimum_terrain_size_for_bounds()。它只用到协议自己的 start_pose，放在适配层会 +# 让不启动仿真的编排进程被迫 import isaaclab（进而因缺 pxr 而失败）。 + + +def configure_upper_limit_env( + env_cfg: Any, + spec: TerrainSpec, + protocol: EvalProtocol, + terrain_types: Sequence[str], + *, + seed: int | None = None, + device: str | None = None, + viewer_env_index: int | None = None, + free_camera: bool = False, +) -> None: + """配置上限评估的固定网格：行 = 难度等级，列 = 地形类型。 + + ``seed`` 同时决定地形生成与重置采样。显式传入是为了让地形本身在不同种子间 + 保持一致（同一块地），而机器人初始状态可以随种子变化。 + + ``viewer_env_index`` / ``free_camera`` 透传给 :func:`configure_eval_env`， + 用于让视口相机跟随某个环境的机器人，或者干脆交给客户端自由控制。 + """ + configure_eval_env( + env_cfg, spec, protocol, seed=seed, device=device, + viewer_env_index=viewer_env_index, free_camera=free_camera, + ) + + generator = env_cfg.scene.terrain.terrain_generator + generator.curriculum = True + generator.num_rows = protocol.num_levels + generator.num_cols = len(terrain_types) + generator.difficulty_range = (0.0, 1.0) + if hasattr(generator, "seed"): + generator.seed = int(seed) if seed is not None else int(env_cfg.seed or 0) + + # 每个格子里放 envs_per_cell 个机器人 + env_cfg.scene.num_envs = protocol.num_levels * len(terrain_types) * protocol.envs_per_cell + env_cfg.scene.terrain.max_init_terrain_level = protocol.num_levels - 1 + + # 只保留这次评估要用到的子地形，且按被选中的次数分配占比 + isaaclab_names = [spec.isaaclab_subterrain[terrain_type] for terrain_type in terrain_types] + original = generator.sub_terrains + generator.sub_terrains = {name: copy.deepcopy(original[name]) for name in isaaclab_names} + for name, sub_cfg in generator.sub_terrains.items(): + sub_cfg.proportion = isaaclab_names.count(name) / len(isaaclab_names) + + # 把等级表里的首尾值写进子地形的难度区间。 + # 用 _set_sub_terrain_field 而不是裸 setattr：字段名写错要立刻报错， + # 静默造一个假属性会让地形保持默认难度，而结果看起来「只是很难」。 + for terrain_type in terrain_types: + sub_cfg = generator.sub_terrains[spec.isaaclab_subterrain[terrain_type]] + levels = spec.levels_for(terrain_type, protocol.num_levels) + _set_sub_terrain_field(sub_cfg, spec.isaaclab_param[terrain_type], (levels[0], levels[-1]), terrain_type) + + # 与等级无关的固定字段（例如粗糙地形的噪声步长） + for terrain_type, extra in spec.isaaclab_extra_params.items(): + if terrain_type not in terrain_types: + continue + sub_cfg = generator.sub_terrains[spec.isaaclab_subterrain[terrain_type]] + for field_name, value in extra.items(): + _set_sub_terrain_field(sub_cfg, field_name, value, terrain_type) + + +def _set_sub_terrain_field(sub_cfg: Any, field_name: str, value: Any, terrain_type: str) -> None: + """写入子地形字段，字段不存在时报错而不是静默创建一个假属性。""" + if not hasattr(sub_cfg, field_name): + raise AttributeError( + f"地形规格为 {terrain_type!r} 指定的字段 {field_name!r} 在 " + f"{type(sub_cfg).__name__} 上不存在。裸 setattr 会静默造一个假属性、" + f"使地形保持默认难度，而那看起来只是「很难」。可用字段：" + f"{sorted(k for k in vars(sub_cfg) if not k.startswith('_'))}" + ) + setattr(sub_cfg, field_name, value) + + +def assign_upper_limit_origins( + env: Any, + spec: TerrainSpec, + protocol: EvalProtocol, + terrain_types: Sequence[str], +) -> tuple[torch.Tensor, torch.Tensor]: + """给每个格子平均分配机器人，并直接把地形等级/类型/起点写进地形对象。 + + 返回 ``(level_ids, type_ids)``，用于事后把逐环境结果折叠成矩阵。 + """ + base_env = env.unwrapped + terrain = base_env.scene.terrain + device = base_env.device + + levels: list[int] = [] + types: list[int] = [] + for type_index in range(len(terrain_types)): + for level_index in range(protocol.num_levels): + levels.extend([level_index] * protocol.envs_per_cell) + types.extend([type_index] * protocol.envs_per_cell) + + level_ids = torch.tensor(levels, dtype=torch.long, device=device) + type_ids = torch.tensor(types, dtype=torch.long, device=device) + terrain.terrain_levels = level_ids + terrain.terrain_types = type_ids + terrain.env_origins = terrain.terrain_origins[level_ids, type_ids] + env.reset() + return level_ids, type_ids diff --git a/eval_kit/isaaclab/policy.py b/eval_kit/isaaclab/policy.py new file mode 100644 index 0000000..2fb468b --- /dev/null +++ b/eval_kit/isaaclab/policy.py @@ -0,0 +1,71 @@ +# Copyright (c) 2025 Deep Robotics +# SPDX-License-Identifier: BSD 3-Clause + +"""checkpoint 解析与策略加载（IsaacLab + RSL-RL）。""" + +from __future__ import annotations + +import os +from typing import Any + +import gymnasium as gym + +from isaaclab.envs import ( + DirectMARLEnv, + DirectMARLEnvCfg, + DirectRLEnvCfg, + ManagerBasedRLEnvCfg, + multi_agent_to_single_agent, +) +from isaaclab_tasks.utils import get_checkpoint_path +from isaaclab_rl.rsl_rl import RslRlVecEnvWrapper + + +def resolve_checkpoint( + log_root_path: str, + load_run: str | None = None, + checkpoint: str | None = None, +) -> str: + """解析出 checkpoint 的绝对路径。 + + 允许 ``load_run`` / ``checkpoint`` 传绝对路径或带目录的相对路径，从而跨实验目录 + 取权重（例如从 Stage 1 的目录里取权重来评估），否则回退到 IsaacLab 的按名查找。 + """ + 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"找不到 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"找不到运行目录：{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 make_env( + task: str, + env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, + clip_actions: float | None = None, +): + """按已配置好的 ``env_cfg`` 建环境，并包成 RSL-RL 的向量环境。""" + env = gym.make(task, cfg=env_cfg) + if isinstance(env.unwrapped, DirectMARLEnv): + env = multi_agent_to_single_agent(env) + return RslRlVecEnvWrapper(env, clip_actions=clip_actions) + + +def load_policy(env, agent_cfg: Any, resume_path: str, device: str | None = None): + """加载 checkpoint 并返回推理策略。""" + from rsl_rl.runners import OnPolicyRunner + + runner = OnPolicyRunner(env, agent_cfg.to_dict(), log_dir=None, device=device or agent_cfg.device) + runner.load(resume_path) + return runner.get_inference_policy(device=device or env.unwrapped.device) diff --git a/eval_kit/isaaclab/recorder.py b/eval_kit/isaaclab/recorder.py new file mode 100644 index 0000000..341cea1 --- /dev/null +++ b/eval_kit/isaaclab/recorder.py @@ -0,0 +1,257 @@ +# Copyright (c) 2025 Deep Robotics +# SPDX-License-Identifier: BSD 3-Clause + +"""跑 rollout 并把仿真信号采集成框架无关的 :class:`~eval_kit.core.schema.Trace`。 + +这是唯一读取 ``env`` 内部状态的地方。采集完之后环境就可以关掉了，后续所有 +指标计算都在内核里对纯数组进行。 + +形态相关的下标（谁是轮、谁是支撑体、哪些关节耗电）全部由 +:class:`~eval_kit.core.schema.RobotProfile` 的正则解析得到，代码里没有硬编码的 +link/joint 名字。解析不到时记录警告并跳过对应指标，而不是抛异常。 +""" + +from __future__ import annotations + +import time +from dataclasses import dataclass, field +from typing import Any + +import numpy as np +import torch + +from eval_kit.core.naming import resolve_matching_indices +from eval_kit.core.schema import EvalProtocol, RobotProfile, Trace + +#: 默认的接触力传感器名字，与 ``velocity_env_cfg.py`` 的场景定义一致。 +DEFAULT_CONTACT_SENSOR = "contact_forces" + + +@dataclass +class RolloutRecorder: + """一次评估里固定不变的东西：形态下标与质量。建一次，跑多次。""" + + feet_ids: list[int] = field(default_factory=list) + feet_names: list[str] = field(default_factory=list) + support_ids: list[int] = field(default_factory=list) + support_names: list[str] = field(default_factory=list) + joint_ids: list[int] = field(default_factory=list) + joint_names: list[str] = field(default_factory=list) + mass_kg: float | None = None + contact_sensor_name: str = DEFAULT_CONTACT_SENSOR + warnings: list[str] = field(default_factory=list) + + def describe(self) -> str: + lines = [ + f"足端 {len(self.feet_ids)} 个: {self.feet_names or '（无）'}", + f"支撑体 {len(self.support_ids)} 个: {self.support_names or '（无）'}", + f"驱动关节 {len(self.joint_ids)} 个: {self.joint_names or '（无）'}", + f"整机质量: {'未知' if self.mass_kg is None else f'{self.mass_kg:.3f} kg'}", + ] + for warning in self.warnings: + lines.append(f"[警告] {warning}") + return "\n".join(lines) + + +def read_total_mass(robot: Any) -> float | None: + """从仿真资产读整机质量。 + + 优先读 PhysX 运行时的真实质量（会反映任何质量随机化），读不到再退回资产默认质量。 + """ + try: + masses = robot.root_physx_view.get_masses() + return float(masses[0].sum().item()) + except Exception: # noqa: BLE001 - 不同版本 API 差异，读不到就退回 + pass + try: + return float(robot.data.default_mass[0].sum().item()) + except Exception: # noqa: BLE001 + return None + + +def build_recorder( + env: Any, + profile: RobotProfile, + *, + contact_sensor_name: str = DEFAULT_CONTACT_SENSOR, +) -> RolloutRecorder: + """按机器人描述解析出足端/支撑体/关节下标，并读一次质量。""" + base_env = env.unwrapped + robot = base_env.scene["robot"] + recorder = RolloutRecorder(contact_sensor_name=contact_sensor_name) + + feet = resolve_matching_indices(profile.feet_body_patterns, robot.body_names) + recorder.feet_ids = feet.indices + recorder.feet_names = feet.names + for pattern in feet.unmatched_patterns: + recorder.warnings.append( + f"机器人 {profile.name!r} 的足端正则 {pattern!r} 在 body_names 里没有命中，" + "支撑裕度指标将被跳过。" + ) + + robot_bodies = set(robot.body_names) + sensor = base_env.scene.sensors.get(contact_sensor_name) if base_env.scene.sensors else None + if sensor is None: + recorder.warnings.append( + f"场景里没有接触力传感器 {contact_sensor_name!r}，支撑接触类指标将被跳过。" + ) + else: + support = resolve_matching_indices(profile.support_body_patterns, sensor.body_names) + recorder.support_ids = support.indices + recorder.support_names = support.names + missing = [name for name in support.names if name not in robot_bodies] + if missing: + recorder.warnings.append( + f"接触力传感器上的 {missing} 不在机器人的 body 列表里，指标可能对不上。" + ) + for pattern in support.unmatched_patterns: + if pattern in profile.support_body_patterns: + recorder.warnings.append( + f"机器人 {profile.name!r} 的支撑体正则 {pattern!r} 没有命中，" + "支撑接触类指标将被跳过。" + ) + + joints = resolve_matching_indices(profile.actuated_joint_patterns, robot.joint_names) + recorder.joint_ids = joints.indices + recorder.joint_names = joints.names + if not joints.indices: + recorder.warnings.append( + "驱动关节正则一个都没命中，能耗与 COT 将无法计算（会返回 nan）。" + ) + + recorder.mass_kg = read_total_mass(robot) + if recorder.mass_kg is None and profile.mass_kg is None: + recorder.warnings.append("既读不到资产质量、描述里也没写兜底值，COT 将返回 nan。") + + return recorder + + +class _FrameCollector: + """逐帧收集张量，最后一次性转成 numpy——每帧都转会明显变慢。""" + + def __init__(self) -> None: + self.frames: dict[str, list[np.ndarray]] = {} + + def add(self, **signals: np.ndarray) -> None: + for key, value in signals.items(): + self.frames.setdefault(key, []).append(value) + + def stack(self) -> dict[str, np.ndarray]: + return {key: np.stack(values, axis=0) for key, values in self.frames.items()} + + +def _capture(robot: Any, sensor: Any, recorder: RolloutRecorder, num_envs: int) -> dict[str, np.ndarray]: + """读一帧信号。所有张量都要 clone，否则会被后续仿真步就地覆写。""" + + def to_numpy(tensor: torch.Tensor) -> np.ndarray: + return tensor.detach().to("cpu", torch.float64).numpy() + + signals = { + "base_pos_w": to_numpy(robot.data.root_pos_w.clone()), + "base_lin_vel_w": to_numpy(robot.data.root_lin_vel_w.clone()), + "base_lin_vel_b": to_numpy(robot.data.root_lin_vel_b.clone()), + "base_ang_vel_b": to_numpy(robot.data.root_ang_vel_b.clone()), + "projected_gravity_b": to_numpy(robot.data.projected_gravity_b.clone()), + "joint_torque": to_numpy(robot.data.applied_torque[:, recorder.joint_ids].clone()) + if recorder.joint_ids + else np.zeros((num_envs, 0)), + "joint_vel": to_numpy(robot.data.joint_vel[:, recorder.joint_ids].clone()) + if recorder.joint_ids + else np.zeros((num_envs, 0)), + "feet_pos_w": to_numpy(robot.data.body_pos_w[:, recorder.feet_ids, :2].clone()) + if recorder.feet_ids + else np.zeros((num_envs, 0, 2)), + } + + if sensor is not None and recorder.support_ids: + forces = sensor.data.net_forces_w + contact_time = sensor.data.current_contact_time + if forces is None: + # 传感器还没更新过（第一帧），按无接触处理 + signals["support_force_n"] = np.zeros((num_envs, len(recorder.support_ids))) + signals["support_contact_time_s"] = np.zeros((num_envs, len(recorder.support_ids))) + else: + selected = forces[:, recorder.support_ids, :] + signals["support_force_n"] = to_numpy(torch.linalg.norm(selected, dim=-1).clone()) + signals["support_contact_time_s"] = to_numpy( + (contact_time[:, recorder.support_ids].clone() if contact_time is not None else torch.zeros_like(selected[..., 0])) + ) + else: + signals["support_force_n"] = np.zeros((num_envs, 0)) + signals["support_contact_time_s"] = np.zeros((num_envs, 0)) + + return signals + + +def record_rollout( + env: Any, + policy: Any, + protocol: EvalProtocol, + profile: RobotProfile, + recorder: RolloutRecorder | None = None, +) -> Trace: + """跑 ``protocol.steps`` 步，采集一条 :class:`Trace`。 + + 第 0 帧记录动作执行之前的起始状态，之后每步记录一次动作之后的状态， + 因此帧数比步数多一。 + """ + base_env = env.unwrapped + robot = base_env.scene["robot"] + if recorder is None: + recorder = build_recorder(env, profile) + sensors = base_env.scene.sensors + sensor = sensors.get(recorder.contact_sensor_name) if sensors else None + + num_envs = int(base_env.num_envs) + collector = _FrameCollector() + initial = _capture(robot, sensor, recorder, num_envs) + # 起始帧还没有任何 episode 结束，置零以保证与后续帧形状一致 + initial["done"] = np.zeros(num_envs, dtype=np.float64) + collector.add(**initial) + + obs = env.get_observations() + # 终端侧进度：可视化时推流画面可能因为渲染/编码出问题而冻住，这时只有终端能 + # 说明仿真到底有没有在推进（推进则打印继续，冻结则这里也停住）。每 10% 打一行。 + progress_every = max(1, protocol.steps // 10) + started = time.monotonic() + for step in range(protocol.steps): + with torch.inference_mode(): + actions = policy(obs) + if isinstance(actions, tuple): + actions = actions[0] + obs, _, dones, _ = env.step(actions) + frame = _capture(robot, sensor, recorder, num_envs) + frame["done"] = dones.detach().to("cpu", torch.float64).numpy().reshape(num_envs) + collector.add(**frame) + if (step + 1) % progress_every == 0 or step + 1 == protocol.steps: + print( + f"[eval_kit] rollout {step + 1}/{protocol.steps} 步，用时 {time.monotonic() - started:.1f} s", + flush=True, + ) + + arrays = collector.stack() + return Trace( + dt=float(base_env.step_dt), + command=np.asarray(protocol.command, dtype=np.float64), + mass_kg=recorder.mass_kg, + **arrays, + ) + + +def summarize_early_resets(trace: Trace) -> dict[str, Any]: + """统计 episode 中途被重置的环境数量。 + + 中途重置会让后续帧来自新 episode，路程/进度类指标会被污染。旧 + ``eval.py`` 没有检查这一点，所以出现过「指标看着正常其实混了两次 rollout」 + 的情况不容易发现。 + """ + # 最后一帧的 done 属于 episode 正常结束，不算早退 + early = trace.done[:-1] > 0.5 + per_env = early.sum(axis=0) + affected = int(np.count_nonzero(per_env > 0)) + return { + "envs_with_early_reset": affected, + "total_early_resets": int(per_env.sum()), + "max_resets_single_env": int(per_env.max()) if per_env.size else 0, + "polluted": affected > 0, + } diff --git a/eval_kit/profiles.py b/eval_kit/profiles.py new file mode 100644 index 0000000..9201417 --- /dev/null +++ b/eval_kit/profiles.py @@ -0,0 +1,56 @@ +# Copyright (c) 2025 Deep Robotics +# SPDX-License-Identifier: BSD 3-Clause + +"""机器人形态注册表。 + +**加一个机器人 = 在这里加一条记录**，评价代码不动。 + +数值来源：火蚁的 link/joint 取自 +``source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/config/wheeled/deeprobotics_m20/rough_env_cfg.py``； +机身尺寸取自 ``docs/solenopsis_urdf_comparison.md`` 里记录的碰撞体尺寸。 +""" + +from __future__ import annotations + +from eval_kit.core.schema import RobotProfile + +SOLENOPSIS = RobotProfile( + name="solenopsis", + # None = 运行时从仿真资产读真实质量（Solenopsis.urdf 当前为 24.74 kg）。 + # 旧 eval.py 的命令行默认值是 25.482，那是更早的 M20 简化模型名义质量， + # 会让 COT 偏低约 4%。 + mass_kg=None, + # 四个轮就是支撑多边形的接触点。 + feet_body_patterns=(".*_wheel",), + # 机身与膝部允许作为越障时的短暂支撑。 + support_body_patterns=("base_link", ".*_knee"), + # 火蚁 8 个关节全部主动驱动；分两组写便于将来出现被动关节时收窄范围。 + actuated_joint_patterns=(".*knee_joint", ".*wheel_joint"), + support_force_threshold_n=15.0, + # 机身碰撞体 0.706 x 0.310 x 0.1714 m，取长边。 + body_length_m=0.706, +) + +#: 四足对照机，仅用于验证「换机器人只改数据」这一点，尚未在仿真里跑过。 +LITE3 = RobotProfile( + name="lite3", + mass_kg=None, + feet_body_patterns=(".*_FOOT",), + # 四足没有「机身支撑」这个语义，留空 → 支撑类指标返回 nan 而不是 0。 + support_body_patterns=(), + actuated_joint_patterns=(".*_joint",), + support_force_threshold_n=15.0, + # 未知，仅用于地形瓦片尺寸校验。 + body_length_m=0.0, +) + +PROFILES: dict[str, RobotProfile] = { + SOLENOPSIS.name: SOLENOPSIS, + LITE3.name: LITE3, +} + + +def get_profile(name: str) -> RobotProfile: + if name not in PROFILES: + raise KeyError(f"未知机器人描述 {name!r}；可选：{sorted(PROFILES)}") + return PROFILES[name] diff --git a/eval_kit/protocols.py b/eval_kit/protocols.py new file mode 100644 index 0000000..3ea26c4 --- /dev/null +++ b/eval_kit/protocols.py @@ -0,0 +1,102 @@ +# Copyright (c) 2025 Deep Robotics +# SPDX-License-Identifier: BSD 3-Clause + +"""评估协议预设。 + +协议就是一次评估的全部「口径」。默认预设 ``standard_v1`` 的数值与旧 +``scripts/reinforcement_learning/rsl_rl/eval.py`` 的命令行默认值一致， +所以两套评估跑出来的数字可以直接对照。 + +要改口径，改预设或另加一个预设，不要在调用处临时传参——预设名会进 +``protocol_fingerprint``，写进输出 JSON。 +""" + +from __future__ import annotations + +from dataclasses import replace + +from eval_kit.core.schema import EvalProtocol + +#: 标准越障评估：10 s、前进 1 m/s、最多 2 m 算成功。与旧 eval.py 默认一致。 +#: +#: **只建议用于与旧 ``eval.py`` 做对照**：它的地砖 2.0 m 会让越界终止项误判最外侧格子， +#: 起始高度 0.0 m 会让机器人在倒金字塔台阶地形里被生成到几何体内部 +#: （``stairs_down`` 因此恒为 0）。要拿可用的数字请用 ``standard_v1_clean``。 +STANDARD_V1 = EvalProtocol( + steps=500, + dt=0.02, + command=(1.0, 0.0, 0.0), + start_pose=(-0.9, 0.0, 0.0), + success_distance_m=2.0, + tilt_threshold_rad=0.70, + acc_threshold_g=2.0 / 9.81, + energy_convention="absolute", + stability_margin_method="foot_bbox", + support_force_threshold_n=15.0, + enable_interval_pushes=False, + terrain_spec="standard_v1", + terrain_size_m=2.0, + num_levels=8, + envs_per_cell=1, + seeds=(0,), +) + +#: 抗扰评估：每 2.01 s 施加一次水平速度扰动，其余同 standard_v1。 +ROBUSTNESS_V1 = replace( + STANDARD_V1, + enable_interval_pushes=True, + push_interval_s=(2.0, 2.01), + push_velocity_range=(-0.5, 0.5), +) + +#: 冒烟测试：步数与网格都缩小，用于确认流水线能跑通，不要用它的数值下结论。 +#: 地形参数沿用 ``standard_v1_clean`` 的修正值，否则它会带上越界误判的假警告。 +SMOKE = replace( + STANDARD_V1, + steps=100, + num_levels=2, + envs_per_cell=1, + terrain_size_m=4.0, + start_z_offset_m=0.5, +) + +#: **推荐用于机器人对比**：修掉两处会让结果失真的问题，其余口径与 ``standard_v1`` 相同， +#: 因此数值仍可与后者对照。 +#: +#: 1. 地砖从 2.0 m 放大到 4.0 m，消除 IsaacLab ``terrain_out_of_bounds`` 的误判 +#: （阈值 3.8 m 由 ``EvalProtocol.minimum_terrain_size_for_bounds()`` 给出）。 +#: 2. 起始高度抬高 0.5 m，让机器人落到地形表面而不是被生成在地形内部。 +#: 实测证据：``--terrain-type stairs_down`` 单独跑，起始高度 0.0 m 时 +#: score=0.00 / 成功率 0.05，抬高后 score 显著恢复。取 0.5 是因为倒金字塔台阶的 +#: 嵌入深度随难度递增（L1 约 0.11 m、L8 约 0.50 m），0.5 m 刚好覆盖最深的一级。 +#: 旧 ``eval.py`` 用 0.0，所以它的倒金字塔台阶那一列恒为 0，是无效的。 +STANDARD_V1_CLEAN = replace(STANDARD_V1, terrain_size_m=4.0, start_z_offset_m=0.5) + +#: 另一种处置方式：保留 2.0 m 地砖，直接去掉越界终止项。 +STANDARD_V1_NO_BOUNDS = replace(STANDARD_V1, enforce_terrain_bounds=False) + +#: 多随机种子的标准评估 —— **对外给结论时用这个**。 +#: +#: 单次 rollout 的噪声足以淹没机器人之间的差异，所以对比必须多种子；同时它基于 +#: ``STANDARD_V1_CLEAN``（地砖 4.0 m、起始高度 +0.5 m），不带已知的 harness 偏差。 +#: +#: 默认 ``envs_per_cell = 1``（每格 1 个机器人），因为起始位姿钉死后同格多个机器人 +#: 起始坐标完全相同、不是独立样本。样本量靠种子数而非每格机器人数量来提升 —— +#: 每个种子会重新生成地形实例，跨种子才是独立的。8 个种子给出每格 1/8 ≈ 12.5% 的 +#: 分辨率；想要更细就显式加 ``--seeds``。 +STANDARD_V1_MULTISEED = replace(STANDARD_V1_CLEAN, seeds=(0, 100, 200, 300, 400, 500, 600, 700)) + +PROTOCOLS: dict[str, EvalProtocol] = { + "standard_v1": STANDARD_V1, + "standard_v1_clean": STANDARD_V1_CLEAN, + "standard_v1_no_bounds": STANDARD_V1_NO_BOUNDS, + "robustness_v1": ROBUSTNESS_V1, + "smoke": SMOKE, + "standard_v1_multiseed": STANDARD_V1_MULTISEED, +} + + +def get_protocol(name: str) -> EvalProtocol: + if name not in PROTOCOLS: + raise KeyError(f"未知评估协议 {name!r}；可选：{sorted(PROTOCOLS)}") + return PROTOCOLS[name] diff --git a/eval_kit/tests/test_cli.py b/eval_kit/tests/test_cli.py new file mode 100644 index 0000000..40ddd0a --- /dev/null +++ b/eval_kit/tests/test_cli.py @@ -0,0 +1,544 @@ +# Copyright (c) 2025 Deep Robotics +# SPDX-License-Identifier: BSD 3-Clause + +"""命令行入口的静态自检：不需要 Isaac Sim，也不需要 GPU。 + +用桩替换 ``isaaclab.app``，从而可以单独加载 ``eval_suite.py``，检查参数解析与 +口径覆盖逻辑。仿真相关的部分（建环境、跑 rollout）不在覆盖范围内——那些必须真跑。 + +运行方式（在仓库根目录）:: + + python3 eval_kit/tests/test_cli.py +""" + +from __future__ import annotations + +import importlib.util +import sys +import types +from pathlib import Path + +if __package__ in (None, ""): + sys.path.insert(0, str(Path(__file__).resolve().parents[2])) + +PROJECT_ROOT = Path(__file__).resolve().parents[2] +EVAL_SUITE_PATH = PROJECT_ROOT / "scripts" / "reinforcement_learning" / "rsl_rl" / "eval_suite.py" + +#: ``AppLauncher.add_app_launcher_args()`` 会添加的参数名，抄自 +#: ``IsaacLab-2.3.2/source/isaaclab/isaaclab/app/app_launcher.py``。 +#: 自己的 parser 里出现任何一个都会让 argparse 抛 +#: ``ArgumentError: conflicting option string``，脚本连启动都做不到。 +#: （只有 device / headless / enable_cameras / experience / livestream / +#: rendering_mode / xr 会被 AppLauncher 自己的显式检查拦下，--verbose / --info / +#: --cpu 等则由 argparse 拦，所以必须把 14 个都列出来。） +APPLAUNCHER_ARG_NAMES = ( + "--headless", + "--livestream", + "--enable_cameras", + "--xr", + "--device", + "--cpu", + "--verbose", + "--info", + "--experience", + "--rendering_mode", + "--kit_args", + "--anim_recording_enabled", + "--anim_recording_start_time", + "--anim_recording_stop_time", +) + + +def _stub_isaaclab() -> None: + """装一个最小的 isaaclab.app 桩，只为让模块级 import 通过。 + + 桩要复刻真实 ``AppLauncher`` 的行为，否则测试放过的正是真跑会崩的地方： + + - ``add_app_launcher_args`` 会添加上面 14 个参数名 + - 如果 parser 里已经有同名参数，它**拒绝启动**而不是静默覆盖 + """ + if "isaaclab.app" in sys.modules: + return + + class AppLauncher: # noqa: D401 - 桩 + def __init__(self, args): + self.app = None + + @staticmethod + def add_app_launcher_args(parser): + existing = {option for action in parser._actions for option in action.option_strings} + conflict = sorted(existing & set(APPLAUNCHER_ARG_NAMES)) + if conflict: + raise ValueError( + f"argument {conflict[0]}: conflicting option string: {conflict[0]}" + ) + parser.add_argument("--headless", action="store_true", default=False) + parser.add_argument("--device", type=str, default=None) + for name in APPLAUNCHER_ARG_NAMES: + dest = name.lstrip("-").replace("-", "_") + if dest in ("headless", "device"): + continue + parser.add_argument(name, dest=dest, nargs="?", const=True, default=None) + return parser + + isaaclab = types.ModuleType("isaaclab") + app = types.ModuleType("isaaclab.app") + app.AppLauncher = AppLauncher # type: ignore[attr-defined] + isaaclab.app = app # type: ignore[attr-defined] + sys.modules["isaaclab"] = isaaclab + sys.modules["isaaclab.app"] = app + + +def _load_eval_suite(): + _stub_isaaclab() + specification = importlib.util.spec_from_file_location("eval_suite_under_test", EVAL_SUITE_PATH) + if specification is None or specification.loader is None: + raise RuntimeError(f"无法加载 {EVAL_SUITE_PATH}") + module = importlib.util.module_from_spec(specification) + specification.loader.exec_module(module) + return module + + +def test_parser_defaults(): + """默认值：协议预设名、机器人描述名、评估类型都在。""" + module = _load_eval_suite() + args = module._build_parser().parse_args([]) + + assert args.task == "Rough-Deeprobotics-M20-v0", args.task + assert args.profile == "solenopsis" + assert args.protocol == "standard_v1" + assert args.eval_type == "all" + assert args.num_envs == 250 + assert args.terrain_type == "all" + # 口径覆盖项默认必须为 None，否则会静默覆盖协议预设 + for name in ( + "episode_length_s", "command_x", "command_y", "command_yaw", "start_x", "start_y", + "terrain_size", "success_distance", "tilt_threshold", "acc_threshold_g", + "energy_convention", "robot_mass", "num_levels", "envs_per_cell", "seeds", + "enforce_terrain_bounds", "start_z_offset", + ): + assert getattr(args, name) is None, f"{name} 默认应为 None，实际 {getattr(args, name)!r}" + return "默认值正确，口径覆盖项默认全部为 None" + + +def test_no_terrain_bounds_flag(): + """--no-terrain-bounds 关闭越界终止项，并改变指纹。""" + module = _load_eval_suite() + from eval_kit.protocols import get_protocol + + preset = get_protocol("standard_v1") + assert preset.enforce_terrain_bounds is True, "standard_v1 应保持与 eval.py 一致（保留终止项）" + + args = module._build_parser().parse_args(["--no-terrain-bounds"]) + protocol = module._apply_overrides(preset, args) + assert protocol.enforce_terrain_bounds is False + assert protocol.fingerprint() != preset.fingerprint(), "关闭终止项必须改变指纹" + + # 不带该开关时不覆盖 + unchanged = module._apply_overrides(preset, module._build_parser().parse_args([])) + assert unchanged.enforce_terrain_bounds is True + return "开关生效且反映在指纹里" + + +def test_clean_protocol_uses_big_tile(): + """standard_v1_clean 通过放大地砖来消除越界误判，而不是关掉终止项。""" + from eval_kit.protocols import get_protocol + + clean = get_protocol("standard_v1_clean") + parity = get_protocol("standard_v1") + no_bounds = get_protocol("standard_v1_no_bounds") + + assert clean.terrain_size_m >= 3.8, f"地砖 {clean.terrain_size_m} 不足以消除误判" + assert clean.enforce_terrain_bounds is True, "clean 方案应保留终止项" + assert no_bounds.enforce_terrain_bounds is False + # 除地砖外其余口径与 parity 方案一致，数值仍可对照 + assert clean.command == parity.command + assert clean.success_distance_m == parity.success_distance_m + assert clean.num_levels == parity.num_levels + return f"clean 地砖 {clean.terrain_size_m} m，保留终止项；no_bounds 关闭终止项" + + +def test_overrides_are_opt_in(): + """不传任何覆盖项时，协议必须与预设逐字段相同。""" + module = _load_eval_suite() + from eval_kit.protocols import get_protocol + + args = module._build_parser().parse_args([]) + preset = get_protocol("standard_v1") + protocol = module._apply_overrides(preset, args) + + assert protocol == preset, "未传覆盖项却改变了协议" + assert protocol.fingerprint() == preset.fingerprint() + return "未传覆盖 → 协议与预设完全相同（含 fingerprint）" + + +def test_overrides_apply_and_update_fingerprint(): + """显式覆盖项生效，且指纹随之改变。""" + module = _load_eval_suite() + from eval_kit.protocols import get_protocol + + args = module._build_parser().parse_args( + [ + "--episode-length-s", "20", + "--command-x", "0.8", + "--command-yaw", "-0.3", + "--success-distance", "3.5", + "--tilt-threshold", "0.5", + "--acc-threshold-g", "0.3", + "--energy-convention", "signed", + "--terrain-size", "1.5", + "--num-levels", "4", + "--envs-per-cell", "2", + "--seeds", "7", "8", + ] + ) + preset = get_protocol("standard_v1") + protocol = module._apply_overrides(preset, args) + + assert protocol.steps == 1000, protocol.steps # 20 s / 0.02 s + assert protocol.command == (0.8, 0.0, -0.3), protocol.command + assert protocol.success_distance_m == 3.5 + assert protocol.tilt_threshold_rad == 0.5 + assert abs(protocol.acc_threshold_g - 0.3) < 1e-12 + assert protocol.energy_convention == "signed" + assert protocol.terrain_size_m == 1.5 + assert protocol.num_levels == 4 + assert protocol.envs_per_cell == 2 + assert protocol.seeds == (7, 8) + assert protocol.fingerprint() != preset.fingerprint() + # 没覆盖的字段保持预设 + assert protocol.start_pose == preset.start_pose + assert protocol.dt == preset.dt + return "11 项覆盖全部生效，未覆盖字段保持预设" + + +def test_command_overrides_compose_independently(): + """单独覆盖某一个速度分量时，另两个分量保持预设。""" + module = _load_eval_suite() + from eval_kit.protocols import get_protocol + + preset = get_protocol("standard_v1") # command = (1.0, 0.0, 0.0) + + only_y = module._apply_overrides(preset, module._build_parser().parse_args(["--command-y", "0.5"])) + assert only_y.command == (1.0, 0.5, 0.0), only_y.command + + only_yaw = module._apply_overrides(preset, module._build_parser().parse_args(["--command-yaw", "1.2"])) + assert only_yaw.command == (1.0, 0.0, 1.2), only_yaw.command + + both = module._apply_overrides( + preset, module._build_parser().parse_args(["--command-y", "0.5", "--command-yaw", "1.2"]) + ) + assert both.command == (1.0, 0.5, 1.2), both.command + + only_x = module._apply_overrides(preset, module._build_parser().parse_args(["--command-x", "2.0"])) + assert only_x.command == (2.0, 0.0, 0.0), only_x.command + return "command 三个分量可独立覆盖" + + +def test_start_pose_overrides_compose(): + """起始位姿同理。""" + module = _load_eval_suite() + from eval_kit.protocols import get_protocol + + preset = get_protocol("standard_v1") + protocol = module._apply_overrides( + preset, module._build_parser().parse_args(["--start-x", "-1.2", "--start-y", "0.4"]) + ) + assert protocol.start_pose == (-1.2, 0.4, 0.0), protocol.start_pose + return f"start_pose = {protocol.start_pose}" + + +def test_start_z_offset_override(): + """起始高度偏移：默认 0.0（与旧 eval.py 一致），显式给出时生效并改变指纹。""" + module = _load_eval_suite() + from eval_kit.protocols import get_protocol + + preset = get_protocol("standard_v1") + assert preset.start_z_offset_m == 0.0, "默认必须为 0.0，否则与 eval.py 不可比" + + parser = module._build_parser() + unchanged = module._apply_overrides(preset, parser.parse_args([])) + assert unchanged.start_z_offset_m == 0.0 + assert unchanged.fingerprint() == preset.fingerprint() + + lifted = module._apply_overrides(preset, parser.parse_args(["--start-z-offset", "0.3"])) + assert lifted.start_z_offset_m == 0.3 + assert lifted.fingerprint() != preset.fingerprint(), "抬高起始高度必须改变指纹" + # 只动高度，不动水平位姿 + assert lifted.start_pose == preset.start_pose + return "默认 0.0，--start-z-offset 生效且反映在指纹里" + + +def test_terrain_types_resolution(): + """地形类型：all 展开成规格里的全部类型，单个类型要校验存在性。""" + module = _load_eval_suite() + from eval_kit.core.terrain_spec import get_terrain_spec + + spec = get_terrain_spec("standard_v1") + + all_args = module._build_parser().parse_args([]) + assert module._terrain_types_for(all_args, spec) == list(spec.types) + + single_args = module._build_parser().parse_args(["--terrain-type", "boxes"]) + assert module._terrain_types_for(single_args, spec) == ["boxes"] + + bad_args = module._build_parser().parse_args(["--terrain-type", "nonexistent"]) + try: + module._terrain_types_for(bad_args, spec) + except SystemExit as exc: + assert "nonexistent" in str(exc) + return "all / 单类型 / 非法类型校验都正确" + raise AssertionError("非法地形类型没有报错") + + +def test_all_profiles_and_protocols_are_loadable(): + """每个机器人描述、每个协议预设都要能被命令行选中。""" + module = _load_eval_suite() + parser = module._build_parser() + from eval_kit.profiles import PROFILES + from eval_kit.protocols import PROTOCOLS + + for name in PROFILES: + args = parser.parse_args(["--profile", name]) + assert args.profile == name + for name in PROTOCOLS: + args = parser.parse_args(["--protocol", name]) + protocol = module._apply_overrides(PROTOCOLS[name], args) + assert protocol.steps >= 1 + assert protocol.seeds, f"协议 {name} 没有种子" + return f"{len(PROFILES)} 个描述、{len(PROTOCOLS)} 个协议预设都可加载" + + +def test_eval_types_match_report_module(): + """命令行的评估类型必须与指标筛选表一致，避免出现选了却没指标的项。""" + module = _load_eval_suite() + from eval_kit.core.report import EVAL_TYPES, METRIC_SELECTION + + args = module._build_parser().parse_args([]) + choices = args.eval_type + assert choices in EVAL_TYPES + for eval_type in EVAL_TYPES: + if eval_type in ("all", "upper_limits"): + continue + assert eval_type in METRIC_SELECTION, f"{eval_type} 没有对应的指标筛选规则" + return f"{len(EVAL_TYPES)} 个评估类型，与 METRIC_SELECTION 对齐" + + +def test_parser_compatible_with_app_launcher(): + """自身的参数不能和 AppLauncher 冲突。 + + main() 的顺序是「先 _build_parser()，再 AppLauncher.add_app_launcher_args()」。 + 这里复刻同样的顺序：任何冲突都会在这一步抛错，也就是脚本连启动都做不到。 + 真跑时曾两次撞上这个坑（先是 --device、后是 --verbose），所以这个检查必须留下。 + """ + module = _load_eval_suite() + parser = module._build_parser() + parser = module.AppLauncher.add_app_launcher_args(parser) + args = parser.parse_args([]) + + # --device / --headless 由 AppLauncher 提供，自己的 parser 里不能再定义 + assert args.device is None + assert args.headless is False + assert args.task == "Rough-Deeprobotics-M20-v0" + assert args.num_envs == 250 + return "与 AppLauncher 的 14 个参数名无冲突，--device/--headless 由其提供" + + +def test_no_option_string_collisions_with_app_launcher(): + """直接枚举检查：自身 parser 的任何 option string 都不能落在 AppLauncher 的名单里。""" + module = _load_eval_suite() + parser = module._build_parser() + own = {option for action in parser._actions for option in action.option_strings} + conflicts = sorted(own & set(APPLAUNCHER_ARG_NAMES)) + assert not conflicts, ( + f"与 AppLauncher 冲突: {conflicts} —— 真实启动时会抛 " + "argparse.ArgumentError: conflicting option string" + ) + return f"自身 {len(own)} 个 option string 与 AppLauncher 的 {len(APPLAUNCHER_ARG_NAMES)} 个均不冲突" + + +def test_verbose_terrain_flag_replaces_verbose(): + """--verbose 归 AppLauncher，自己的逐格调试用 --verbose-terrain。""" + module = _load_eval_suite() + parser = module._build_parser() + args = parser.parse_args(["--verbose-terrain"]) + assert args.verbose_terrain is True + plain = parser.parse_args([]) + assert plain.verbose_terrain is False + return "--verbose-terrain 生效，未占用 --verbose" + + +def test_output_path_defaults_next_to_checkpoint(): + """默认输出落在 checkpoint 同目录的 eval_kit/ 下，不覆盖 eval.py 的 eval/。""" + module = _load_eval_suite() + import tempfile + + with tempfile.TemporaryDirectory() as directory: + checkpoint = str(Path(directory) / "model_100.pt") + Path(checkpoint).write_bytes(b"") + args = module._build_parser().parse_args([]) + path = module._resolve_output_path(args, checkpoint, "all") + + assert "/eval_kit/" in path, path + assert path.endswith(".json") + assert "all_" in Path(path).name + return f"默认输出目录为 /eval_kit/" + + +def test_livestream_value_falls_back_to_env_var(): + """--livestream 默认 -1 表示「听 LIVESTREAM 环境变量」，不能当成 0。""" + import os + + module = _load_eval_suite() + parser = module._build_parser() + + args = parser.parse_args([]) + args.livestream = -1 + previous = os.environ.get("LIVESTREAM") + os.environ.pop("LIVESTREAM", None) + try: + assert module._livestream_value(args) == 0, "没设环境变量时应视为关闭" + os.environ["LIVESTREAM"] = "2" + assert module._livestream_value(args) == 2, "应回退读 LIVESTREAM 环境变量" + + explicit = parser.parse_args([]) + explicit.livestream = 1 + assert module._livestream_value(explicit) == 1, "命令行显式值优先于环境变量" + finally: + if previous is None: + os.environ.pop("LIVESTREAM", None) + else: + os.environ["LIVESTREAM"] = previous + return "命令行 > LIVESTREAM 环境变量 > 关闭" + + +def test_viewer_env_index_covers_livestream(): + """推流也要设视口相机；真 headless（且不推流）才返回 None。 + + 这条决定推流画面能不能看到机器人：``AppLauncher`` 在开启推流时会把 + ``headless`` 改写成 True，所以不能按 ``headless`` 判断「用户要不要看画面」。 + """ + module = _load_eval_suite() + # --headless 由 AppLauncher 提供，所以这里要按 main() 的顺序加上 + parser = module.AppLauncher.add_app_launcher_args(module._build_parser()) + + plain_headless = parser.parse_args(["--headless"]) + plain_headless.livestream = 0 + assert module._viewer_env_index(plain_headless) is None, "headless 且不推流时不应设视口相机" + + streamed = parser.parse_args(["--headless", "--view-env", "7"]) + streamed.livestream = 2 + assert module._viewer_env_index(streamed) == 7, "推流时视口相机应生效并跟随 --view-env" + + default_target = parser.parse_args([]) + default_target.livestream = 2 + assert module._viewer_env_index(default_target) == 0, "不传 --view-env 时默认跟随 env 0" + + gui = parser.parse_args([]) + gui.livestream = 0 + assert module._viewer_env_index(gui) == 0, "本地 GUI（非 headless）仍然跟随 env 0" + + # 自由视角：不跟随，且 --free-camera 优先于 --view-env + free = parser.parse_args(["--headless", "--free-camera"]) + free.livestream = 2 + assert module._viewer_env_index(free) is None, "--free-camera 时不应跟随任何环境" + both = parser.parse_args(["--headless", "--free-camera", "--view-env", "3"]) + both.livestream = 2 + assert module._viewer_env_index(both) is None, "--free-camera 与 --view-env 同时给出时以前者为准" + return "headless→None；推流 / 本地 GUI→跟随 --view-env（默认 0）；--free-camera 优先于跟随" + + +def test_free_camera_only_applies_with_a_picture(): + """--free-camera 只在有画面时生效；给定 --view-env 也让位于它。""" + module = _load_eval_suite() + parser = module.AppLauncher.add_app_launcher_args(module._build_parser()) + + headless = parser.parse_args(["--headless", "--free-camera"]) + headless.livestream = 0 + assert module._display_enabled(headless) is False, "headless 且不推流时没有画面" + assert module._viewer_env_index(headless) is None + + streamed = parser.parse_args(["--free-camera"]) + streamed.livestream = 2 + assert module._display_enabled(streamed) is True, "推流算有画面" + assert module._viewer_env_index(streamed) is None, "有画面 + --free-camera → 不跟随" + + # 显式给出 --view-env 但没有 --free-camera：仍然跟随 + following = parser.parse_args(["--view-env", "2"]) + following.livestream = 2 + assert module._viewer_env_index(following) == 2 + return "无画面时 --free-camera 无效果；有画面时不跟随；未给该参数时行为不变" + + +def test_pause_defaults_to_connection_window_only_when_streaming(): + """--pause-s 是给推流客户端留连接时间的，默认值在 main() 里按是否推流决定。""" + module = _load_eval_suite() + # --livestream 由 AppLauncher 提供，所以这里要按 main() 的顺序加上 + parser = module.AppLauncher.add_app_launcher_args(module._build_parser()) + + assert parser.parse_args([]).pause_s is None, "默认必须是 None，否则区分不出用户显式传的 0" + assert parser.parse_args(["--pause-s", "0"]).pause_s == 0.0 + assert parser.parse_args(["--pause-s", "12.5"]).pause_s == 12.5 + + def resolve(arguments): + """复刻 main() 里的取值规则。""" + parsed = parser.parse_args(arguments) + streaming = module._livestream_value(parsed) > 0 + return parsed.pause_s if parsed.pause_s is not None else (30.0 if streaming else 0.0) + + assert resolve([]) == 0.0, "不推流时不该停顿" + assert resolve(["--livestream", "2"]) == 30.0, "推流时默认留 30 s 连接时间" + assert resolve(["--livestream", "2", "--pause-s", "0"]) == 0.0, "显式 0 可以关掉停顿" + assert resolve(["--livestream", "1"]) == 30.0, "公网推流同样要留连接时间" + return "不推流默认 0 s；推流默认 30 s；显式值优先" + + +def test_local_ipv4_addresses_skips_loopback(): + """打印给 WebRTC 客户端的地址里不能出现回环地址。""" + module = _load_eval_suite() + addresses = module._local_ipv4_addresses() + assert all(not address.startswith("127.") for address in addresses), addresses + assert all(address.count(".") == 3 for address in addresses), addresses + return f"拿到 {len(addresses)} 个非回环地址：{addresses}" + + +def test_record_frames_defaults_off(): + """抓帧是可选功能：默认不抓，给了目录和间隔才生效。""" + module = _load_eval_suite() + parser = module._build_parser() + + plain = parser.parse_args([]) + assert plain.record_frames is None, "默认不应抓帧" + assert plain.record_every == 5, "默认每 5 步一张（0.1 s 仿真时间，合成 10 fps 视频）" + + on = parser.parse_args(["--record-frames", "/tmp/frames", "--record-every", "1"]) + assert on.record_frames == "/tmp/frames" + assert on.record_every == 1 + return "默认关闭；--record-frames 指定目录、--record-every 指定间隔" + + +def main() -> int: + tests = [value for name, value in sorted(globals().items()) if name.startswith("test_")] + failures: list[tuple[str, str]] = [] + for test in tests: + name = test.__name__ + try: + note = test() + except (Exception, SystemExit) as exc: # noqa: BLE001 + # 也要接住 SystemExit：argparse 参数写错会直接 SystemExit(2)， + # 不接的话整个测试文件当场退出，后面的测试一条都不跑、也没有汇总行。 + failures.append((name, f"{type(exc).__name__}: {exc}")) + print(f"FAIL {name}\n {type(exc).__name__}: {exc}") + else: + print(f"ok {name}" + (f" — {note}" if note else "")) + + print() + if failures: + print(f"{len(failures)}/{len(tests)} 个测试失败") + return 1 + print(f"全部 {len(tests)} 个测试通过") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/eval_kit/tests/test_core.py b/eval_kit/tests/test_core.py new file mode 100644 index 0000000..3da4229 --- /dev/null +++ b/eval_kit/tests/test_core.py @@ -0,0 +1,645 @@ +# Copyright (c) 2025 Deep Robotics +# SPDX-License-Identifier: BSD 3-Clause + +"""评价内核的单元测试：用合成轨迹验证指标，不需要仿真环境。 + +运行方式（在仓库根目录）:: + + python3 eval_kit/tests/test_core.py + +合成轨迹的价值在于：指标的期望值可以手算，因此能证明指标实现是对的， +而不是「跑起来没报错」。这也是把指标做成纯函数的主要理由。 +""" + +from __future__ import annotations + +import math +import sys +from dataclasses import replace +from pathlib import Path + +import numpy as np + +if __package__ in (None, ""): + sys.path.insert(0, str(Path(__file__).resolve().parents[2])) + +from eval_kit.core.aggregation import ( + accumulated_success, + aggregate_runs, + build_upper_limit_matrix, + upper_limit_score, +) +from eval_kit.core.metrics import compute_metrics, support_margin_xy, tilt_from_projected_gravity +from eval_kit.core.naming import resolve_matching_indices +from eval_kit.core.report import ( + filter_metrics, + format_aggregate_summary, + format_matrix, + format_metric_result, + format_numpy_matrix, + format_rollout_summary, + format_upper_limit_summary, +) +from eval_kit.core.schema import EvalProtocol, RobotProfile, Trace +from eval_kit.core.terrain_spec import get_terrain_spec +from eval_kit.profiles import get_profile +from eval_kit.protocols import get_protocol + +G = 9.81 + + +# --------------------------------------------------------------------------- 构造器 + + +def make_trace( + *, + num_envs: int = 3, + steps: int = 10, + dt: float = 0.02, + speed: float = 1.0, + tilt_rad: float = 0.0, + torque: float = 1.0, + joint_vel: float = 1.0, + num_joints: int = 2, + num_feet: int = 4, + num_support: int = 2, + support_force: float = 0.0, + support_contact_time: float = 0.0, + mass: float | None = 26.581, + start_x: float = -0.9, + command: tuple[float, float, float] = (1.0, 0.0, 0.0), + accel_spike_at: int | None = None, + accel_spike: float = 0.1, +) -> Trace: + """造一条可控的合成轨迹。 + + 机器人沿 +x 匀速前进；``accel_spike_at`` 指定在第几帧把速度永久抬高 + ``accel_spike``，用于制造一次可预期的不稳定计数。 + """ + frames = steps + 1 + time = np.arange(frames, dtype=np.float64) * dt + + speed_profile = np.full(frames, float(speed), dtype=np.float64) + if accel_spike_at is not None: + speed_profile[accel_spike_at:] += accel_spike + + x = start_x + np.concatenate([[0.0], np.cumsum(speed_profile[1:] * dt)]) + + base_pos_w = np.zeros((frames, num_envs, 3), dtype=np.float64) + base_pos_w[:, :, 0] = x[:, None] + + base_lin_vel_w = np.zeros((frames, num_envs, 3), dtype=np.float64) + base_lin_vel_w[:, :, 0] = speed_profile[:, None] + + projected_gravity_b = np.zeros((frames, num_envs, 3), dtype=np.float64) + projected_gravity_b[:, :, 0] = -math.sin(tilt_rad) + projected_gravity_b[:, :, 2] = -math.cos(tilt_rad) + + joint_torque = np.full((frames, num_envs, num_joints), float(torque), dtype=np.float64) + joint_vel_arr = np.full((frames, num_envs, num_joints), float(joint_vel), dtype=np.float64) + + feet_pos_w = np.zeros((frames, num_envs, num_feet, 2), dtype=np.float64) + if num_feet: + offsets = np.linspace(-0.3, 0.3, num_feet) + for foot_index, offset_x in enumerate(offsets): + feet_pos_w[:, :, foot_index, 0] = x[:, None] + offset_x + feet_pos_w[:, :, foot_index, 1] = 0.1 if foot_index % 2 else -0.1 + + support_force_n = np.full((frames, num_envs, num_support), float(support_force), dtype=np.float64) + support_contact_time_s = np.full( + (frames, num_envs, num_support), float(support_contact_time), dtype=np.float64 + ) + + return Trace( + dt=dt, + command=np.asarray(command, dtype=np.float64), + base_pos_w=base_pos_w, + base_lin_vel_w=base_lin_vel_w, + base_lin_vel_b=base_lin_vel_w.copy(), + base_ang_vel_b=np.zeros((frames, num_envs, 3), dtype=np.float64), + projected_gravity_b=projected_gravity_b, + joint_torque=joint_torque, + joint_vel=joint_vel_arr, + feet_pos_w=feet_pos_w, + support_force_n=support_force_n, + support_contact_time_s=support_contact_time_s, + done=np.zeros((frames, num_envs), dtype=np.float64), + mass_kg=mass, + ) + + +def short_protocol(**overrides) -> EvalProtocol: + """时间很短、成功距离很小的协议，便于手算期望值。""" + kwargs: dict[str, object] = {"steps": 10} + kwargs.update(overrides) + return replace(get_protocol("standard_v1"), **kwargs) # type: ignore[arg-type] + + +# ------------------------------------------------------------------------------- 测试 + + +def test_perfect_tracker_succeeds(): + """匀速前进、姿态水平：成功率 1，跟踪误差 0。""" + protocol = short_protocol(success_distance_m=0.09) + metrics = compute_metrics(make_trace(), get_profile("solenopsis"), protocol) + + assert metrics["success_rate"] == 1.0, metrics["success_rate"] + assert metrics["success_count"] == 3.0 + assert abs(metrics["mean_lin_vel_error"]) < 1e-9, metrics["mean_lin_vel_error"] + assert abs(metrics["mean_yaw_vel_error"]) < 1e-9 + assert abs(metrics["avg_max_tilt_rad"]) < 1e-9 + assert metrics["unstability_rate"] == 0.0 + # speed=1 m/s 从 start_x=-0.9 出发：第 4 帧进度 0.08 < 0.09，第 5 帧 0.10 >= 0.09， + # 因此首次越线在第 5 帧，耗时 5*0.02 = 0.10 s。 + assert abs(metrics["avg_time_consumption_s"] - 0.10) < 1e-9, metrics["avg_time_consumption_s"] + return "success_rate=1.00, t=0.10 s, 跟踪误差 0" + + +def test_stationary_robot_fails(): + """原地不动：成功率 0，最终进度 0。""" + protocol = short_protocol(success_distance_m=0.1) + metrics = compute_metrics(make_trace(speed=0.0), get_profile("solenopsis"), protocol) + + assert metrics["success_rate"] == 0.0 + assert abs(metrics["avg_final_progress_m"]) < 1e-9 + assert math.isnan(metrics["avg_time_consumption_s"]), metrics["avg_time_consumption_s"] + assert math.isnan(metrics["time_efficiency"]) + return "success_rate=0.00, 进度 0, 耗时 nan" + + +def test_faller_is_not_success(): + """走够了距离但倾角超阈值：不算成功。""" + protocol = short_protocol(success_distance_m=0.1, tilt_threshold_rad=0.70) + # 倾角 1.0 rad > 0.70 rad + metrics = compute_metrics(make_trace(tilt_rad=1.0), get_profile("solenopsis"), protocol) + + assert abs(metrics["avg_final_progress_m"] - 0.2) < 1e-9 + assert metrics["success_rate"] == 0.0, "倾角超限却算成功了" + assert abs(metrics["avg_max_tilt_rad"] - 1.0) < 1e-9 + assert metrics["avg_stable_time_s"] == 0.0 + return "进度达成但 success_rate=0.00，max_tilt=1.00 rad" + + +def test_energy_and_cot_match_hand_calculation(): + """能耗与 COT 与手算值一致。""" + steps, dt, joints = 10, 0.02, 2 + mass = 26.581 + protocol = short_protocol(steps=steps) + trace = make_trace(steps=steps, dt=dt, num_joints=joints, mass=mass) + metrics = compute_metrics(trace, get_profile("solenopsis"), protocol) + + expected_energy = 1.0 * 1.0 * joints * steps * dt # |tau*vel| * 关节数 * 步数 * dt + expected_route = 1.0 * steps * dt + # 分母带 1e-6 保护项，与旧 eval.py 的 `route + 1e-6` 一致 + expected_cot = expected_energy / ((expected_route + 1e-6) * mass * G) + + assert abs(metrics["avg_energy_j"] - expected_energy) < 1e-12, (metrics["avg_energy_j"], expected_energy) + assert abs(metrics["avg_cost_of_transport"] - expected_cot) < 1e-12, ( + metrics["avg_cost_of_transport"], + expected_cot, + ) + assert abs(metrics["robot_mass_kg"] - mass) < 1e-12 + return f"energy={expected_energy:.4f} J, COT={expected_cot:.5f}, mass={mass} kg（自动取自 Trace）" + + +def test_signed_energy_convention_differs(): + """有符号口径在存在负功时给出更小的能耗；两种约定都被记录。""" + steps, dt = 10, 0.02 + trace = make_trace(steps=steps, dt=dt, num_joints=2, joint_vel=1.0) + # 让后一半控制步的力矩反向，制造负功：帧 1..5 为 +1，帧 6..10 为 -1。 + half = steps // 2 + trace.joint_torque[1 + half :] = -1.0 + + absolute = compute_metrics(trace, get_profile("solenopsis"), short_protocol(energy_convention="absolute")) + signed = compute_metrics(trace, get_profile("solenopsis"), short_protocol(energy_convention="signed")) + + assert absolute["avg_energy_j"] > signed["avg_energy_j"], ( + absolute["avg_energy_j"], + signed["avg_energy_j"], + ) + assert abs(signed["avg_energy_j"]) < 1e-12, "正负功抵消后应为 0" + assert absolute["energy_convention"] == "absolute" + assert signed["energy_convention"] == "signed" + return f"absolute={absolute['avg_energy_j']:.4f} J, signed={signed['avg_energy_j']:.4f} J" + + +def test_instability_counts_acceleration_spikes(): + """一次速度跳变恰好计一次不稳定。""" + steps, dt = 10, 0.02 + protocol = replace( + short_protocol(steps=steps), acc_threshold_g=2.0 / G + ) + # 第 5 帧速度永久 +0.1 m/s → 该处加速度 0.1/0.02 = 5 m/s² > 2 m/s² + metrics = compute_metrics( + make_trace(steps=steps, dt=dt, accel_spike_at=5, accel_spike=0.1), + get_profile("solenopsis"), + protocol, + ) + assert abs(metrics["unstability_rate"] - 1.0 / steps) < 1e-12, metrics["unstability_rate"] + return f"unstability_rate={metrics['unstability_rate']:.3f} (= 1/{steps})" + + +def test_cost_of_transport_uses_trace_mass_over_profile(): + """Trace 里的质量优先于 profile 兜底值。""" + profile = replace(get_profile("solenopsis"), mass_kg=999.0) + metrics = compute_metrics(make_trace(mass=26.581), profile, short_protocol()) + assert abs(metrics["robot_mass_kg"] - 26.581) < 1e-12 + + trace_no_mass = make_trace(mass=None) + metrics_fallback = compute_metrics(trace_no_mass, profile, short_protocol()) + assert abs(metrics_fallback["robot_mass_kg"] - 999.0) < 1e-12 + return "Trace 优先，profile 兜底" + + +def test_report_formatters_run_on_real_shapes(): + """报告层三个格式化函数都要能真跑——它们之前在没有任何测试覆盖的情况下 + 把一个不存在的 numpy 参数留在里面，直到实跑才炸。""" + spec = get_terrain_spec("standard_v1") + terrain_types = ["stairs_up", "boxes"] + level_ids = np.array([0, 0, 1, 1, 0, 0, 1, 1]) + type_ids = np.array([0, 0, 0, 0, 1, 1, 1, 1]) + per_env = { + "success": [1.0, 1.0, 0.0, 0.0, 1.0, 0.0, 1.0, 1.0], + "max_progress": [3.0, 3.0, 1.0, 1.0, 2.5, 1.5, 3.0, 3.0], + "max_tilt": [0.1] * 8, + } + matrix = build_upper_limit_matrix(per_env, level_ids, type_ids, terrain_types, spec, 2) + + # 普通数值 + text = format_numpy_matrix(matrix["matrix"]["success_rate"], title="success rate") + assert "success rate" in text and "1.00" in text, text + + # 含 nan：把 title 里的 "nan" 排除在外，只检查矩阵主体 + empty = build_upper_limit_matrix( + per_env, level_ids, type_ids, terrain_types, spec, 4 + ) + body_only = format_numpy_matrix(empty["matrix"]["success_rate"], title="unassigned levels").split("\n", 1)[1] + assert "-" in body_only, f"nan 没有被替换成 -：{body_only}" + assert "nan" not in body_only, f"主体里仍出现 nan：{body_only}" + + # 上限汇总：走 build_upper_limit_matrix 的真实输出结构 + summary = format_upper_limit_summary(matrix) + assert "success rate" in summary and "score" in summary, summary + assert "difficulty units" in summary + + # 逐等级矩阵（format_matrix） + table = format_matrix(matrix["matrix"]["success_rate"], terrain_types, title="table") + assert "level" in table and "stairs_up" in table, table + + # 单指标与汇总 + metrics = compute_metrics(make_trace(), get_profile("solenopsis"), short_protocol(success_distance_m=0.05)) + assert "success rate" in format_metric_result("access_rate", metrics) + assert "protocol" in format_rollout_summary(metrics) + return "四个格式化函数在真实结构上都能跑，nan 正确显示为 -" + + +def test_minimum_terrain_size_for_bounds(): + """地砖下限公式：与网格行数无关，取默认值时是 3.8 m。 + + 这个函数必须在纯内核里——编排进程不启动仿真，无法 import 任何依赖 isaaclab + 的模块（会因缺 pxr 失败）。之前它放在适配层，父进程一调用就崩。 + """ + protocol = get_protocol("standard_v1") + # start_x = -0.9, buffer = 3.0, border = 2.0 → 2*(0.9+3.0-2.0) = 3.8 + assert abs(protocol.minimum_terrain_size_for_bounds() - 3.8) < 1e-12 + + # clean 预设的地砖应当不小于该下限 + clean = get_protocol("standard_v1_clean") + assert clean.terrain_size_m >= clean.minimum_terrain_size_for_bounds() + + # 与网格行数无关 + for levels in (2, 5, 8, 16): + assert abs(replace(protocol, num_levels=levels).minimum_terrain_size_for_bounds() - 3.8) < 1e-12 + + # border 与 buffer 可调，且不会返回负数 + assert abs(protocol.minimum_terrain_size_for_bounds(border_width=1.0, distance_buffer=1.0) - 1.8) < 1e-12 + assert protocol.minimum_terrain_size_for_bounds(border_width=10.0, distance_buffer=0.0) == 0.0 + + # 起始位置越靠外，要求越大 + far = replace(protocol, start_pose=(-1.5, 0.0, 0.0)) + assert far.minimum_terrain_size_for_bounds() > protocol.minimum_terrain_size_for_bounds() + return "默认 3.8 m，与行数无关，参数可调且非负" + + +def test_missing_morphology_returns_nan_not_zero(): + """没有足端 / 没有支撑体的机器人应返回 nan，而不是看起来像测出来的 0。""" + profile = RobotProfile( + name="no_morphology", + mass_kg=10.0, + feet_body_patterns=(), + support_body_patterns=(), + actuated_joint_patterns=(".*",), + ) + metrics = compute_metrics(make_trace(num_feet=0, num_support=0), profile, short_protocol()) + assert math.isnan(metrics["avg_stability_margin_m"]), metrics["avg_stability_margin_m"] + assert math.isnan(metrics["avg_support_contact_ratio"]) + assert math.isnan(metrics["successful_runs_with_support_ratio"]) + assert metrics["stability_margin_method"] == "unavailable" + return "缺失形态 → nan（旧实现给 0.00）" + + +def test_support_contact_events_and_ratio(): + """支撑接触比例与事件计数可手算。""" + steps, dt = 10, 0.02 + trace = make_trace(steps=steps, dt=dt, support_force=0.0) + # 连续 3 帧有超过阈值的支撑接触 → 1 次事件，占 3/10 步 + trace.support_force_n[2:5] = 100.0 + trace.support_contact_time_s[2:5] = 0.05 + + protocol = replace(short_protocol(steps=steps, success_distance_m=0.1)) + metrics = compute_metrics(trace, get_profile("solenopsis"), protocol) + + assert abs(metrics["avg_support_contact_ratio"] - 3.0 / steps) < 1e-12, metrics["avg_support_contact_ratio"] + assert abs(metrics["avg_support_events"] - 1.0) < 1e-12, metrics["avg_support_events"] + assert abs(metrics["successful_runs_with_support_ratio"] - 1.0) < 1e-12 + return "contact_ratio=0.300, events=1.0, 成功率内用支撑=100%" + + +def test_support_margin_sign_and_tilt_formula(): + """支撑裕度在包围盒内外符号正确；倾角公式反解一致。""" + root_inside = np.array([[0.0, 0.0]]) + feet = np.array([[[-1.0, -1.0], [1.0, -1.0], [1.0, 1.0], [-1.0, 1.0]]], dtype=np.float64) + inside = support_margin_xy(root_inside, feet) + assert abs(inside[0] - 1.0) < 1e-12, inside + + root_outside = np.array([[2.0, 0.0]]) + outside = support_margin_xy(root_outside, feet) + assert abs(outside[0] - (-1.0)) < 1e-12, outside + + theta = 0.5 + gravity_b = np.array([[[-math.sin(theta), 0.0, -math.cos(theta)]]], dtype=np.float64) + tilt = tilt_from_projected_gravity(gravity_b) + assert abs(tilt[0, 0] - theta) < 1e-9 + return "盒内 +1.000，盒外 -1.000，倾角反解 0.500 rad" + + +def test_accumulated_and_score_match_legged_gym_formula(): + """累积成功率与 score 与 legged_gym criterion.py 的公式一致。""" + rates = [1.0, 0.8, 0.5, 0.0] + accumulated = accumulated_success(rates) + assert np.allclose(accumulated, [1.0, 0.8, 0.4, 0.0]), accumulated + + # 手算：i=0 → 1.0*1*(1-0.8)=0.2；i=1 → 0.8*2*(1-0.5)=0.8； + # i=2 → 0.4*3*(1-0.0)=1.2；i=3 → 0.0*4*1=0.0 + expected = 0.2 + 0.8 + 1.2 + 0.0 + assert abs(upper_limit_score(rates) - expected) < 1e-12, (upper_limit_score(rates), expected) + + # 全通过与全失败是两个极端 + assert abs(upper_limit_score([1.0] * 4) - 4.0) < 1e-12 + assert abs(upper_limit_score([0.0] * 4)) < 1e-12 + return f"score={expected:.2f}（手算），全通过=4.00，全失败=0.00" + + +def test_upper_limit_matrix_cells(): + """上限矩阵按「等级 × 地形」正确分格。""" + spec = get_terrain_spec("standard_v1") + terrain_types = ["stairs_up", "boxes"] + num_levels = 2 + envs_per_cell = 2 + + level_ids = np.array([0, 0, 1, 1, 0, 0, 1, 1]) + type_ids = np.array([0, 0, 0, 0, 1, 1, 1, 1]) + # stairs_up: 等级0 两个都成功，等级1 两个都失败 + # boxes: 等级0 一个成功，等级1 两个都成功 + per_env = { + "success": [1.0, 1.0, 0.0, 0.0, 1.0, 0.0, 1.0, 1.0], + "max_progress": [3.0, 3.0, 1.0, 1.0, 2.5, 1.5, 3.0, 3.0], + "max_tilt": [0.1] * 8, + } + result = build_upper_limit_matrix( + per_env, level_ids, type_ids, terrain_types, spec, num_levels + ) + + stairs = result["per_type"]["stairs_up"] + boxes = result["per_type"]["boxes"] + assert stairs["success_rate"] == [1.0, 0.0], stairs["success_rate"] + assert boxes["success_rate"] == [0.5, 1.0], boxes["success_rate"] + assert np.allclose(stairs["success_rate_accumulated"], [1.0, 0.0]) + assert np.allclose(boxes["success_rate_accumulated"], [0.5, 0.5]) + # 行=等级、列=地形 + assert result["matrix"]["success_rate"] == [[1.0, 0.5], [0.0, 1.0]] + # 难度值取自规格 + assert result["matrix"]["difficulty"][0] == [0.04, 0.04] + assert result["matrix"]["difficulty"][1] == [0.06, 0.06] + assert result["num_levels"] == num_levels + assert result["scores"]["stairs_up"] == upper_limit_score([1.0, 0.0]) + return "矩阵 [[1.00,0.50],[0.00,1.00]]，难度 [[0.04,0.04],[0.06,0.06]]" + + +def test_fingerprint_tracks_protocol_changes(): + """改变口径会改变指纹；改变种子不会。""" + base = get_protocol("standard_v1") + same = get_protocol("standard_v1") + assert base.fingerprint() == same.fingerprint() + + changed_energy = replace(base, energy_convention="signed") + changed_tilt = replace(base, tilt_threshold_rad=0.5) + pushed = replace(base, enable_interval_pushes=True) + reseeded = replace(base, seeds=(1, 2, 3)) + + assert base.fingerprint() != changed_energy.fingerprint() + assert base.fingerprint() != changed_tilt.fingerprint() + assert base.fingerprint() != pushed.fingerprint() + assert base.fingerprint() == reseeded.fingerprint(), "种子不影响单次 rollout 的口径" + return "能耗约定/倾角阈值/推力开关都改变指纹，种子不改变" + + +def test_trace_validation_rejects_mismatched_shapes(): + """形状不一致要早失败。""" + trace = make_trace(num_envs=2, steps=5, num_joints=3) + trace.joint_vel = np.zeros((6, 3, 3)) # 环境数写成 3，应为 2 + try: + trace.validate() + except ValueError as exc: + assert "joint_vel" in str(exc), str(exc) + return "捕获到形状错误" + raise AssertionError("形状不一致却没有报错") + + +def test_trace_roundtrip(tmp_path: Path | None = None): + """轨迹存盘再读回，指标不变——这是「事后重新分析」的前提。""" + import tempfile + + trace = make_trace(steps=8) + protocol = short_protocol(steps=8, success_distance_m=0.05) + profile = get_profile("solenopsis") + before = compute_metrics(trace, profile, protocol) + + with tempfile.TemporaryDirectory() as directory: + path = str(Path(directory) / "trace.npz") + trace.save(path) + reloaded = Trace.load(path) + after = compute_metrics(reloaded, profile, protocol) + + for key in before: + if key.startswith("_"): + continue + left, right = before[key], after[key] + if isinstance(left, float) and math.isnan(left): + assert math.isnan(right), key + else: + assert left == right, (key, left, right) + return f"存盘往返后 {len([k for k in before if not k.startswith('_')])} 个指标完全一致" + + +def test_report_filters_and_formats(): + """输出层按评估类型裁剪，且 nan 不炸格式串。""" + metrics = compute_metrics(make_trace(num_feet=0, num_support=0), get_profile("solenopsis"), short_protocol()) + selected = filter_metrics(metrics, "stability_margin") + assert list(selected.keys()) == ["avg_stability_margin_m"] + + text = format_metric_result("stability_margin", metrics) + assert "nan" in text, text + text = format_metric_result("access_rate", metrics) + assert "success rate" in text + return "filter_metrics 裁剪正确，nan 格式化安全" + + +def test_profile_registry(): + """加机器人只改数据：火蚁有轮与支撑体，四足没有支撑体语义。""" + solenopsis = get_profile("solenopsis") + assert solenopsis.feet_body_patterns == (".*_wheel",) + assert solenopsis.support_body_patterns == ("base_link", ".*_knee") + + lite3 = get_profile("lite3") + assert lite3.support_body_patterns == () + + try: + get_profile("nonexistent") + except KeyError: + return "solenopsis / lite3 描述可用，未知名字报错" + raise AssertionError("未知机器人没有报错") + + +def test_naming_resolves_fire_ant_bodies(): + """火蚁的 link/joint 名字能被描述里的正则解析出来。""" + body_names = [ + "base_link", + "fl_knee", "fr_knee", "hl_knee", "hr_knee", + "fl_wheel", "fr_wheel", "hl_wheel", "hr_wheel", + ] + resolution = resolve_matching_indices(get_profile("solenopsis").feet_body_patterns, body_names) + assert resolution.names == ["fl_wheel", "fr_wheel", "hl_wheel", "hr_wheel"], resolution.names + assert resolution.indices == [5, 6, 7, 8] + assert not resolution.unmatched_patterns + + support = resolve_matching_indices(get_profile("solenopsis").support_body_patterns, body_names) + assert support.names == ["base_link", "fl_knee", "fr_knee", "hl_knee", "hr_knee"], support.names + + joint_names = [ + "fl_knee_joint", "fr_knee_joint", "hl_knee_joint", "hr_knee_joint", + "fl_wheel_joint", "fr_wheel_joint", "hl_wheel_joint", "hr_wheel_joint", + ] + joints = resolve_matching_indices(get_profile("solenopsis").actuated_joint_patterns, joint_names) + assert len(joints) == 8, joints.names + return f"足端 4 个、支撑体 5 个、驱动关节 {len(joints)} 个" + + +def test_naming_does_not_crash_when_morphology_is_missing(): + """换一个没有轮子的机器人：未命中的正则被记录下来，而不是抛异常。""" + quadruped_bodies = ["base", "FL_hip", "FL_thigh", "FL_shank", "FL_FOOT"] + resolution = resolve_matching_indices(get_profile("solenopsis").feet_body_patterns, quadruped_bodies) + assert len(resolution) == 0 + assert resolution.unmatched_patterns == [".*_wheel"] + assert not resolution + + # 四足自己的描述能解析 + lite3 = resolve_matching_indices(get_profile("lite3").feet_body_patterns, quadruped_bodies) + assert lite3.names == ["FL_FOOT"] + # 空描述合法，得到空结果 + empty = resolve_matching_indices(get_profile("lite3").support_body_patterns, quadruped_bodies) + assert len(empty) == 0 and not empty.unmatched_patterns + return "无轮机器人：未命中被记录，不崩溃" + + +def test_naming_reports_ambiguous_and_bad_regex(): + """一个名字命中多个正则会被记录；正则写错要立刻报错。""" + resolution = resolve_matching_indices((".*_FOOT", "FL_.*"), ["FL_FOOT", "FR_FOOT"]) + assert resolution.names == ["FL_FOOT", "FR_FOOT"] + assert resolution.ambiguous_names == [("FL_FOOT", [".*_FOOT", "FL_.*"])], resolution.ambiguous_names + + try: + resolve_matching_indices(("[unclosed",), ["anything"]) + except ValueError as exc: + assert "无法编译" in str(exc) + return "多重命中被记录，坏正则报错" + raise AssertionError("坏正则没有报错") + + +def test_aggregate_runs_mean_std_and_nan_handling(): + """多次试验聚合：mean/std 正确，nan 不污染结果，非数值字段单独记录。""" + runs = [ + (0, {"success_rate": 0.5, "avg_energy_j": 100.0, "profile": "solenopsis", "avg_cost_of_transport": float("nan")}), + (1, {"success_rate": 1.0, "avg_energy_j": 200.0, "profile": "solenopsis", "avg_cost_of_transport": 0.4}), + (2, {"success_rate": 0.75, "avg_energy_j": 150.0, "profile": "solenopsis", "avg_cost_of_transport": float("nan")}), + ] + aggregate = aggregate_runs(runs) + + assert aggregate["num_runs"] == 3 + assert aggregate["seeds"] == [0, 1, 2] + + success = aggregate["summary"]["success_rate"] + assert abs(success["mean"] - 0.75) < 1e-12, success + assert abs(success["std"] - 0.25) < 1e-12, success + assert success["min"] == 0.5 and success["max"] == 1.0 + assert success["values"] == [0.5, 1.0, 0.75] + + # nan 被剔除，只剩 1 个有效值 → std 记 0 + cot = aggregate["summary"]["avg_cost_of_transport"] + assert abs(cot["mean"] - 0.4) < 1e-12, cot + assert cot["std"] == 0.0 + assert cot["finite_runs"] == 1 + + # 非数值字段单独放在 constant 里 + assert aggregate["constant"]["profile"] == "solenopsis" + assert "profile_INCONSISTENT" not in aggregate["constant"] + + # 全部无效 → mean 是 nan，且记录了有效次数 + all_nan = aggregate_runs([(0, {"x": float("nan")}), (1, {"x": float("nan")})]) + assert math.isnan(all_nan["summary"]["x"]["mean"]) + assert all_nan["summary"]["x"]["finite_runs"] == 0 + + text = format_aggregate_summary(aggregate) + assert "success_rate" in text and "±" in text + return "mean/std 正确，nan 被剔除，非数值字段单独登记" + + +def test_aggregate_runs_flags_inconsistent_metadata(): + """两次 run 的口径不一致时要显式暴露，而不是悄悄取第一个。""" + aggregate = aggregate_runs( + [ + (0, {"success_rate": 1.0, "energy_convention": "absolute"}), + (1, {"success_rate": 1.0, "energy_convention": "signed"}), + ] + ) + assert "energy_convention_INCONSISTENT" in aggregate["constant"], aggregate["constant"] + return "口径不一致会显式标记" + + +# ------------------------------------------------------------------------------- 运行器 + + +def main() -> int: + tests = [value for name, value in sorted(globals().items()) if name.startswith("test_")] + failures: list[tuple[str, str]] = [] + for test in tests: + name = test.__name__ + try: + note = test() + except Exception as exc: # noqa: BLE001 - 测试运行器需要捕获全部异常 + failures.append((name, f"{type(exc).__name__}: {exc}")) + print(f"FAIL {name}\n {type(exc).__name__}: {exc}") + else: + suffix = f" — {note}" if note else "" + print(f"ok {name}{suffix}") + + print() + if failures: + print(f"{len(failures)}/{len(tests)} 个测试失败") + return 1 + print(f"全部 {len(tests)} 个测试通过") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) 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/eval_suite.py b/scripts/reinforcement_learning/rsl_rl/eval_suite.py new file mode 100644 index 0000000..1a97dd4 --- /dev/null +++ b/scripts/reinforcement_learning/rsl_rl/eval_suite.py @@ -0,0 +1,893 @@ +# Copyright (c) 2025 Deep Robotics +# SPDX-License-Identifier: BSD 3-Clause + +"""评价套件的命令行入口。 + +与 ``eval.py`` 并存、互不影响：本脚本不导入也不修改 ``eval.py``，两者可以并行使用， +输出 JSON 结构相似（指标名完全一致），便于互相对照。 + +用法示例:: + + # 标准评估（默认协议与火蚁描述） + python scripts/reinforcement_learning/rsl_rl/eval_suite.py \\ + --task Rough-Deeprobotics-M20-v0 \\ + --checkpoint /绝对路径/model_19999.pt \\ + --headless + + # 上限评估，5 个随机种子，顺带把轨迹存下来 + python scripts/reinforcement_learning/rsl_rl/eval_suite.py \\ + --task Rough-Deeprobotics-M20-v0 \\ + --checkpoint /绝对路径/model_19999.pt \\ + --protocol standard_v1_multiseed --save-trace /tmp/traces --headless + + # 抗扰评估 + python scripts/reinforcement_learning/rsl_rl/eval_suite.py \\ + --task Rough-Deeprobotics-M20-v0 \\ + --checkpoint /绝对路径/model_19999.pt \\ + --eval-type robustness --headless + + # 用 WebRTC 看画面（不需要显示器；局域网内用 Isaac Sim WebRTC Streaming Client 连） + python scripts/reinforcement_learning/rsl_rl/eval_suite.py \\ + --task Rough-Deeprobotics-M20-v0 \\ + --checkpoint /绝对路径/model_19999.pt \\ + --eval-type upper_limits --terrain-type stairs_up \\ + --livestream 2 + + ``--livestream 1`` 走公网（NVIDIA NVCF 信令，需要 ``PUBLIC_IP``）， + ``--livestream 2`` 走局域网直连；两者都由 ``AppLauncher`` 提供，详见 + ``isaaclab.app.AppLauncher`` 的文档。本脚本只额外保证：视口相机跟着被评估的 + 机器人（``--view-env``）、开始 rollout 前留出连接时间（``--pause-s``）、 + 以及把「客户端该连哪里」打印出来。 + + 协议预设与机器人描述都在 ``eval_kit/`` 里：改口径改 ``protocols.py``， + 加机器人改 ``profiles.py``。 +""" + +from __future__ import annotations + +import argparse +import copy +import json +import os +import subprocess +import sys +import tempfile +from dataclasses import replace +from typing import Any + +#: 原始命令行，用于在需要多个环境时 fork 子进程（见文件末尾的说明）。 +_ORIGINAL_ARGV = list(sys.argv[1:]) + +# --- 路径引导：必须在导入 eval_kit / rl_training 之前完成 -------------------------- +_HERE = os.path.dirname(os.path.abspath(__file__)) +PROJECT_ROOT = os.path.abspath(os.path.join(_HERE, "..", "..", "..")) +PROJECT_SOURCE_PATH = os.path.join(PROJECT_ROOT, "source", "rl_training") +for path in (PROJECT_SOURCE_PATH, PROJECT_ROOT): + if path not in sys.path: + sys.path.insert(0, path) + +from isaaclab.app import AppLauncher # noqa: E402 + +from eval_kit.core.report import ( # noqa: E402 + EVAL_TYPES, + filter_metrics, + format_aggregate_summary, + format_metric_result, + format_rollout_summary, + format_upper_limit_summary, +) +from eval_kit.core.schema import EvalProtocol # noqa: E402 +from eval_kit.core.terrain_spec import get_terrain_spec # noqa: E402 +from eval_kit.profiles import PROFILES, get_profile # noqa: E402 +from eval_kit.protocols import PROTOCOLS, get_protocol # noqa: E402 + + +def _build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description="基于 eval_kit 的机器人性能评估（与 eval.py 并存，指标口径一致）。", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + parser.add_argument("--task", type=str, default="Rough-Deeprobotics-M20-v0", help="任务名。") + parser.add_argument( + "--agent", type=str, default="rsl_rl_cfg_entry_point", help="agent 配置入口名（配合 Hydra）。" + ) + parser.add_argument( + "--profile", type=str, default="solenopsis", choices=sorted(PROFILES), help="机器人形态描述。" + ) + parser.add_argument( + "--protocol", type=str, default="standard_v1", choices=sorted(PROTOCOLS), help="评估协议预设。" + ) + parser.add_argument( + "--eval-type", dest="eval_type", type=str, default="all", choices=list(EVAL_TYPES), help="评估类型。" + ) + + checkpoint = parser.add_argument_group("checkpoint") + checkpoint.add_argument("--checkpoint", type=str, default=None, help="checkpoint 路径（支持绝对路径）。") + checkpoint.add_argument("--load-run", dest="load_run", type=str, default=None, help="运行目录或目录名。") + checkpoint.add_argument( + "--experiment-name", dest="experiment_name", type=str, default=None, + help="实验名，用于在 logs/rsl_rl/<实验名> 下查找；不传则用 agent 配置里的。", + ) + + run = parser.add_argument_group("运行规模（不传则取协议预设值）") + run.add_argument( + "--num-envs", dest="num_envs", type=int, default=250, + help="并行环境数（上限评估时会按网格规模自动覆盖）。", + ) + # 注意：--device 由 AppLauncher.add_app_launcher_args() 提供，这里不能再加。 + # 它保留了 device / headless / enable_cameras / experience / livestream / + # rendering_mode / xr 这几个名字，重复定义会直接抛 ValueError。 + run.add_argument( + "--seeds", type=int, nargs="+", default=None, + help="重复试验的随机种子，覆盖协议预设；给多个种子时会输出 mean±std。", + ) + run.add_argument("--save-trace", dest="save_trace", type=str, default=None, help="把轨迹存到该目录。") + run.add_argument( + "--view-env", dest="view_env", type=int, default=None, + help="让视口相机跟随第 N 个环境的机器人（非 headless 或开启推流时有效）。" + "不传时默认跟随 env 0。环境编号 0 起，" + "编号与（地形, 难度）的对应关系会在启动时打印。", + ) + run.add_argument( + "--free-camera", dest="free_camera", action="store_true", default=False, + help="不要让视口相机跟随任何机器人：起始摆在俯视全场的位置，之后在 WebRTC 客户端里" + "自己移动（脚本与 IsaacLab 都不会再改写相机）。同时给了 --view-env 时以本参数为准。", + ) + run.add_argument( + "--pause-s", dest="pause_s", type=float, default=None, + help="开始 rollout 之前停顿的秒数：期间画面照常推流、物理冻结，按回车可提前开始。" + "用来给 WebRTC 客户端留连接时间。开启推流时默认 30 s，其余情况默认 0（不停顿）。", + ) + run.add_argument( + "--record-frames", dest="record_frames", type=str, default=None, + help="把 rollout 期间的视口画面按帧存成 PNG 到该目录（例如 frame_000150.png）。" + "推流传输不可用时用它离线看评估过程；抓帧间隔见 --record-every。", + ) + run.add_argument( + "--record-every", dest="record_every", type=int, default=5, + help="--record-frames 每多少步抓一帧。", + ) + + overrides = parser.add_argument_group("口径覆盖（不传则取协议预设值，改动会反映到 fingerprint）") + overrides.add_argument("--episode-length-s", dest="episode_length_s", type=float, default=None) + overrides.add_argument("--command-x", dest="command_x", type=float, default=None) + overrides.add_argument("--command-y", dest="command_y", type=float, default=None) + overrides.add_argument("--command-yaw", dest="command_yaw", type=float, default=None) + overrides.add_argument("--start-x", dest="start_x", type=float, default=None) + overrides.add_argument("--start-y", dest="start_y", type=float, default=None) + overrides.add_argument( + "--start-z-offset", dest="start_z_offset", type=float, default=None, + help="起始高度相对地形瓦片中心的抬升 [m]。倒金字塔台阶这类「中心最低」的地形需要正值，" + "否则机器人会被生成在地形内部（默认 0.0，与旧 eval.py 一致）。", + ) + overrides.add_argument("--terrain-size", dest="terrain_size", type=float, default=None, help="地形瓦片边长 [m]。") + overrides.add_argument("--success-distance", dest="success_distance", type=float, default=None) + overrides.add_argument("--tilt-threshold", dest="tilt_threshold", type=float, default=None, help="摔倒倾角 [rad]。") + overrides.add_argument( + "--acc-threshold-g", dest="acc_threshold_g", type=float, default=None, + help="不稳定判定的加速度阈值，单位 g。", + ) + overrides.add_argument( + "--energy-convention", dest="energy_convention", type=str, default=None, choices=["absolute", "signed"] + ) + overrides.add_argument( + "--robot-mass", dest="robot_mass", type=float, default=None, + help="覆盖 COT 用质量；默认从资产自动读取。", + ) + overrides.add_argument( + "--terrain-type", dest="terrain_type", type=str, default="all", + help="上限评估只用某一种地形：all 或地形规格里的类型名。", + ) + overrides.add_argument("--num-levels", dest="num_levels", type=int, default=None, help="上限评估的难度等级数。") + overrides.add_argument("--envs-per-cell", dest="envs_per_cell", type=int, default=None, help="每格机器人数。") + overrides.add_argument( + "--no-terrain-bounds", dest="enforce_terrain_bounds", action="store_false", default=None, + help="去掉越界终止项。地砖小于 3.8 m 时它会误判最外侧格子（详见 --help 里 protocol 的说明）。", + ) + + output = parser.add_argument_group("输出") + output.add_argument("--output", type=str, default=None, help="输出 JSON 路径；默认写在 checkpoint 同目录的 eval_kit/ 下。") + output.add_argument("--print-json", dest="print_json", action="store_true", default=False, help="把 JSON 打到 stdout。") + # 不能叫 --verbose：AppLauncher.add_app_launcher_args() 已经占用了这个名字， + # 重复定义会让 argparse 直接抛 ArgumentError（旧 eval.py 就卡在这里）。 + output.add_argument( + "--verbose-terrain", dest="verbose_terrain", action="store_true", default=False, + help="打印逐格地形/难度调试信息。（注意 --verbose 属于 AppLauncher，行为不同）", + ) + output.add_argument( + "--include-summary", dest="include_summary", action="store_true", default=False, + help="eval-type=all 时额外跑一次混合地形汇总。", + ) + # 内部参数：一个进程只能建一个环境，需要多个环境时由父进程 fork 自己并传这些参数。 + # 不建议手工使用。 + internal = parser.add_argument_group("内部（子进程用，勿手工指定）") + internal.add_argument("--stage", type=str, default="", choices=["", "upper", "plain"]) + internal.add_argument("--stage-seed", dest="stage_seed", type=int, default=0) + internal.add_argument("--stage-out", dest="stage_out", type=str, default=None) + return parser + + +def _apply_overrides(protocol: EvalProtocol, args: argparse.Namespace) -> EvalProtocol: + """把命令行显式给出的口径覆盖到协议预设上。没给的不动，保持预设原值。""" + updates: dict[str, object] = {} + if args.episode_length_s is not None: + updates["steps"] = int(round(args.episode_length_s / protocol.dt)) + if args.command_x is not None: + updates["command"] = (args.command_x, protocol.command[1], protocol.command[2]) + if args.command_y is not None: + current = updates.get("command", protocol.command) + updates["command"] = (current[0], args.command_y, current[2]) + if args.command_yaw is not None: + current = updates.get("command", protocol.command) + updates["command"] = (current[0], current[1], args.command_yaw) + if args.start_x is not None: + updates["start_pose"] = (args.start_x, protocol.start_pose[1], protocol.start_pose[2]) + if args.start_y is not None: + current = updates.get("start_pose", protocol.start_pose) + updates["start_pose"] = (current[0], args.start_y, current[2]) + if args.start_z_offset is not None: + updates["start_z_offset_m"] = args.start_z_offset + if args.terrain_size is not None: + updates["terrain_size_m"] = args.terrain_size + if args.success_distance is not None: + updates["success_distance_m"] = args.success_distance + if args.tilt_threshold is not None: + updates["tilt_threshold_rad"] = args.tilt_threshold + if args.acc_threshold_g is not None: + updates["acc_threshold_g"] = args.acc_threshold_g + if args.energy_convention is not None: + updates["energy_convention"] = args.energy_convention + if args.num_levels is not None: + updates["num_levels"] = args.num_levels + if args.envs_per_cell is not None: + updates["envs_per_cell"] = args.envs_per_cell + if args.seeds is not None: + updates["seeds"] = tuple(args.seeds) + if args.enforce_terrain_bounds is not None: + updates["enforce_terrain_bounds"] = args.enforce_terrain_bounds + return replace(protocol, **updates) # type: ignore[arg-type] + + +def _terrain_types_for(args: argparse.Namespace, spec) -> list[str]: + if args.terrain_type == "all": + return list(spec.types) + if args.terrain_type not in spec.types: + raise SystemExit( + f"地形类型 {args.terrain_type!r} 不在规格 {spec.name!r} 里；可选：{list(spec.types)}" + ) + return [args.terrain_type] + + +def _resolve_output_path(args: argparse.Namespace, resume_path: str, eval_type: str) -> str: + if args.output is not None: + path = os.path.abspath(os.path.expanduser(args.output)) + else: + from datetime import datetime + + directory = os.path.join(os.path.dirname(resume_path), "eval_kit") + suffix = datetime.now().strftime("%Y-%m-%d_%H-%M-%S") + path = os.path.join(directory, f"{eval_type}_{suffix}.json") + parent = os.path.dirname(path) + if parent: + os.makedirs(parent, exist_ok=True) + return path + + +def _livestream_value(args_cli: argparse.Namespace) -> int: + """解析实际生效的推流模式：0 关、1 公网 WebRTC、2 局域网 WebRTC。 + + ``--livestream`` 的默认值是 -1，含义是「听 ``LIVESTREAM`` 环境变量的」， + 所以 -1 时要回退去读环境变量，而不是当成 0。 + """ + try: + value = int(getattr(args_cli, "livestream", -1)) + except (TypeError, ValueError): + value = -1 + if value < 0: + try: + value = int(os.environ.get("LIVESTREAM", 0) or 0) + except ValueError: + value = 0 + return value + + +def _display_enabled(args_cli: argparse.Namespace) -> bool: + """这次运行会不会给用户看画面（本地 GUI 或推流）。""" + return not getattr(args_cli, "headless", True) or _livestream_value(args_cli) > 0 + + +def _viewer_env_index(args_cli: argparse.Namespace) -> int | None: + """视口相机跟随哪个环境；不跟随时（自由视角、或完全不需要画面）返回 None。 + + 必须在构造 ``AppLauncher`` **之前**调用：``AppLauncher`` 会把 ``livestream`` + 从 namespace 里 pop 掉，还会在开启推流时把 ``headless`` 改写成 True（推流在 + 服务端本来就按 headless 启动），等它跑完再读这两个值就分不清「用户要不要看画面」了。 + + 推流时 ``SimulationContext`` 的 ``has_gui()``（本地 GUI 或推流 GUI）仍为真、 + ``env_cfg.viewer`` 也仍然生效，所以推流和本地 GUI 两种情况都要设视口相机， + 否则推流画面停在默认相机上、看不到机器人。 + """ + if not _display_enabled(args_cli) or getattr(args_cli, "free_camera", False): + return None + return 0 if args_cli.view_env is None else int(args_cli.view_env) + + +def _set_free_camera_overview(env: Any) -> None: + """自由视角模式：把相机摆到能看见整片地形的位置。 + + ``env_cfg.viewer.origin_type`` 是 ``world`` 时，``ViewportCameraController`` + 只在初始化时设一次相机、之后不再改写，所以这里设的视角会一直保留， + 客户端里可以随便移动鼠标/键盘来观察。 + """ + import numpy as np + + origins = env.unwrapped.scene.env_origins + center = origins.mean(dim=0) + radius = float((origins - center).norm(dim=1).max().item()) + center_np = center.detach().cpu().numpy() + eye = center_np + np.array([0.0, -1.6 * radius, 1.25 * radius], dtype=center_np.dtype) + env.unwrapped.sim.set_camera_view(eye=eye, target=center_np) + print( + f"[eval_kit] 自由视角：相机不跟随机器人，起始俯视全场" + f"（地形中心 {np.round(center_np, 1).tolist()}，半径 {radius:.1f} m），" + f"之后在客户端里自己移动", + flush=True, + ) + + +def _allow_client_resize(args_cli: argparse.Namespace) -> None: + """让 WebRTC 客户端可以改画面尺寸（默认关掉会冻住画面）。 + + 客户端一旦改窗口大小/全屏，就会发 ``onCustomResizeMessage``；服务端默认拒绝 + （``... must both be positive and /app/livestream/allowDynamicResize enabled``）， + 同时推流用的 ``shared swapchain buffer`` 进入坏状态 —— 之后画面不再更新，而仿真 + 照常在跑（终端里的 rollout 进度还在涨）。这两个设置通过 ``--kit_args`` 传给 Kit。 + + 用户自己显式给了同名设置就不覆盖（想关掉就传 ``--kit_args=--/app/livestream/allowDynamicResize=false``）。 + """ + wanted = ("--/app/livestream/allowResize=true", "--/app/livestream/allowDynamicResize=true") + existing = args_cli.kit_args or "" + extra = [flag for flag in wanted if flag.split("=")[0] not in existing] + if not extra: + return + args_cli.kit_args = " ".join(filter(None, [existing, *extra])) + print( + f"[eval_kit] 推流追加 kit 设置：{' '.join(extra)}（允许客户端调整窗口大小/全屏）", + flush=True, + ) + + +def _grab_viewport_frame(state: dict, directory: str, every: int) -> None: + """每 ``every`` 步向视口要一张 PNG。 + + 抓帧是异步的（等下一次渲染产出），所以这里不阻塞 rollout：上一张还没写完就跳过 + 这一帧。写完的文件名形如 ``frame_000150.png``，可用 ffmpeg 合成视频。 + """ + if every <= 0 or state["step"] % every: + return + pending = state["pending"] + if pending is not None and not os.path.exists(pending): + return + from omni.kit.viewport.utility import capture_viewport_to_file, get_active_viewport + + path = os.path.join(directory, f"frame_{state['step']:06d}.png") + state["pending"] = path + capture_viewport_to_file(get_active_viewport(), path) + + +def _with_frame_capture(policy: Any, directory: str, every: int) -> Any: + """把策略包一层，顺带按间隔抓视口帧（策略每被调用一次就是 rollout 推进一步）。""" + os.makedirs(directory, exist_ok=True) + state: dict = {"step": 0, "pending": None} + print(f"[eval_kit] rollout 期间抓帧到 {directory}（每 {every} 步一张 PNG）", flush=True) + + def wrapped(obs: Any) -> Any: + state["step"] += 1 + _grab_viewport_frame(state, directory, every) + return policy(obs) + + return wrapped + + +def _report_recorded_frames(directory: str, every: int, dt: float) -> None: + """抓帧结束后提示怎么合成视频（帧是按仿真时间等间隔抓的，fps = 1/(every*dt)）。""" + frames = sorted(name for name in os.listdir(directory) if name.endswith(".png")) + if not frames: + print(f"[eval_kit] 抓帧目录 {directory} 里没有帧（视口没渲染？）", flush=True) + return + fps = 1.0 / (every * dt) if every > 0 and dt > 0 else 10.0 + print(f"[eval_kit] 抓到 {len(frames)} 帧：{directory}", flush=True) + print( + f"[eval_kit] 合成视频（按仿真速度播放）：cd {directory} && " + f"ffmpeg -y -framerate {fps:.2f} -pattern_type glob -i '*.png' " + f"-pix_fmt yuv420p -c:v libx264 -vf 'pad=ceil(iw/2)*2:ceil(ih/2)*2' /tmp/eval.mp4", + flush=True, + ) + + +def _local_ipv4_addresses() -> list[str]: + """本机的非回环 IPv4 地址，用来告诉用户 WebRTC 客户端该连哪里。""" + import socket + + addresses: list[str] = [] + try: + for info in socket.getaddrinfo(socket.gethostname(), None, socket.AF_INET): + address = info[4][0] + if address not in addresses: + addresses.append(address) + except OSError: + pass + try: + # 不真的发包：只是让内核按路由表挑一个默认出口地址 + with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as probe: + probe.connect(("8.8.8.8", 80)) + address = probe.getsockname()[0] + if address not in addresses: + addresses.append(address) + except OSError: + pass + return [address for address in addresses if not address.startswith("127.")] + + +def _print_livestream_info(livestream: int) -> None: + """打印客户端该连到哪里。要在 Kit 起来之后调用，端口从 carb 设置里读实际值。""" + import carb.settings + + port = carb.settings.get_settings().get("/app/livestream/port") + addresses = _local_ipv4_addresses() + if livestream == 1: + public_ip = os.environ.get("PUBLIC_IP") or (addresses[0] if addresses else "127.0.0.1") + print(f"[eval_kit] WebRTC 推流已启用（--livestream 1，公网，PUBLIC_IP={public_ip}）", flush=True) + else: + print("[eval_kit] WebRTC 推流已启用（--livestream 2，局域网直连）", flush=True) + if addresses: + targets = " ".join(f"{address}:{port}" for address in addresses) + print(f"[eval_kit] 客户端连接地址：{targets}", flush=True) + else: + print(f"[eval_kit] 客户端连接地址：<本机 IP>:{port}", flush=True) + print( + "[eval_kit] 本子进程结束后推流会断开（每个阶段/种子都是一个新进程），" + "客户端重连即可", + flush=True, + ) + + +def _hold_before_rollout(sim: Any, seconds: float, *, start_height_note: str) -> None: + """停顿 ``seconds`` 秒再开始 rollout：画面继续推流、物理冻结。 + + 用 ``sim.render()`` 而不是 ``time.sleep``：独立进程模式下 Kit 的更新循环由我们 + 调用 ``app.update()`` 驱动，光 sleep 推流画面会冻住。``render()`` 内部会先把 + ``/app/player/playSimulations`` 关掉再 update，因此画面在刷新、时间不前进。 + 另外把 timeline 暂停，避免物理在等待期间自己走掉（那会让评估起点不是刚重置的状态）。 + """ + import select + import time as _time + + sim.pause() + started = _time.monotonic() + deadline = started + seconds + # 明确说明「悬空不动」是正常的：这一窗物理是冻结的，机器人正停在法规的起始位姿上 + print( + f"[eval_kit] 画面已就绪，先停 {seconds:.0f} s 再开始 rollout —— 期间物理是冻结的，" + f"机器人在起始位姿上不动{start_height_note}，看起来「悬空不动」是正常的。", + flush=True, + ) + if sys.stdin is not None and sys.stdin.isatty(): + print("[eval_kit] 在终端按回车可立即开始 rollout", flush=True) + next_report = started + 5.0 + while _time.monotonic() < deadline: + sim.render() + if sys.stdin is not None and sys.stdin.isatty() and select.select([sys.stdin], [], [], 0)[0]: + sys.stdin.readline() + break + if _time.monotonic() >= next_report: + print(f"[eval_kit] 剩余 {max(0.0, deadline - _time.monotonic()):.0f} s（物理仍冻结）", flush=True) + next_report += 5.0 + sim.play() + print(f"[eval_kit] 停顿 {_time.monotonic() - started:.1f} s 结束，开始 rollout（画面从这里才动）", flush=True) + + +def _format_early_reset_warning(early: dict, protocol: EvalProtocol, min_terrain_size: float) -> str: + """把早退统计变成可执行的建议。""" + import json as _json + + lines = [ + "[警告] 有环境在 episode 中途被重置，路程/进度类指标可能被污染：" + f"{_json.dumps(early, ensure_ascii=False)}" + ] + if protocol.enforce_terrain_bounds: + lines.append(" 最常见的原因是 IsaacLab 的 terrain_out_of_bounds 终止项：它比较机器人的") + lines.append(" 世界坐标与整张地图的范围，而不是机器人自己那块地砖，于是最外侧的格子") + lines.append(f" 从第 0 步起就被判越界。当前地砖 {protocol.terrain_size_m} m，") + lines.append(f" 要保留该终止项需要 >= {min_terrain_size:.1f} m。三种处置：") + lines.append(f" 1) --terrain-size {max(min_terrain_size, protocol.terrain_size_m):.1f} 保留终止项，消除误判") + lines.append(" 2) --no-terrain-bounds 去掉该终止项") + lines.append(" 3) --protocol standard_v1_clean 等价于方案 1") + return "\n".join(lines) + + +def _stage_argv(stage: str, seed: int, out_path: str) -> list[str]: + """构造子进程命令行：在原始参数上追加阶段参数。""" + return [ + sys.executable, + os.path.abspath(__file__), + *_ORIGINAL_ARGV, + "--stage", stage, + "--stage-seed", str(seed), + "--stage-out", out_path, + ] + + +def _spawn_stage(stage: str, seed: int) -> dict[str, object]: + """在子进程里跑一个「一个环境」的阶段，读回它写的 JSON。 + + 为什么必须开子进程：IsaacLab 在本配置下一个进程只能创建一个环境。第一个环境 + 建好之后，``ROUGH_TERRAINS_CFG`` 这类模块级单例会被运行时长出自引用，之后 + 任何 cfg 构造（新实例 ``env_cfg_cls()``、开场存的 pristine 深拷贝、重新深拷贝） + 都会在 ``cfg.validate()`` 里 RecursionError。三种补救方式都实测失败，所以只能 + 一个环境一个进程。 + """ + handle = tempfile.NamedTemporaryFile(prefix=f"eval_kit_{stage}_", suffix=".json", delete=False) + handle.close() + out_path = handle.name + argv = _stage_argv(stage, seed, out_path) + print(f"[eval_kit] 子进程 stage={stage} seed={seed}", flush=True) + completed = subprocess.run(argv) + if completed.returncode != 0: + os.unlink(out_path) + raise SystemExit(f"子进程 stage={stage} seed={seed} 失败（退出码 {completed.returncode}）") + try: + with open(out_path, encoding="utf-8") as file_handle: + return json.load(file_handle) + finally: + os.unlink(out_path) + + +def _recorder_report(recorder) -> dict[str, object]: + return { + "feet": recorder.feet_names if recorder else [], + "support": recorder.support_names if recorder else [], + "joints": recorder.joint_names if recorder else [], + "mass_kg": recorder.mass_kg if recorder else None, + "warnings": recorder.warnings if recorder else [], + } + + +def _run_parent(args_cli: argparse.Namespace) -> None: + """父进程：不启动仿真，只做编排与合并。每个环境交给一个子进程。""" + import numpy as np + + from eval_kit.core.aggregation import aggregate_runs, build_upper_limit_matrix + from eval_kit.core.report import filter_metrics, format_aggregate_summary, format_metric_result + from eval_kit.core.report import format_rollout_summary, format_upper_limit_summary + + profile = get_profile(args_cli.profile) + protocol = _apply_overrides(get_protocol(args_cli.protocol), args_cli) + spec = get_terrain_spec(protocol.terrain_spec) + if args_cli.robot_mass is not None: + profile = replace(profile, mass_kg=args_cli.robot_mass) + + if not args_cli.checkpoint: + raise SystemExit( + "编排进程需要 --checkpoint 才能派发子进程（父进程不加载 hydra，" + "因此无法用 --load-run 按实验名查找）。请传入 checkpoint 路径。" + ) + resume_path = os.path.abspath(os.path.expanduser(args_cli.checkpoint)) + if not os.path.isfile(resume_path): + raise SystemExit(f"找不到 checkpoint：{resume_path}") + + print("[eval_kit] 协议预设:", args_cli.protocol) + print("[eval_kit] profile :", profile.name) + print("[eval_kit] checkpoint:", resume_path) + print("[eval_kit] 口径 :", json.dumps(protocol.to_dict(), ensure_ascii=False)) + print("[eval_kit] 地形规格 :", spec.to_dict()) + print("[eval_kit] 种子 :", list(protocol.seeds), "（每个种子一个子进程）") + + result: dict[str, object] = { + "suite": "eval_kit", + "task": args_cli.task, + "eval_type": args_cli.eval_type, + "checkpoint": resume_path, + "profile": profile.to_dict(), + "protocol": protocol.to_dict(), + "terrain_spec": spec.to_dict(), + "seeds": list(protocol.seeds), + "runner": "one-env-per-process", + } + + wants_upper = args_cli.eval_type in ("upper_limits", "all") + + if wants_upper: + terrain_types = _terrain_types_for(args_cli, spec) + per_seed: list[dict[str, object]] = [] + pooled: dict[str, list[float]] = {"success": [], "max_progress": [], "max_tilt": []} + pooled_levels: list[np.ndarray] = [] + pooled_types: list[np.ndarray] = [] + recorder_report: dict[str, object] = {} + + for seed in protocol.seeds: + data = _spawn_stage("upper", seed) + recorder_report = data["recorder"] # type: ignore[assignment] + per_env = data["per_env"] # type: ignore[assignment] + level_ids = np.asarray(data["level_ids"]) # type: ignore[arg-type] + type_ids = np.asarray(data["type_ids"]) # type: ignore[arg-type] + for key in pooled: + pooled[key].extend(per_env[key]) # type: ignore[index] + pooled_levels.append(level_ids) + pooled_types.append(type_ids) + matrix = build_upper_limit_matrix( + per_env, level_ids, type_ids, terrain_types, spec, protocol.num_levels # type: ignore[arg-type] + ) + per_seed.append( + { + "seed": seed, + "scores": matrix["scores"], + "matrix": matrix["matrix"], + "metrics": data["metrics"], + "early_resets": data["early_resets"], + } + ) + if args_cli.verbose_terrain: + print(f"[eval_kit] seed {seed}: " + " ".join(f"{k}={v:.2f}" for k, v in matrix["scores"].items())) + + pooled_matrix = build_upper_limit_matrix( + pooled, np.concatenate(pooled_levels), np.concatenate(pooled_types), terrain_types, spec, protocol.num_levels + ) + result["terrain_types"] = terrain_types + result["upper_limits"] = pooled_matrix + result["upper_limits_per_seed"] = per_seed + result["recorder"] = recorder_report + print() + print(format_upper_limit_summary(pooled_matrix)) + first_metrics = per_seed[0]["metrics"] # type: ignore[index] + print() + print(format_rollout_summary(first_metrics)) # type: ignore[arg-type] + + for entry in per_seed: + early = entry["early_resets"] # type: ignore[index] + if early and early["polluted"]: # type: ignore[index] + print() + print(_format_early_reset_warning(early, protocol, protocol.minimum_terrain_size_for_bounds())) # type: ignore[index,arg-type] + break + + if args_cli.eval_type == "all" and args_cli.include_summary: + runs = [] + for seed in protocol.seeds: + data = _spawn_stage("plain", seed) + runs.append((seed, data["metrics"])) # type: ignore[arg-type] + result["summary"] = aggregate_runs(runs) # type: ignore[arg-type] + print() + print("[混合地形汇总]") + print(format_aggregate_summary(result["summary"])) # type: ignore[arg-type] + if runs: + result["summary_metrics"] = filter_metrics(runs[0][1], args_cli.eval_type) + else: + runs = [] + recorder_report = {} + early_any = None + for seed in protocol.seeds: + data = _spawn_stage("plain", seed) + runs.append((seed, data["metrics"])) # type: ignore[arg-type] + recorder_report = data["recorder"] # type: ignore[assignment] + early_any = data["early_resets"] + aggregate = aggregate_runs(runs) # type: ignore[arg-type] + result["aggregate"] = aggregate + result["metrics"] = filter_metrics(runs[0][1], args_cli.eval_type) + result["recorder"] = recorder_report + result["early_resets"] = early_any + print() + print(format_metric_result(args_cli.eval_type, runs[0][1])) + if len(runs) > 1: + print() + print(format_aggregate_summary(aggregate)) + if early_any and early_any["polluted"]: # type: ignore[index] + print() + print(_format_early_reset_warning(early_any, protocol, protocol.minimum_terrain_size_for_bounds())) # type: ignore[index,arg-type] + + warnings = result.get("recorder", {}).get("warnings") # type: ignore[union-attr] + if warnings: + print("\n[形态解析警告]") + for warning in warnings: # type: ignore[union-attr] + print(" -", warning) + + output_path = _resolve_output_path(args_cli, resume_path, args_cli.eval_type) + with open(output_path, "w", encoding="utf-8") as handle: + json.dump(result, handle, indent=2, ensure_ascii=False) + print(f"\n[eval_kit] 结果已保存：{output_path}") + if args_cli.print_json: + print(json.dumps(result, indent=2, ensure_ascii=False)) + + +def _run_child( + args_cli: argparse.Namespace, + *, + viewer_env_index: int | None, + free_camera: bool, + livestream: int, + pause_s: float, +) -> None: + """子进程：启动仿真，只建一个环境，跑完把结果写进 ``--stage-out``。 + + ``viewer_env_index`` / ``free_camera`` / ``livestream`` / ``pause_s`` 都由 + :func:`main` 在 ``AppLauncher`` 之前算好传进来：构造 ``AppLauncher`` 会改写 + ``headless`` 并把 ``livestream`` 从 namespace 里 pop 掉，之后再读就晚了。 + """ + print(f"[eval_kit][child] stage={args_cli.stage} seed={args_cli.stage_seed} 开始", flush=True) + import numpy as np + + from isaaclab_rl.rsl_rl import RslRlOnPolicyRunnerCfg + from isaaclab_tasks.utils.hydra import hydra_task_config + + import rl_training # noqa: F401 触发任务注册 + import rl_training.tasks # noqa: F401 + + from eval_kit.core.aggregation import build_upper_limit_matrix + from eval_kit.core.metrics import compute_metrics + from eval_kit.isaaclab import env_setup + from eval_kit.isaaclab.policy import load_policy, make_env, resolve_checkpoint + from eval_kit.isaaclab.recorder import build_recorder, record_rollout, summarize_early_resets + + profile = get_profile(args_cli.profile) + # 子进程只跑一个种子：种子的覆盖由父进程通过 --stage-seed 传进来 + protocol = replace(_apply_overrides(get_protocol(args_cli.protocol), args_cli), seeds=(args_cli.stage_seed,)) + spec = get_terrain_spec(protocol.terrain_spec) + if args_cli.robot_mass is not None: + profile = replace(profile, mass_kg=args_cli.robot_mass) + + @hydra_task_config(args_cli.task, args_cli.agent) + def _child(env_cfg, agent_cfg: RslRlOnPolicyRunnerCfg): + experiment_name = args_cli.experiment_name or agent_cfg.experiment_name + log_root_path = os.path.abspath(os.path.join("logs", "rsl_rl", experiment_name)) + resume_path = resolve_checkpoint(log_root_path, args_cli.load_run, args_cli.checkpoint) + + cfg = copy.deepcopy(env_cfg) + if args_cli.stage == "upper": + terrain_types = _terrain_types_for(args_cli, spec) + env_setup.configure_upper_limit_env( + cfg, spec, protocol, terrain_types, + seed=args_cli.stage_seed, device=args_cli.device, + viewer_env_index=viewer_env_index, free_camera=free_camera, + ) + else: + terrain_types = [] + env_setup.configure_eval_env( + cfg, spec, protocol, num_envs=args_cli.num_envs, + device=args_cli.device, seed=args_cli.stage_seed, + viewer_env_index=viewer_env_index, free_camera=free_camera, + ) + + if viewer_env_index is not None: + num_envs = int(cfg.scene.num_envs) + if not (0 <= viewer_env_index < num_envs): + raise SystemExit( + f"--view-env {viewer_env_index} 超出范围：本次共 {num_envs} 个环境（编号 0~{num_envs - 1}）。" + ) + if terrain_types: + per_type = protocol.num_levels * protocol.envs_per_cell + type_index, remainder = divmod(viewer_env_index, per_type) + level_index = remainder // protocol.envs_per_cell + print( + f"[eval_kit] 视口跟随 env {viewer_env_index} → 地形 {terrain_types[type_index]}" + f"，难度等级 L{level_index + 1}（共 {num_envs} 个环境）", + flush=True, + ) + else: + print(f"[eval_kit] 视口跟随 env {viewer_env_index}（共 {num_envs} 个环境）", flush=True) + + env = make_env(args_cli.task, cfg, agent_cfg.clip_actions) + recorder = build_recorder(env, profile) + try: + # 策略跟仿真放在同一张卡上：默认的 --device 也会覆盖 agent 配置里的 + # device，否则在别人占着 cuda:0 的机器上会莫名其妙 OOM。 + policy = load_policy(env, agent_cfg, resume_path, device=args_cli.device) + if args_cli.record_frames: + policy = _with_frame_capture(policy, args_cli.record_frames, int(args_cli.record_every)) + if args_cli.stage == "upper": + level_ids, type_ids = env_setup.assign_upper_limit_origins(env, spec, protocol, terrain_types) + if free_camera: + _set_free_camera_overview(env) + if pause_s > 0: + _hold_before_rollout( + env.unwrapped.sim, + pause_s, + start_height_note=( + f"（本次起始高度被抬高了 {protocol.start_z_offset_m:.2f} m，" + f"所以是离地悬着的）" + if protocol.start_z_offset_m > 0 + else "" + ), + ) + trace = record_rollout(env, policy, protocol, profile, recorder) + if args_cli.record_frames: + _report_recorded_frames(args_cli.record_frames, int(args_cli.record_every), protocol.dt) + metrics = compute_metrics( + trace, profile, protocol, include_per_env=(args_cli.stage == "upper") + ) + early = summarize_early_resets(trace) + finally: + env.close() + + if args_cli.save_trace: + os.makedirs(args_cli.save_trace, exist_ok=True) + tag = "upper" if args_cli.stage == "upper" else args_cli.eval_type + trace.save(os.path.join(args_cli.save_trace, f"{tag}_seed{args_cli.stage_seed}.npz")) + + payload: dict[str, object] = { + "seed": args_cli.stage_seed, + "metrics": metrics, + "early_resets": early, + "recorder": _recorder_report(recorder), + } + if args_cli.stage == "upper": + payload["per_env"] = metrics.pop("_per_env") + payload["level_ids"] = level_ids.cpu().numpy().tolist() + payload["type_ids"] = type_ids.cpu().numpy().tolist() + with open(args_cli.stage_out, "w", encoding="utf-8") as handle: + json.dump(payload, handle) + + _child() + print("[eval_kit][child] 完成", flush=True) + + +def main() -> None: + parser = _build_parser() + AppLauncher.add_app_launcher_args(parser) + args_cli, hydra_args = parser.parse_known_args() + + # 父进程只做编排：它不建环境，因此也不需要启动仿真。 + if args_cli.stage == "": + if _livestream_value(args_cli) > 0 and not args_cli.headless: + print( + "[eval_kit] 注意：--livestream 会自动进入 headless 模式（推流由服务端渲染），" + "不需要也不应该再传 --headless。", + flush=True, + ) + _run_parent(args_cli) + return + + # 这三个值必须在 AppLauncher 之前定下来（见 _viewer_env_index 的说明）。 + livestream = _livestream_value(args_cli) + viewer_env_index = _viewer_env_index(args_cli) + free_camera = bool(getattr(args_cli, "free_camera", False)) and _display_enabled(args_cli) + pause_s = args_cli.pause_s if args_cli.pause_s is not None else (30.0 if livestream > 0 else 0.0) + if args_cli.free_camera and not free_camera: + print( + "[eval_kit] 提示：--free-camera 只在有画面时（本地 GUI 或 --livestream）有效，" + "本次是 headless 且未推流，忽略该参数。", + flush=True, + ) + if livestream > 0: + _allow_client_resize(args_cli) + + # 子进程：启动仿真，跑一个环境，然后退出。 + sys.argv = [sys.argv[0]] + hydra_args + simulation_app = AppLauncher(args_cli).app + try: + if livestream > 0: + _print_livestream_info(livestream) + _run_child( + args_cli, + viewer_env_index=viewer_env_index, + free_camera=free_camera, + livestream=livestream, + pause_s=pause_s, + ) + except BaseException: + # 关键：不能走 simulation_app.close()。它会调用 os._exit(0)，把「子进程失败」 + # 伪装成「成功」，并且丢掉还没 flush 的 stderr —— 于是父进程只看到退出码 0 + # 和一个空的 JSON，真实报错（例如资产配置里的 SyntaxError）完全看不到。 + # 先把报错刷出去，再以非零码强制退出，父进程才会把真实错误报上来。 + import traceback + + traceback.print_exc() + for stream in (sys.stdout, sys.stderr): + stream.flush() + os._exit(1) + for stream in (sys.stdout, sys.stderr): + stream.flush() + simulation_app.close() + + + +if __name__ == "__main__": + main() diff --git a/scripts/reinforcement_learning/rsl_rl/play.py b/scripts/reinforcement_learning/rsl_rl/play.py new file mode 100644 index 0000000..3e387c3 --- /dev/null +++ b/scripts/reinforcement_learning/rsl_rl/play.py @@ -0,0 +1,285 @@ +# 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 + +# 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 +try: + from isaaclab.utils.pretrained_checkpoint import get_published_pretrained_checkpoint +except ModuleNotFoundError: + # Isaac Lab 2.3.x and earlier keep this helper under isaaclab_rl. + from isaaclab_rl.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 + +from rl_utils import camera_follow + +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..6ab8110 --- /dev/null +++ b/source/rl_training/rl_training/assets/deeprobotics.py @@ -0,0 +1,125 @@ +# 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", + usd_dir="/home/suyufei/tmp/IsaacLab_m20_wheel220", + 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.32), + 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..5f120fd --- /dev/null +++ b/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/config/wheeled/deeprobotics_m20/__init__.py @@ -0,0 +1,68 @@ +# 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", + }, +) + +gym.register( + id="Rough-Deeprobotics-M20-Bridge-v0", + entry_point="isaaclab.envs:ManagerBasedRLEnv", + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": f"{__name__}.bridge_env_cfg:DeeproboticsM20BridgeEnvCfg", + "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:DeeproboticsM20BridgePPORunnerCfg", + "cusrl_cfg_entry_point": f"{agents.__name__}.cusrl_ppo_cfg:DeeproboticsM20RoughTrainerCfg", + }, +) + +gym.register( + id="Smooth-Deeprobotics-M20-v0", + entry_point="isaaclab.envs:ManagerBasedRLEnv", + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": f"{__name__}.smooth_env_cfg:DeeproboticsM20SmoothEnvCfg", + "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:DeeproboticsM20SmoothPPORunnerCfg", + "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..7bce2cd --- /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,78 @@ +# 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" + + +@configclass +class DeeproboticsM20BridgePPORunnerCfg(DeeproboticsM20RoughPPORunnerCfg): + """Single-side bridge fine-tuning, forward-only command.""" + + def __post_init__(self): + super().__post_init__() + + self.max_iterations = 5000 + self.experiment_name = "huoyi_bridge" + + +@configclass +class DeeproboticsM20SmoothPPORunnerCfg(DeeproboticsM20RoughPPORunnerCfg): + """Smooth random-undulation fine-tuning, forward-only command.""" + + def __post_init__(self): + super().__post_init__() + + self.max_iterations = 5000 + self.experiment_name = "huoyi_smooth" diff --git a/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/config/wheeled/deeprobotics_m20/bridge_env_cfg.py b/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/config/wheeled/deeprobotics_m20/bridge_env_cfg.py new file mode 100644 index 0000000..612abc7 --- /dev/null +++ b/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/config/wheeled/deeprobotics_m20/bridge_env_cfg.py @@ -0,0 +1,76 @@ +# Copyright (c) 2025 Deep Robotics +# SPDX-License-Identifier: BSD 3-Clause + +"""Single-side bridge task: left legs on a long raised bridge, right legs on flat ground.""" + +from isaaclab.terrains import TerrainGeneratorCfg +from isaaclab.utils import configclass + +import rl_training.tasks.manager_based.locomotion.velocity.mdp as mdp + +from .rough_env_cfg import DeeproboticsM20RoughEnvCfg + + +@configclass +class DeeproboticsM20BridgeEnvCfg(DeeproboticsM20RoughEnvCfg): + """Stage 1 policy on a long single-side bridge, forward-only command.""" + + def __post_init__(self): + super().__post_init__() + + # ------------------------------Terrain------------------------------ + # One long bridge cell type. The bridge runs along +x, so the robot can + # keep walking forward without leaving the split-level configuration. + self.scene.terrain.terrain_generator = TerrainGeneratorCfg( + size=(12.0, 4.0), + border_width=2.0, + num_rows=10, + num_cols=20, + horizontal_scale=0.1, + vertical_scale=0.005, + slope_threshold=0.75, + use_cache=False, + curriculum=True, + sub_terrains={ + "single_side_bridge": mdp.SingleSideBridgeTerrainCfg( + proportion=1.0, + bridge_height_range=(0.06, 0.14), + bridge_width=1.5, + ) + }, + ) + self.scene.terrain.max_init_terrain_level = 0 + + # ------------------------------Spawn------------------------------ + # Spawn exactly on the bridge edge: y=0 and yaw=0 put the left wheels + # on the bridge and the right wheels on the flat ground. + self.events.randomize_reset_base.params = { + "pose_range": { + "x": (-1.0, 1.0), + "y": (0.0, 0.0), + "z": (0.30, 0.30), + "roll": (-0.03, 0.03), + "pitch": (-0.03, 0.03), + "yaw": (0.0, 0.0), + }, + "velocity_range": { + "x": (-0.1, 0.1), + "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), + }, + } + + # ------------------------------Commands------------------------------ + # Forward-only: no lateral and no yaw command. + 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.2, 0.8) + self.commands.base_velocity.ranges.lin_vel_y = (0.0, 0.0) + self.commands.base_velocity.ranges.ang_vel_z = (0.0, 0.0) + + # Keep the same 33-dim policy / 8-dim action layout as the other tasks. + self.disable_zero_weight_rewards() 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..47653ce --- /dev/null +++ b/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/config/wheeled/deeprobotics_m20/rough_env_cfg.py @@ -0,0 +1,475 @@ +# 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 SceneEntityCfg +from isaaclab.managers import TerminationTermCfg as DoneTerm +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.11, + "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.11, + "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, + }, + ) + + knee_inward_fold_penalty = RewTerm( + func=mdp.knee_inward_fold_penalty, + weight=0.0, + params={ + "front_asset_cfg": SceneEntityCfg("robot", joint_names=""), + "hind_asset_cfg": SceneEntityCfg("robot", joint_names=""), + "max_inward_tilt": 0.0, + }, + ) + + +@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", + ] + front_knee_joint_names = ["fl_knee_joint", "fr_knee_joint"] + hind_knee_joint_names = ["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 + + # ------------------------------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.05 + self.rewards.flat_orientation_l2.weight = -0.35 + self.rewards.roll_l2.weight = -0.45 + 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 = -3e-6 + 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 + # Inward fold limit: with the knee link vertical at q=0, front knees fold inward for + # q < 0 and hind knees for q > 0. Do not allow the legs to pass the vertical + # (straight-down) pose on the inward side. + self.rewards.knee_inward_fold_penalty.weight = -1.0 + self.rewards.knee_inward_fold_penalty.params["front_asset_cfg"] = SceneEntityCfg( + "robot", joint_names=self.front_knee_joint_names + ) + self.rewards.knee_inward_fold_penalty.params["hind_asset_cfg"] = SceneEntityCfg( + "robot", joint_names=self.hind_knee_joint_names + ) + self.rewards.knee_inward_fold_penalty.params["max_inward_tilt"] = 0.0 + 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.0 + 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.015 + + # 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 = DoneTerm( + func=mdp.bad_orientation, + params={"limit_angle": 0.9, "asset_cfg": SceneEntityCfg("robot")}, + ) + + # ------------------------------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 = 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.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.35 + self.rewards.roll_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/config/wheeled/deeprobotics_m20/smooth_env_cfg.py b/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/config/wheeled/deeprobotics_m20/smooth_env_cfg.py new file mode 100644 index 0000000..222e0a7 --- /dev/null +++ b/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/config/wheeled/deeprobotics_m20/smooth_env_cfg.py @@ -0,0 +1,79 @@ +# Copyright (c) 2025 Deep Robotics +# SPDX-License-Identifier: BSD 3-Clause + +"""Smooth random-undulation task: train both pitch and roll to stay at zero.""" + +from isaaclab.terrains import TerrainGeneratorCfg +from isaaclab.utils import configclass + +import rl_training.tasks.manager_based.locomotion.velocity.mdp as mdp + +from .rough_env_cfg import DeeproboticsM20RoughEnvCfg + + +@configclass +class DeeproboticsM20SmoothEnvCfg(DeeproboticsM20RoughEnvCfg): + """Stage 1 policy on a smooth, randomly undulating terrain, forward-only command.""" + + def __post_init__(self): + super().__post_init__() + + # ------------------------------Terrain------------------------------ + # Low-frequency sinusoids only: no steps, no sharp edges. + self.scene.terrain.terrain_generator = TerrainGeneratorCfg( + size=(8.0, 8.0), + border_width=2.0, + num_rows=10, + num_cols=20, + horizontal_scale=0.05, + vertical_scale=0.005, + slope_threshold=None, + use_cache=False, + curriculum=True, + sub_terrains={ + "smooth_random": mdp.SmoothRandomTerrainCfg( + proportion=1.0, + border_width=0.25, + height_range=(0.10, 0.20), + wavelength_range=(1.0, 2.2), + num_waves=3, + ) + }, + ) + self.scene.terrain.max_init_terrain_level = 9 + + # ------------------------------Spawn------------------------------ + self.events.randomize_reset_base.params = { + "pose_range": { + "x": (-2.0, 2.0), + "y": (-2.0, 2.0), + "z": (0.30, 0.30), + "roll": (-0.05, 0.05), + "pitch": (-0.05, 0.05), + "yaw": (-3.14, 3.14), + }, + "velocity_range": { + "x": (-0.1, 0.1), + "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), + }, + } + + # Commands keep the original Stage-1 distribution inherited from + # DeeproboticsM20RoughEnvCfg: x in [-2, 2] m/s, y = 0, yaw in [-1.5, 1.5] + # rad/s, heading command enabled, 2% standing and 100% heading envs. + + # ------------------------------Rewards------------------------------ + # Split the flat-orientation penalty into true Euler pitch and roll in + # radians. Giving -0.4 to each keeps the total strength of the previous + # -0.8 flat_orientation_l2 term while letting the two axes be tuned + # independently. + self.rewards.flat_orientation_l2.weight = 0.0 + self.rewards.pitch_l2.weight = -0.4 + self.rewards.roll_l2.weight = -0.4 + self.rewards.ang_vel_xy_l2.weight = -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..216d286 --- /dev/null +++ b/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/mdp/__init__.py @@ -0,0 +1,22 @@ +# 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 +from .terrain 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..160754f --- /dev/null +++ b/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/mdp/rewards.py @@ -0,0 +1,1027 @@ +# 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 knee_inward_fold_penalty( + env: ManagerBasedRLEnv, + front_asset_cfg: SceneEntityCfg, + hind_asset_cfg: SceneEntityCfg, + max_inward_tilt: float = 0.0, +) -> torch.Tensor: + """Penalize front/hind knee joints that fold inward past ``max_inward_tilt``. + + The knee link points straight down when its joint angle is zero, so the signed joint + position is the leg's tilt from vertical (all four knee axes are ``-Y`` in the URDF). + Front legs fold inward when the angle decreases, hind legs fold inward when it + increases. ``max_inward_tilt`` is the largest inward tilt allowed from the vertical + (straight-down) pose: 0.0 stops the leg at vertical, ``pi / 2`` would allow it to + reach horizontal. Only the inward direction is limited; outward folding is left to + the regular joint position limits. The penalty is quadratic in the excess angle and + is zero while the joint stays inside the limit. + """ + front_asset: Articulation = env.scene[front_asset_cfg.name] + hind_asset: Articulation = env.scene[hind_asset_cfg.name] + front_pos = front_asset.data.joint_pos[:, front_asset_cfg.joint_ids] + hind_pos = hind_asset.data.joint_pos[:, hind_asset_cfg.joint_ids] + # front: q >= -max_inward_tilt, hind: q <= max_inward_tilt + front_excess = torch.clamp(-max_inward_tilt - front_pos, min=0.0) + hind_excess = torch.clamp(hind_pos - max_inward_tilt, min=0.0) + return torch.sum(front_excess.square() + hind_excess.square(), dim=1) + + +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 roll_l2(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: + """Penalize body roll angle in radians (squared).""" + asset: RigidObject = env.scene[asset_cfg.name] + roll, _, _ = math_utils.euler_xyz_from_quat(asset.data.root_quat_w) + return torch.square(roll) + + +def pitch_l2(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: + """Penalize body pitch angle in radians (squared).""" + asset: RigidObject = env.scene[asset_cfg.name] + _, pitch, _ = math_utils.euler_xyz_from_quat(asset.data.root_quat_w) + return torch.square(pitch) + + +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/mdp/terrain.py b/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/mdp/terrain.py new file mode 100644 index 0000000..755ce1d --- /dev/null +++ b/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/mdp/terrain.py @@ -0,0 +1,107 @@ +"""Custom terrains for the FireAnt locomotion tasks.""" + +from __future__ import annotations + +import numpy as np +import trimesh + +from isaaclab.terrains import HfTerrainBaseCfg, SubTerrainBaseCfg +from isaaclab.terrains.height_field.utils import height_field_to_mesh +from isaaclab.utils import configclass + + +def single_side_bridge_terrain( + difficulty: float, cfg: "SingleSideBridgeTerrainCfg" +) -> tuple[list[trimesh.Trimesh], np.ndarray]: + """Long raised bridge on the +y half of the cell; flat ground on the -y half. + + The bridge edge is exactly at the cell center line (local ``y = size[1] / 2``), + so a robot spawned at the cell center with zero yaw has its left wheels on the + bridge and its right wheels on the flat ground. + """ + height = cfg.bridge_height_range[0] + difficulty * ( + cfg.bridge_height_range[1] - cfg.bridge_height_range[0] + ) + + # Ground slab: top surface at z=0, extends over the whole cell. + ground = trimesh.creation.box( + extents=(cfg.size[0], cfg.size[1], 0.05), + transform=trimesh.transformations.translation_matrix( + (cfg.size[0] / 2.0, cfg.size[1] / 2.0, -0.025) + ), + ) + + # Bridge: full cell length along x, raised strip on the +y side. + bridge = trimesh.creation.box( + extents=(cfg.size[0], cfg.bridge_width, height), + transform=trimesh.transformations.translation_matrix( + ( + cfg.size[0] / 2.0, + cfg.size[1] / 2.0 + cfg.bridge_width / 2.0, + height / 2.0, + ) + ), + ) + + origin = np.array((cfg.size[0] / 2.0, cfg.size[1] / 2.0, 0.0)) + return [ground, bridge], origin + + +@configclass +class SingleSideBridgeTerrainCfg(SubTerrainBaseCfg): + """Configuration for a long single-side bridge terrain.""" + + function = single_side_bridge_terrain + + bridge_height_range: tuple[float, float] = (0.06, 0.14) + """Bridge height at the start and end of the difficulty curriculum (m).""" + + bridge_width: float = 1.5 + """Width of the raised bridge on the +y side of the cell (m).""" + + +@height_field_to_mesh +def smooth_random_terrain(difficulty: float, cfg: "SmoothRandomTerrainCfg") -> np.ndarray: + """Smooth low-frequency random height field. + + The surface is a sum of a few low-frequency sinusoids. It has no steps or + sharp edges, so the robot has to keep its body level over gentle, random + pitch and roll disturbances. + """ + width_pixels = int(cfg.size[0] / cfg.horizontal_scale) + length_pixels = int(cfg.size[1] / cfg.horizontal_scale) + x = np.arange(width_pixels, dtype=np.float64) * cfg.horizontal_scale + y = np.arange(length_pixels, dtype=np.float64) * cfg.horizontal_scale + grid_x, grid_y = np.meshgrid(x, y, indexing="ij") + + amplitude = cfg.height_range[0] + difficulty * (cfg.height_range[1] - cfg.height_range[0]) + heights = np.zeros_like(grid_x) + for _ in range(cfg.num_waves): + direction = np.random.uniform(0.0, 2.0 * np.pi) + wavelength = np.random.uniform(*cfg.wavelength_range) + phase = np.random.uniform(0.0, 2.0 * np.pi) + weight = np.random.uniform(0.5, 1.0) + projection = grid_x * np.cos(direction) + grid_y * np.sin(direction) + heights += weight * np.sin(2.0 * np.pi * projection / wavelength + phase) + + heights -= heights.mean() + peak = float(np.max(np.abs(heights))) + if peak > 1e-9: + heights *= amplitude / peak + return np.rint(heights / cfg.vertical_scale).astype(np.int16) + + +@configclass +class SmoothRandomTerrainCfg(HfTerrainBaseCfg): + """Configuration for a smooth random-undulation terrain.""" + + function = smooth_random_terrain + + height_range: tuple[float, float] = (0.02, 0.05) + """Peak height of the undulations at the start and end of the curriculum (m).""" + + wavelength_range: tuple[float, float] = (1.5, 3.0) + """Spatial wavelength range of the undulations (m).""" + + num_waves: int = 3 + """Number of sinusoids summed to form the surface.""" 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..46a9866 --- /dev/null +++ b/source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/velocity_env_cfg.py @@ -0,0 +1,868 @@ +# 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) + roll_l2 = RewTerm(func=mdp.roll_l2, weight=0.0) + pitch_l2 = RewTerm(func=mdp.pitch_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, +)