Add standalone Huoyi training project

2026-10-13 21:25:03 +08:00
commit 1aeeb34653
69 changed files with 7312 additions and 0 deletions
@@ -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
@@ -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/*
@@ -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
@@ -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()
@@ -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
@@ -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.
@@ -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.
@@ -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 | <img src="./docs/imgs/deeprobotics_lite3.png" alt="Lite3" width="300">
 | [Deeprobotics M20](https://www.deeprobotics.cn/robot/index/lynx.html) | Rough-Deeprobotics-M20-v0 | <img src="./docs/imgs/deeprobotics_m20.png" alt="deeprobotics_m20" width="300">
 > [!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
  ```
 <details>
 <summary>Setup as Omniverse Extension (Optional, click to expand)</summary>
 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.
 </details>
 ## 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=<ENV_NAME> --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=<ENV_NAME> --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=<ENV_NAME> --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 `<path-to-isaac-lab>` with your own IsaacLab path.**
 ```json
 {
    "python.languageServer": "Pylance",
    "python.analysis.extraPaths": [
        "${workspaceFolder}/source/rl_training",
        "/<path-to-isaac-lab>/source/isaaclab",
        "/<path-to-isaac-lab>/source/isaaclab_assets",
        "/<path-to-isaac-lab>/source/isaaclab_mimic",
        "/<path-to-isaac-lab>/source/isaaclab_rl",
        "/<path-to-isaac-lab>/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)
@@ -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。
@@ -0,0 +1 @@
 1.0.0
@@ -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.
@@ -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
    }
  ]
 }
@@ -0,0 +1,292 @@
 <?xml version="1.0" encoding="utf-8"?>
 <!--
  Solenopsis fire-ant, generated from the 2026 SolidWorks STEP assembly.
  Units: metre, kilogram, radian, kg*m^2.
  Coordinate conversion: (x,y,z)_URDF = (z,x,y)_SolidWorks.
  All inertia tensors are about each rigid body's own centre of mass.
  Visual geometry uses the STEP-derived STL files in
  ../meshes_solenopsis_visual at a millimetre-to-metre scale of 0.001.
  Collision geometry remains primitive boxes and cylinders.
  The 0.127 kg gap between the supplied rigid-group masses and the supplied
  whole-robot mass is lumped into base_link at the reported base COM.  Thus
  base mass is 17.274 + 0.127 = 17.401 kg, while its reported COM and inertia
  tensor are unchanged.  No spatial distribution is invented for that mass.
 -->
 <robot name="Solenopsis_fire_ant">
  <link name="base_link">
    <inertial>
      <origin xyz="-0.000072 0.000171 0.076892" rpy="0 0 0"/>
      <mass value="17.401"/>
      <inertia ixx="0.133964838" ixy="-0.000732466" ixz="0.000157104"
               iyy="0.363696493" iyz="0.000165399" izz="0.450717599"/>
    </inertial>
    <visual>
      <geometry>
        <mesh filename="../meshes_solenopsis_visual/base_link.STL" scale="0.001 0.001 0.001"/>
      </geometry>
    </visual>
    <collision>
      <origin xyz="0 0 0.0788" rpy="0 0 0"/>
      <geometry>
        <box size="0.706 0.310 0.1714"/>
      </geometry>
    </collision>
  </link>
  <!-- Front-left.  Left-link tensors use the exact proper rotation
       (x,y,z)_link = (-z,-x,y)_SolidWorks-module. -->
  <link name="fl_knee">
    <inertial>
      <origin xyz="-0.000040 -0.001088 -0.158988" rpy="0 0 0"/>
      <mass value="0.955"/>
      <inertia ixx="0.014124479" ixy="0.000002919" ixz="0.000004851"
               iyy="0.014470094" iyz="-0.000261898" izz="0.000478899"/>
    </inertial>
    <visual>
      <geometry>
        <mesh filename="../meshes_solenopsis_visual/fl_knee.STL" scale="0.001 0.001 0.001"/>
      </geometry>
    </visual>
    <collision>
      <origin xyz="0 -0.0005 -0.15575" rpy="0 0 0"/>
      <geometry>
        <box size="0.088 0.031 0.3765"/>
      </geometry>
    </collision>
    <collision>
      <origin xyz="0 -0.02625 0" rpy="1.570796326795 0 0"/>
      <geometry>
        <cylinder length="0.0205" radius="0.01995"/>
      </geometry>
    </collision>
  </link>
  <joint name="fl_knee_joint" type="revolute">
    <origin xyz="0.117577164466 0.167 0.070" rpy="0 0 0"/>
    <parent link="base_link"/>
    <child link="fl_knee"/>
    <axis xyz="0 -1 0"/>
    <limit lower="-2.792" upper="2.809" effort="75.0" velocity="17.8024"/>
  </joint>
  <link name="fl_wheel">
    <inertial>
      <origin xyz="-0.000002 0.059326 -0.000009" rpy="0 0 0"/>
      <mass value="1.340"/>
      <inertia ixx="0.003296552" ixy="0.000000144" ixz="-0.000000077"
               iyy="0.005826562" iyz="0.000000424" izz="0.003296995"/>
    </inertial>
    <visual>
      <geometry>
        <mesh filename="../meshes_solenopsis_visual/fl_wheel.STL" scale="0.001 0.001 0.001"/>
      </geometry>
    </visual>
    <collision>
      <origin xyz="0 0.0457795 0" rpy="1.570796326795 0 0"/>
      <geometry>
        <cylinder length="0.082959" radius="0.098804"/>
      </geometry>
    </collision>
  </link>
  <joint name="fl_wheel_joint" type="continuous">
    <origin xyz="0 0 -0.300" rpy="0 0 0"/>
    <parent link="fl_knee"/>
    <child link="fl_wheel"/>
    <axis xyz="0 -1 0"/>
    <limit effort="28.68" velocity="109.956"/>
  </joint>
  <!-- Front-right.  Right-link tensors use the exact proper rotation
       (x,y,z)_link = (z,x,y)_SolidWorks-module. -->
  <link name="fr_knee">
    <inertial>
      <origin xyz="0.000040 0.001088 -0.158988" rpy="0 0 0"/>
      <mass value="0.955"/>
      <inertia ixx="0.014124479" ixy="0.000002919" ixz="-0.000004851"
               iyy="0.014470094" iyz="0.000261898" izz="0.000478899"/>
    </inertial>
    <visual>
      <geometry>
        <mesh filename="../meshes_solenopsis_visual/fr_knee.STL" scale="0.001 0.001 0.001"/>
      </geometry>
    </visual>
    <collision>
      <origin xyz="0 0.0005 -0.15575" rpy="0 0 0"/>
      <geometry>
        <box size="0.088 0.031 0.3765"/>
      </geometry>
    </collision>
    <collision>
      <origin xyz="0 0.02625 0" rpy="1.570796326795 0 0"/>
      <geometry>
        <cylinder length="0.0205" radius="0.01995"/>
      </geometry>
    </collision>
  </link>
  <joint name="fr_knee_joint" type="revolute">
    <origin xyz="0.117577164466 -0.167 0.070" rpy="0 0 0"/>
    <parent link="base_link"/>
    <child link="fr_knee"/>
    <axis xyz="0 -1 0"/>
    <limit lower="-2.792" upper="2.809" effort="75.0" velocity="17.8024"/>
  </joint>
  <link name="fr_wheel">
    <inertial>
      <origin xyz="0.000002 -0.059326 -0.000009" rpy="0 0 0"/>
      <mass value="1.340"/>
      <inertia ixx="0.003296552" ixy="0.000000144" ixz="0.000000077"
               iyy="0.005826562" iyz="-0.000000424" izz="0.003296995"/>
    </inertial>
    <visual>
      <geometry>
        <mesh filename="../meshes_solenopsis_visual/fr_wheel.STL" scale="0.001 0.001 0.001"/>
      </geometry>
    </visual>
    <collision>
      <origin xyz="0 -0.0457795 0" rpy="1.570796326795 0 0"/>
      <geometry>
        <cylinder length="0.082959" radius="0.098804"/>
      </geometry>
    </collision>
  </link>
  <joint name="fr_wheel_joint" type="continuous">
    <origin xyz="0 0 -0.300" rpy="0 0 0"/>
    <parent link="fr_knee"/>
    <child link="fr_wheel"/>
    <axis xyz="0 -1 0"/>
    <limit effort="28.68" velocity="109.956"/>
  </joint>
  <link name="hl_knee">
    <inertial>
      <origin xyz="-0.000040 -0.001088 -0.158988" rpy="0 0 0"/>
      <mass value="0.955"/>
      <inertia ixx="0.014124479" ixy="0.000002919" ixz="0.000004851"
               iyy="0.014470094" iyz="-0.000261898" izz="0.000478899"/>
    </inertial>
    <visual>
      <geometry>
        <mesh filename="../meshes_solenopsis_visual/hl_knee.STL" scale="0.001 0.001 0.001"/>
      </geometry>
    </visual>
    <collision>
      <origin xyz="0 -0.0005 -0.15575" rpy="0 0 0"/>
      <geometry>
        <box size="0.088 0.031 0.3765"/>
      </geometry>
    </collision>
    <collision>
      <origin xyz="0 -0.02625 0" rpy="1.570796326795 0 0"/>
      <geometry>
        <cylinder length="0.0205" radius="0.01995"/>
      </geometry>
    </collision>
  </link>
  <joint name="hl_knee_joint" type="revolute">
    <origin xyz="-0.117577164466 0.167 0.070" rpy="0 0 0"/>
    <parent link="base_link"/>
    <child link="hl_knee"/>
    <axis xyz="0 -1 0"/>
    <limit lower="-2.809" upper="2.792" effort="75.0" velocity="17.8024"/>
  </joint>
  <link name="hl_wheel">
    <inertial>
      <origin xyz="-0.000002 0.059326 -0.000009" rpy="0 0 0"/>
      <mass value="1.340"/>
      <inertia ixx="0.003296552" ixy="0.000000144" ixz="-0.000000077"
               iyy="0.005826562" iyz="0.000000424" izz="0.003296995"/>
    </inertial>
    <visual>
      <geometry>
        <mesh filename="../meshes_solenopsis_visual/hl_wheel.STL" scale="0.001 0.001 0.001"/>
      </geometry>
    </visual>
    <collision>
      <origin xyz="0 0.0457795 0" rpy="1.570796326795 0 0"/>
      <geometry>
        <cylinder length="0.082959" radius="0.098804"/>
      </geometry>
    </collision>
  </link>
  <joint name="hl_wheel_joint" type="continuous">
    <origin xyz="0 0 -0.300" rpy="0 0 0"/>
    <parent link="hl_knee"/>
    <child link="hl_wheel"/>
    <axis xyz="0 -1 0"/>
    <limit effort="28.68" velocity="109.956"/>
  </joint>
  <link name="hr_knee">
    <inertial>
      <origin xyz="0.000040 0.001088 -0.158988" rpy="0 0 0"/>
      <mass value="0.955"/>
      <inertia ixx="0.014124479" ixy="0.000002919" ixz="-0.000004851"
               iyy="0.014470094" iyz="0.000261898" izz="0.000478899"/>
    </inertial>
    <visual>
      <geometry>
        <mesh filename="../meshes_solenopsis_visual/hr_knee.STL" scale="0.001 0.001 0.001"/>
      </geometry>
    </visual>
    <collision>
      <origin xyz="0 0.0005 -0.15575" rpy="0 0 0"/>
      <geometry>
        <box size="0.088 0.031 0.3765"/>
      </geometry>
    </collision>
    <collision>
      <origin xyz="0 0.02625 0" rpy="1.570796326795 0 0"/>
      <geometry>
        <cylinder length="0.0205" radius="0.01995"/>
      </geometry>
    </collision>
  </link>
  <joint name="hr_knee_joint" type="revolute">
    <origin xyz="-0.117577164466 -0.167 0.070" rpy="0 0 0"/>
    <parent link="base_link"/>
    <child link="hr_knee"/>
    <axis xyz="0 -1 0"/>
    <limit lower="-2.809" upper="2.792" effort="75.0" velocity="17.8024"/>
  </joint>
  <link name="hr_wheel">
    <inertial>
      <origin xyz="0.000002 -0.059326 -0.000009" rpy="0 0 0"/>
      <mass value="1.340"/>
      <inertia ixx="0.003296552" ixy="0.000000144" ixz="0.000000077"
               iyy="0.005826562" iyz="-0.000000424" izz="0.003296995"/>
    </inertial>
    <visual>
      <geometry>
        <mesh filename="../meshes_solenopsis_visual/hr_wheel.STL" scale="0.001 0.001 0.001"/>
      </geometry>
    </visual>
    <collision>
      <origin xyz="0 -0.0457795 0" rpy="1.570796326795 0 0"/>
      <geometry>
        <cylinder length="0.082959" radius="0.098804"/>
      </geometry>
    </collision>
  </link>
  <joint name="hr_wheel_joint" type="continuous">
    <origin xyz="0 0 -0.300" rpy="0 0 0"/>
    <parent link="hr_knee"/>
    <child link="hr_wheel"/>
    <axis xyz="0 -1 0"/>
    <limit effort="28.68" velocity="109.956"/>
  </joint>
 </robot>
@@ -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.
 <span style="color: red;">**Note**</span>: 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**   | <img src="./images/lite3_high_res.png" width="200" /><br>Lite3/Lite3_urdf/urdf/Lite3_high_res.urdf | <img src="./images/lite3.png" width="200" /><br>Lite3/Lite3_urdf/urdf/Lite3.urdf |
 | **X30**     | <img src="./images/x30_high_res.png" width="200" /><br>X30/X30_urdf/urdf/X30_high_res.urdf | <img src="./images/x30.png" width="200" /><br>X30/X30_urdf/urdf/X30.urdf |
 | **M20**     | <img src="./images/m20_high_res.png" width="200" /><br>M20/M20_urdf/urdf/M20_high_res.urdf | <img src="./images/m20.png" width="200" /><br>M20/M20_urdf/urdf/M20.urdf |
 ## Contributors
 See the [Contributors](Contributors.md) page for a list of contributors.
@@ -0,0 +1,178 @@
 # 火蚁样机参数简表
 本文档依据当前火蚁仿真模型与训练配置整理：
 - URDF: `deep_robotics_model/M20/M20_urdf/urdf/M20.urdf`
 - 执行器配置: `source/rl_training/rl_training/assets/deeprobotics.py`
 - 训练配置: `source/rl_training/rl_training/tasks/manager_based/locomotion/velocity/config/wheeled/deeprobotics_m20/rough_env_cfg.py`
 ## 1. 结构形式
 火蚁当前为 8 自由度轮腿机器人
 ## 2. 质量参数
 整机仿真总质量约：23.690 kg
 | 部件 | 单个质量 | 数量 | 小计 |
 | --- | ---: | ---: | ---: |
 | 机身 base_link | 15.882 kg | 1 | 15.882 kg |
 | 膝部连杆（按 P65MAX-S 质量建模） | 0.612 kg | 4 | 2.448 kg |
 | 轮毂电机与轮胎（W190-S） | 1.340 kg | 4 | 5.360 kg |
 ## 3. 尺寸参数
 机身主体：
 长 x 宽 x 高 = 0.60 m x 0.32 m x 0.12 m
 轮子：
 半径 = 0.096 m
 直径 = 0.192 m
 轴向总尺寸 = 0.083 m（手册工程图标注 82.95 mm）
 ```
 膝部连杆简化尺寸：
 主连杆 box = 0.05 m x 0.024 m x 0.20 m
 末端圆柱半径 = 0.04 m
 末端圆柱长度 = 0.09 m
 零关节角几何尺寸：
 ```text
 轴距 = 0.60 m
 轮距 = 0.50 m
 轮心相对机身中心高度 = -0.21 m
 ```
 训练初始姿态下，前膝为 `+0.5 rad`，后膝为 `-0.5 rad`：
 ```text
 有效轴距约 = 0.801 m
 轮距 = 0.50 m
 机身中心接地高度约 = 0.280 m
 机身底部离地约 = 0.220 m
 ```
 ## 4. 电机与执行器参数
 ### 膝关节电机
 仿真限制：
 ```text
 型号 = P65MAX-S
 额定扭矩 = 21.3 Nm
 峰值扭矩 = 75.0 Nm
 速度上限 = 17.8024 rad/s ≈ 170 rpm
 ```
 IsaacLab 控制参数：
 ```text
 位置刚度 stiffness = 80.0
 阻尼 damping = 2.0
 控制延迟 = 0 ~ 5 step
 ```
 当前 URDF 与 Isaac Lab actuator 均使用相同的峰值扭矩和速度上限。
 ### 轮端电机
 仿真实际执行器限制：
 ```text
 型号 = W190-S
 额定扭矩 = 8.7 Nm
 峰值扭矩 = 28.68 Nm
 额定点转速 = 94.457 rad/s ≈ 902 rpm
 速度上限 = 109.956 rad/s ≈ 1050 rpm
 理论空载轮缘线速度 = 10.56 m/s
 ```
 URDF 中轮关节写有：
 ```text
 effort = 28.68 Nm
 velocity = 109.956 rad/s
 ```
 URDF 与 Isaac Lab actuator 的限制现已保持一致。
 ## 5. 训练需求
 Stage 1 基础运动：
 ```text
 前后速度 lin_vel_x = -2.0 ~ 2.0 m/s
 侧向速度 lin_vel_y = 0
 转向 yaw_rate = -1.5 ~ 1.5 rad/s
 ```
 Stage 2 超支撑越障：
 ```text
 前进速度 lin_vel_x = 0.25 ~ 1.2 m/s
 侧向速度 lin_vel_y = 0
 转向 yaw_rate = -0.6 ~ 0.6 rad/s
 ```
 Stage 2 主要地形难度：
 ```text
 boxes 高度 = 0.03 ~ 0.22 m
 上台阶高度 = 0.04 ~ 0.24 m
 下台阶高度 = 0.04 ~ 0.20 m
 随机粗糙高度 = 0.01 ~ 0.14 m
 ```
 ## 6. 样机设计重点
 优先级建议：
 1. 膝关节峰值扭矩优先于轮端极速。
 2. 轮端要保证低速大扭矩，避免顶障碍时堵转。
 3. 机身底部、膝部、轮架必须能承受超支撑接触冲击。
 4. 轮胎材料要保证摩擦和一定缓冲，不能只看刚性轮模型。
 5. 电机驱动器需要考虑峰值电流、连续散热和限流策略。
 ## 7. 当前仿真限制
 接触材质配置：
 | 接触对象 | 静摩擦随机范围 | 动摩擦随机范围 | 恢复系数随机范围 |
 | --- | ---: | ---: | ---: |
 | W190-S 轮子 | 0.9 ~ 1.4 | 0.6 ~ 0.9 | 0.0 ~ 0.05 |
 | base_link 与 knee links | 0.4 ~ 0.8 | 0.25 ~ 0.4 | 0.0 ~ 0.05 |
 地形静、动摩擦均为 1.0，摩擦使用乘法组合，因此轮地有效摩擦范围与轮子材质随机范围一致。火蚁任务将地形恢复系数覆盖为 0.0，因此轮地碰撞的有效恢复系数为 0。质量与惯量随机化保持启用。
 当前模型不能直接等价为最终硬件设计：
 - 几何为简化 box/cylinder，不是完整 CAD。
 - 没有真实电机扭矩-转速曲线。
 - 没有电池、电流限制、热衰减。
 - 没有减速器效率、背隙、柔性结构。
 - 没有轮胎变形模型。
 - 接触与摩擦比真实世界更理想。
 - 轮关节 URDF effort 与 IsaacLab actuator effort 不一致，训练实际以 actuator 为准。
 ## 8. 初版样机目标
 建议以如下指标作为第一版样机基准：
 ```text
 整机仿真质量：约 23.69 kg
 轮径：0.192 m
 轮距：0.50 m
 有效轴距：约 0.80 m
 膝关节峰值扭矩：75 Nm
 轮端峰值扭矩：28.68 Nm
 机身底部离地：约 0.22 m
 目标越障高度：0.20 ~ 0.24 m
 ```
@@ -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 的轴向总尺寸，简化圆柱模型将其作为轮宽。
 - 当前执行器仍使用固定扭矩与固定速度上限，没有拟合手册中的完整扭矩—转速曲线、热衰减或电流限制。
@@ -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只加载方块和圆柱参与物理计算。
@@ -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"
@@ -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]]
@@ -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
@@ -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()
@@ -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 <checkpoint.pt>"
  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
@@ -0,0 +1,280 @@
 # Copyright (c) 2025 Deep Robotics
 # SPDX-License-Identifier: BSD 3-Clause
 # Copyright (c) 2024-2025 Ziqi Fan
 # SPDX-License-Identifier: Apache-2.0
 # Copyright (c) 2024-2025, The Isaac Lab Project Developers.
 # All rights reserved.
 #
 # SPDX-License-Identifier: Apache-2.0
 """Script to play a checkpoint if an RL agent from RSL-RL."""
 """Launch Isaac Sim Simulator first."""
 import argparse
 import os
 import sys
 PROJECT_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", ".."))
 PROJECT_SOURCE_PATH = os.path.join(PROJECT_ROOT, "source", "rl_training")
 if PROJECT_SOURCE_PATH not in sys.path:
    sys.path.insert(0, PROJECT_SOURCE_PATH)
 from isaaclab.app import AppLauncher
 # local imports
 sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
 import cli_args
 from rl_utils import camera_follow
 # add argparse arguments
 parser = argparse.ArgumentParser(description="Train an RL agent with RSL-RL.")
 parser.add_argument("--video", action="store_true", default=False, help="Record videos during training.")
 parser.add_argument("--video_length", type=int, default=200, help="Length of the recorded video (in steps).")
 parser.add_argument(
    "--disable_fabric", action="store_true", default=False, help="Disable fabric and use USD I/O operations."
 )
 parser.add_argument("--num_envs", type=int, default=None, help="Number of environments to simulate.")
 parser.add_argument("--task", type=str, default=None, help="Name of the task.")
 parser.add_argument(
    "--agent", type=str, default="rsl_rl_cfg_entry_point", help="Name of the RL agent configuration entry point."
 )
 parser.add_argument("--seed", type=int, default=None, help="Seed used for the environment")
 parser.add_argument(
    "--use_pretrained_checkpoint",
    action="store_true",
    help="Use the pre-trained checkpoint from Nucleus.",
 )
 parser.add_argument("--real-time", action="store_true", default=False, help="Run in real-time, if possible.")
 parser.add_argument("--keyboard", action="store_true", default=False, help="Whether to use keyboard.")
 parser.add_argument(
    "--visual-only",
    action="store_true",
    default=False,
    help="Display the environment with zero actions without loading an RL checkpoint.",
 )
 # append RSL-RL cli arguments
 cli_args.add_rsl_rl_args(parser)
 # append AppLauncher cli args
 AppLauncher.add_app_launcher_args(parser)
 # parse the arguments
 args_cli, hydra_args = parser.parse_known_args()
 # always enable cameras to record video
 if args_cli.video:
    args_cli.enable_cameras = True
 # clear out sys.argv for Hydra
 sys.argv = [sys.argv[0]] + hydra_args
 # launch omniverse app
 app_launcher = AppLauncher(args_cli)
 simulation_app = app_launcher.app
 """Rest everything follows."""
 import gymnasium as gym
 import time
 import torch
 from rsl_rl.runners import OnPolicyRunner
 from isaaclab.devices import Se2Keyboard, Se2KeyboardCfg
 from isaaclab.envs import (
    DirectMARLEnv,
    DirectMARLEnvCfg,
    DirectRLEnvCfg,
    ManagerBasedRLEnvCfg,
    multi_agent_to_single_agent,
 )
 from isaaclab.managers import ObservationTermCfg as ObsTerm
 from isaaclab.utils.assets import retrieve_file_path
 from isaaclab.utils.dict import print_dict
 from isaaclab.utils.pretrained_checkpoint import get_published_pretrained_checkpoint
 from isaaclab_rl.rsl_rl import RslRlOnPolicyRunnerCfg, RslRlVecEnvWrapper, export_policy_as_jit, export_policy_as_onnx
 from isaaclab_tasks.utils import get_checkpoint_path
 from isaaclab_tasks.utils.hydra import hydra_task_config
 import rl_training
 import rl_training.tasks  # noqa: F401
 print(f"[INFO] Using rl_training package from: {rl_training.__file__}")
@hydra_task_config(args_cli.task, args_cli.agent)
 def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agent_cfg: RslRlOnPolicyRunnerCfg):
    """Play with RSL-RL agent."""
    task_name = args_cli.task.split(":")[-1]
    # override configurations with non-hydra CLI arguments
    agent_cfg = cli_args.update_rsl_rl_cfg(agent_cfg, args_cli)
    env_cfg.scene.num_envs = args_cli.num_envs if args_cli.num_envs is not None else 50
    # set the environment seed
    # note: certain randomizations occur in the environment initialization so we set the seed here
    env_cfg.seed = agent_cfg.seed
    env_cfg.sim.device = args_cli.device if args_cli.device is not None else env_cfg.sim.device
    # spawn the robot randomly in the grid (instead of their terrain levels)
    env_cfg.scene.terrain.max_init_terrain_level = None
    # reduce the number of terrains to save memory
    if env_cfg.scene.terrain.terrain_generator is not None:
        env_cfg.scene.terrain.terrain_generator.num_rows = 5
        env_cfg.scene.terrain.terrain_generator.num_cols = 5
        env_cfg.scene.terrain.terrain_generator.curriculum = False
    # disable randomization for play
    env_cfg.observations.policy.enable_corruption = False
    # Remove perturbations and parameter randomization for reproducible playback.
    for event_name in (
        "randomize_rigid_body_material",
        "randomize_wheel_material",
        "randomize_support_material",
        "randomize_rigid_body_mass",
        "randomize_rigid_body_mass_base",
        "randomize_rigid_body_inertia",
        "randomize_com_positions",
        "randomize_apply_external_force_torque",
        "randomize_actuator_gains",
        "randomize_push_robot",
    ):
        if hasattr(env_cfg.events, event_name):
            setattr(env_cfg.events, event_name, None)
    env_cfg.curriculum.command_levels = None
    if args_cli.keyboard:
        env_cfg.scene.num_envs = 1
        env_cfg.terminations.time_out = None
        env_cfg.commands.base_velocity.debug_vis = False
        config = Se2KeyboardCfg(
            v_x_sensitivity=env_cfg.commands.base_velocity.ranges.lin_vel_x[1]/2,
            v_y_sensitivity=env_cfg.commands.base_velocity.ranges.lin_vel_y[1],
            omega_z_sensitivity=env_cfg.commands.base_velocity.ranges.ang_vel_z[1],
        )
        controller = Se2Keyboard(config)
        env_cfg.observations.policy.velocity_commands = ObsTerm(
            func=lambda env: torch.tensor(controller.advance(), dtype=torch.float32).unsqueeze(0).to(env.device),
        )
    # specify directory for logging experiments
    log_root_path = os.path.join("logs", "rsl_rl", agent_cfg.experiment_name)
    log_root_path = os.path.abspath(log_root_path)
    if args_cli.visual_only:
        resume_path = None
        log_dir = os.path.join(log_root_path, "visual_only")
        print("[INFO] Visual-only mode: no policy checkpoint will be loaded.")
    else:
        print(f"[INFO] Loading experiment from directory: {log_root_path}")
        if args_cli.use_pretrained_checkpoint:
            resume_path = get_published_pretrained_checkpoint("rsl_rl", task_name)
            if not resume_path:
                print("[INFO] Unfortunately a pre-trained checkpoint is currently unavailable for this task.")
                return
        elif args_cli.load_run is not None:
            resume_path = get_checkpoint_path(log_root_path, agent_cfg.load_run, agent_cfg.load_checkpoint)
        elif args_cli.checkpoint:
            resume_path = retrieve_file_path(args_cli.checkpoint)
        else:
            resume_path = get_checkpoint_path(log_root_path, agent_cfg.load_run, agent_cfg.load_checkpoint)
        log_dir = os.path.dirname(resume_path)
    # create isaac environment
    env = gym.make(args_cli.task, cfg=env_cfg, render_mode="rgb_array" if args_cli.video else None)
    # convert to single-agent instance if required by the RL algorithm
    if isinstance(env.unwrapped, DirectMARLEnv):
        env = multi_agent_to_single_agent(env)
    # wrap for video recording
    if args_cli.video:
        video_kwargs = {
            "video_folder": os.path.join(log_dir, "videos", "play"),
            "step_trigger": lambda step: step == 0,
            "video_length": args_cli.video_length,
            "disable_logger": True,
        }
        print("[INFO] Recording videos during training.")
        print_dict(video_kwargs, nesting=4)
        env = gym.wrappers.RecordVideo(env, **video_kwargs)
    # wrap around environment for rsl-rl
    env = RslRlVecEnvWrapper(env, clip_actions=agent_cfg.clip_actions)
    if args_cli.visual_only:
        policy = None
    else:
        print(f"[INFO]: Loading model checkpoint from: {resume_path}")
        # load previously trained model
        ppo_runner = OnPolicyRunner(env, agent_cfg.to_dict(), log_dir=None, device=agent_cfg.device)
        ppo_runner.load(resume_path)
        # obtain the trained policy for inference
        policy = ppo_runner.get_inference_policy(device=env.unwrapped.device)
        # extract the neural network module
        # we do this in a try-except to maintain backwards compatibility.
        try:
            # version 2.3 onwards
            policy_nn = ppo_runner.alg.policy
        except AttributeError:
            # version 2.2 and below
            policy_nn = ppo_runner.alg.actor_critic
        # export policy to onnx/jit
        export_model_dir = os.path.join(os.path.dirname(resume_path), "exported")
        export_policy_as_onnx(
            policy=policy_nn,
            normalizer=None,
            path=export_model_dir,
            filename="policy.onnx",
        )
        export_policy_as_jit(
            policy=policy_nn,
            normalizer=None,
            path=export_model_dir,
            filename="policy.pt",
        )
    dt = env.unwrapped.step_dt
    # print(dt, "dt")
    # reset environment
    obs = env.get_observations()
    timestep = 0
    # simulate environment
    while simulation_app.is_running():
        start_time = time.time()
        # run everything in inference mode
        with torch.inference_mode():
            # agent stepping
            if policy is None:
                actions = torch.zeros((env.num_envs, env.num_actions), device=env.unwrapped.device)
            else:
                actions = policy(obs)
            # env stepping
            obs, _, _, _ = env.step(actions)
        if args_cli.video:
            timestep += 1
            # Exit the play loop after recording one video
            if timestep == args_cli.video_length:
                break
        if args_cli.keyboard:
            camera_follow(env)
        # time delay for real-time evaluation
        sleep_time = dt - (time.time() - start_time)
        if args_cli.real_time and sleep_time > 0:
            time.sleep(sleep_time)
    # close the simulator
    env.close()
 if __name__ == "__main__":
    # run the main function
    main()
    # close sim app
    simulation_app.close()
@@ -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()
@@ -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()
@@ -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()
@@ -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()
@@ -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"
@@ -0,0 +1,3 @@
 [build-system]
 requires = ["setuptools", "wheel", "toml"]
 build-backend = "setuptools.build_meta"
@@ -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 *
@@ -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"]
@@ -0,0 +1,124 @@
 # Copyright (c) 2025 Deep Robotics
 # SPDX-License-Identifier: BSD 3-Clause
 # Copyright (c) 2024-2025 Ziqi Fan
 # SPDX-License-Identifier: Apache-2.0
 import isaaclab.sim as sim_utils
 from isaaclab.actuators import DCMotorCfg, DelayedPDActuatorCfg
 from isaaclab.assets.articulation import ArticulationCfg
 from rl_training.assets import ISAACLAB_ASSETS_DATA_DIR
 DEEPROBOTICS_LITE3_CFG = ArticulationCfg(
    spawn=sim_utils.UsdFileCfg(
        usd_path=f"{ISAACLAB_ASSETS_DATA_DIR}/Lite3/Lite3_usd/Lite3.usd",
        activate_contact_sensors=True,
        rigid_props=sim_utils.RigidBodyPropertiesCfg(
            disable_gravity=False,
            retain_accelerations=False,
            linear_damping=0.0,
            angular_damping=0.0,
            max_linear_velocity=1000.0,
            max_angular_velocity=1000.0,
            max_depenetration_velocity=1.0,
        ),
        articulation_props=sim_utils.ArticulationRootPropertiesCfg(
            enabled_self_collisions=False, solver_position_iteration_count=4, solver_velocity_iteration_count=1
        ),
    ),
    init_state=ArticulationCfg.InitialStateCfg(
        pos=(0.0, 0.0, 0.35),
        joint_pos={
            ".*Knee_joint": 0.0,
        },
        joint_vel={".*": 0.0},
    ),
    soft_joint_pos_limit_factor=0.99,
    actuators={
        # "Hip": DelayedPDActuatorCfg(
        #     joint_names_expr=[".*_Hip[X,Y]_joint"],
        #     effort_limit=24.0,
        #     velocity_limit=26.2,
        #     stiffness=30.0,
        #     damping=1.0,
        #     friction=0.0,
        #     armature=0.0,
        #     min_delay=0,
        #     max_delay=5,
        # ),
        "Knee": DelayedPDActuatorCfg(
            joint_names_expr=[".*_Knee_joint"],
            effort_limit=36.0,
            velocity_limit=26.2,
            stiffness=30.0,
            damping=1.0,
            friction=0.0,
            armature=0.0,
            min_delay=0,
            max_delay=5,
        ),
    },
 )
 DEEPROBOTICS_M20_CFG = ArticulationCfg(
    spawn=sim_utils.UrdfFileCfg(
        asset_path=f"{ISAACLAB_ASSETS_DATA_DIR}/M20/M20_urdf/urdf/Solenopsis.urdf",
        fix_base=False,
        merge_fixed_joints=True,
        joint_drive=None,
        activate_contact_sensors=True,
        visual_material=sim_utils.PreviewSurfaceCfg(
            diffuse_color=(0.42, 0.45, 0.50),
            roughness=0.65,
            metallic=0.15,
        ),
        rigid_props=sim_utils.RigidBodyPropertiesCfg(
            disable_gravity=False,
            retain_accelerations=False,
            linear_damping=0.0,
            angular_damping=0.0,
            max_linear_velocity=1000.0,
            max_angular_velocity=1000.0,
            max_depenetration_velocity=1.0,
        ),
        articulation_props=sim_utils.ArticulationRootPropertiesCfg(
            enabled_self_collisions=False, solver_position_iteration_count=4, solver_velocity_iteration_count=1
        ),
    ),
    init_state=ArticulationCfg.InitialStateCfg(
        pos=(0.0, 0.0, 0.306),
        joint_pos={
            "f[l,r]_knee_joint": 0.5,
            "h[l,r]_knee_joint": -0.5,
            ".*wheel_joint": 0.0,
        },
        joint_vel={".*": 0.0},
    ),
    soft_joint_pos_limit_factor=0.9,
    actuators={
        "joint": DelayedPDActuatorCfg(
            joint_names_expr=[".*knee_joint"],
            effort_limit=75.0,
            velocity_limit=17.8024,
            stiffness=80.0,
            damping=2.0,
            friction=0.0,
            armature=0.0,
            min_delay=0,
            max_delay=5,
        ),
        "wheel": DelayedPDActuatorCfg(
            joint_names_expr=[".*_wheel_joint"],
            effort_limit=28.68,
            velocity_limit=109.956,
            stiffness=0.0,
            damping=0.6,
            friction=1.0,
            armature=0.00243216,
            min_delay=0,
            max_delay=5,
        ),
    },
 )
@@ -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)
@@ -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
@@ -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
@@ -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
 """
@@ -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.
@@ -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.
@@ -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",
    },
 )
@@ -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
@@ -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"
@@ -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()
@@ -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)
@@ -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.
@@ -0,0 +1,46 @@
 # Copyright (c) 2025 Deep Robotics
 # SPDX-License-Identifier: BSD 3-Clause
 # 
 # # Copyright (c) 2024-2025 Ziqi Fan
 # SPDX-License-Identifier: Apache-2.0
 import gymnasium as gym
 from . import agents
 ##
 # Register Gym environments.
 ##
 gym.register(
    id="Flat-Deeprobotics-M20-v0",
    entry_point="isaaclab.envs:ManagerBasedRLEnv",
    disable_env_checker=True,
    kwargs={
        "env_cfg_entry_point": f"{__name__}.flat_env_cfg:DeeproboticsM20FlatEnvCfg",
        "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:DeeproboticsM20FlatPPORunnerCfg",
        "cusrl_cfg_entry_point": f"{agents.__name__}.cusrl_ppo_cfg:DeeproboticsM20FlatTrainerCfg",
    },
 )
 gym.register(
    id="Rough-Deeprobotics-M20-v0",
    entry_point="isaaclab.envs:ManagerBasedRLEnv",
    disable_env_checker=True,
    kwargs={
        "env_cfg_entry_point": f"{__name__}.rough_env_cfg:DeeproboticsM20RoughEnvCfg",
        "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:DeeproboticsM20RoughPPORunnerCfg",
        "cusrl_cfg_entry_point": f"{agents.__name__}.cusrl_ppo_cfg:DeeproboticsM20RoughTrainerCfg",
    },
 )
 gym.register(
    id="Rough-Deeprobotics-M20-Support-v0",
    entry_point="isaaclab.envs:ManagerBasedRLEnv",
    disable_env_checker=True,
    kwargs={
        "env_cfg_entry_point": f"{__name__}.rough_env_cfg:DeeproboticsM20SupportRoughEnvCfg",
        "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:DeeproboticsM20SupportRoughPPORunnerCfg",
        "cusrl_cfg_entry_point": f"{agents.__name__}.cusrl_ppo_cfg:DeeproboticsM20RoughTrainerCfg",
    },
 )
@@ -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
@@ -0,0 +1,56 @@
 # Copyright (c) 2025 Deep Robotics
 # SPDX-License-Identifier: BSD 3-Clause
 # Copyright (c) 2024-2025 Ziqi Fan
 # SPDX-License-Identifier: Apache-2.0
 from isaaclab.utils import configclass
 from isaaclab_rl.rsl_rl import RslRlOnPolicyRunnerCfg, RslRlPpoActorCriticCfg, RslRlPpoAlgorithmCfg
@configclass
 class DeeproboticsM20RoughPPORunnerCfg(RslRlOnPolicyRunnerCfg):
    num_steps_per_env = 24
    max_iterations = 20000
    save_interval = 100
    experiment_name = "deeprobotics_m20_rough"
    empirical_normalization = False
    clip_actions = 100
    policy = RslRlPpoActorCriticCfg(
        init_noise_std=1.0,
        noise_std_type="log",
        actor_hidden_dims=[512, 256, 128],
        critic_hidden_dims=[512, 256, 128],
        activation="elu",
    )
    algorithm = RslRlPpoAlgorithmCfg(
        value_loss_coef=1.0,
        use_clipped_value_loss=True,
        clip_param=0.2,
        entropy_coef=0.01,
        num_learning_epochs=5,
        num_mini_batches=4,
        learning_rate=1.0e-3,
        schedule="adaptive",
        gamma=0.99,
        lam=0.95,
        desired_kl=0.01,
        max_grad_norm=1.0,
    )
@configclass
 class DeeproboticsM20FlatPPORunnerCfg(DeeproboticsM20RoughPPORunnerCfg):
    def __post_init__(self):
        super().__post_init__()
        self.max_iterations = 5000
        self.experiment_name = "deeprobotics_m20_flat"
@configclass
 class DeeproboticsM20SupportRoughPPORunnerCfg(DeeproboticsM20RoughPPORunnerCfg):
    def __post_init__(self):
        super().__post_init__()
        self.experiment_name = "deeprobotics_m20_support"
@@ -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()
@@ -0,0 +1,504 @@
 # Copyright (c) 2025 Deep Robotics
 # SPDX-License-Identifier: BSD 3-Clause
 # Copyright (c) 2024-2025 Ziqi Fan
 # SPDX-License-Identifier: Apache-2.0
 from isaaclab.managers import CurriculumTermCfg as CurrTerm
 from isaaclab.managers import EventTermCfg as EventTerm
 from isaaclab.managers import RewardTermCfg as RewTerm
 from isaaclab.managers import ObservationTermCfg as ObsTerm
 from isaaclab.managers import SceneEntityCfg
 from isaaclab.utils import configclass
 import rl_training.tasks.manager_based.locomotion.velocity.mdp as mdp
 from rl_training.tasks.manager_based.locomotion.velocity.velocity_env_cfg import (
    ActionsCfg,
    EventCfg,
    LocomotionVelocityRoughEnvCfg,
    RewardsCfg,
 )
 ##
 # Pre-defined configs
 ##
 from rl_training.assets.deeprobotics import DEEPROBOTICS_M20_CFG  # isort: skip
@configclass
 class DeeproboticsM20ActionsCfg(ActionsCfg):
    """Action specifications for the MDP."""
    joint_pos = mdp.JointPositionActionCfg(
        asset_name="robot", joint_names=[""], scale=0.25, use_default_offset=True, clip=None, preserve_order=True
    )
    joint_vel = mdp.JointVelocityActionCfg(
        asset_name="robot", joint_names=[""], scale=20.0, use_default_offset=True, clip=None, preserve_order=True
    )
@configclass
 class DeeproboticsM20EventCfg(EventCfg):
    """Domain randomization with tire-specific contact material."""
    # Replace the inherited all-body material randomization with separate tire and support materials.
    randomize_rigid_body_material = None
    randomize_wheel_material = EventTerm(
        func=mdp.randomize_rigid_body_material,
        mode="startup",
        params={
            "asset_cfg": SceneEntityCfg("robot", body_names=".*_wheel"),
            "static_friction_range": (0.9, 1.4),
            "dynamic_friction_range": (0.6, 0.9),
            "restitution_range": (0.0, 0.05),
            "num_buckets": 1024,
        },
    )
    randomize_support_material = EventTerm(
        func=mdp.randomize_rigid_body_material,
        mode="startup",
        params={
            "asset_cfg": SceneEntityCfg("robot", body_names=["base_link", ".*_knee"]),
            "static_friction_range": (0.4, 0.8),
            "dynamic_friction_range": (0.25, 0.4),
            "restitution_range": (0.0, 0.05),
            "num_buckets": 1024,
        },
    )
@configclass
 class DeeproboticsM20RewardsCfg(RewardsCfg):
    """Reward terms for the MDP."""
    joint_vel_wheel_l2 = RewTerm(
        func=mdp.joint_vel_l2, weight=0.0, params={"asset_cfg": SceneEntityCfg("robot", joint_names="")}
    )
    joint_acc_wheel_l2 = RewTerm(
        func=mdp.joint_acc_l2, weight=0.0, params={"asset_cfg": SceneEntityCfg("robot", joint_names="")}
    )
    joint_torques_wheel_l2 = RewTerm(
        func=mdp.joint_torques_l2, weight=0.0, params={"asset_cfg": SceneEntityCfg("robot", joint_names="")}
    )
    track_diff_wheel_yaw_exp = RewTerm(
        func=mdp.track_diff_wheel_yaw_exp,
        weight=1.0,
        params={
            "command_name": "base_velocity",
            "std": 0.6,
            "wheel_radius": 0.096,
            "track_width": 0.50,
            "left_wheel_indices": (0, 2),
            "right_wheel_indices": (1, 3),
            "command_threshold": 0.05,
            "asset_cfg": SceneEntityCfg("robot", joint_names=""),
        },
    )
    lin_vel_xy_yaw_command_l2 = RewTerm(
        func=mdp.lin_vel_xy_yaw_command_l2,
        weight=-0.4,
        params={
            "command_name": "base_velocity",
            "yaw_command_threshold": 0.25,
            "lin_command_threshold": 0.15,
            "asset_cfg": SceneEntityCfg("robot"),
        },
    )
    command_progress = RewTerm(
        func=mdp.command_progress,
        weight=0.0,
        params={
            "command_name": "base_velocity",
            "min_command": 0.12,
            "max_reward_speed": 0.8,
            "asset_cfg": SceneEntityCfg("robot"),
        },
    )
    support_contact_progress = RewTerm(
        func=mdp.ObstacleSupportEscapeReward,
        weight=0.0,
        params={
            "command_name": "base_velocity",
            "sensor_cfg": SceneEntityCfg("contact_forces", body_names=""),
            "wheel_sensor_cfg": SceneEntityCfg("contact_forces", body_names=""),
            "asset_cfg": SceneEntityCfg("robot"),
            "wheel_asset_cfg": SceneEntityCfg("robot", joint_names=""),
            "wheel_radius": 0.096,
            "support_force_threshold": 15.0,
            "wheel_obstacle_force_threshold": 35.0,
            "max_contact_time": 0.45,
            "min_command": 0.15,
            "stall_speed_ratio": 0.35,
            "min_stall_speed": 0.12,
            "slip_threshold": 0.30,
            "obstacle_memory_time": 0.60,
            "support_window_time": 0.80,
            "recovery_speed_ratio": 0.60,
            "min_recovery_speed": 0.18,
            "recovery_reward_time": 0.25,
            "cooldown_time": 0.60,
            "recovery_bonus": 0.25,
            "max_reward_speed": 0.8,
        },
    )
    support_contact_excess = RewTerm(
        func=mdp.support_contact_excess,
        weight=0.0,
        params={
            "sensor_cfg": SceneEntityCfg("contact_forces", body_names=""),
            "force_threshold": 120.0,
            "max_contact_time": 0.6,
        },
    )
@configclass
 class DeeproboticsM20RoughEnvCfg(LocomotionVelocityRoughEnvCfg):
    """Stage 1: blind tactile locomotion without body-support rewards."""
    actions: DeeproboticsM20ActionsCfg = DeeproboticsM20ActionsCfg()
    events: DeeproboticsM20EventCfg = DeeproboticsM20EventCfg()
    rewards: DeeproboticsM20RewardsCfg = DeeproboticsM20RewardsCfg()
    base_link_name = "base_link"
    foot_link_name = ".*_wheel"
    support_link_names = ["base_link", ".*_knee"]
    # fmt: off
    leg_joint_names = [
        "fl_knee_joint",
        "fr_knee_joint",
        "hl_knee_joint",
        "hr_knee_joint",
    ]
    wheel_joint_names = [
        "fl_wheel_joint", "fr_wheel_joint", "hl_wheel_joint", "hr_wheel_joint",
    ]
    knee_joint_names = [
        "fl_knee_joint", "fr_knee_joint", "hl_knee_joint", "hr_knee_joint",
    ]
    joint_names = leg_joint_names + wheel_joint_names
    # fmt: on
    def __post_init__(self):
        # post init of parent
        super().__post_init__()
        # ------------------------------Sence------------------------------
        self.scene.robot = DEEPROBOTICS_M20_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot")
        self.scene.height_scanner.prim_path = "{ENV_REGEX_NS}/Robot/" + self.base_link_name
        self.scene.height_scanner_base.prim_path = "{ENV_REGEX_NS}/Robot/" + self.base_link_name
        self.scene.terrain.physics_material.restitution = 0.0
        self.sim.physics_material = self.scene.terrain.physics_material
        # ------------------------------Observations------------------------------
        self.observations.policy.joint_pos.func = mdp.joint_pos_rel_without_wheel
        self.observations.policy.joint_pos.params["wheel_asset_cfg"] = SceneEntityCfg(
            "robot", joint_names=self.wheel_joint_names
        )
        self.observations.critic.joint_pos.func = mdp.joint_pos_rel_without_wheel
        self.observations.critic.joint_pos.params["wheel_asset_cfg"] = SceneEntityCfg(
            "robot", joint_names=self.wheel_joint_names
        )
        self.observations.policy.base_lin_vel.scale = 2.0
        self.observations.policy.base_ang_vel.scale = 0.25
        self.observations.policy.joint_pos.scale = 1.0
        self.observations.policy.joint_vel.scale = 0.05
        self.observations.policy.base_lin_vel = None
        self.observations.policy.height_scan = None
        self.observations.policy.joint_pos.params["asset_cfg"].joint_names = self.joint_names
        self.observations.policy.joint_vel.params["asset_cfg"].joint_names = self.joint_names
        self.observations.policy.wheel_contact = ObsTerm(
            func=mdp.contact_state,
            params={
                "sensor_cfg": SceneEntityCfg("contact_forces", body_names=[self.foot_link_name]),
                "threshold": 2.0,
            },
            clip=(0.0, 1.0),
            scale=1.0,
        )
        self.observations.policy.support_contact = ObsTerm(
            func=mdp.contact_state,
            params={
                "sensor_cfg": SceneEntityCfg("contact_forces", body_names=self.support_link_names),
                "threshold": 2.0,
            },
            clip=(0.0, 1.0),
            scale=1.0,
        )
        self.observations.policy.wheel_contact_force = ObsTerm(
            func=mdp.contact_force_norm,
            params={
                "sensor_cfg": SceneEntityCfg("contact_forces", body_names=[self.foot_link_name]),
                "scale": 100.0,
                "clip": 2.0,
            },
            clip=(0.0, 2.0),
            scale=1.0,
        )
        self.observations.policy.support_contact_force = ObsTerm(
            func=mdp.contact_force_norm,
            params={
                "sensor_cfg": SceneEntityCfg("contact_forces", body_names=self.support_link_names),
                "scale": 100.0,
                "clip": 2.0,
            },
            clip=(0.0, 2.0),
            scale=1.0,
        )
        self.observations.policy.wheel_slip = ObsTerm(
            func=mdp.wheel_slip_proxy,
            params={
                "asset_cfg": SceneEntityCfg("robot", joint_names=self.wheel_joint_names, preserve_order=True),
                "wheel_radius": 0.096,
            },
            clip=(0.0, 3.0),
            scale=0.5,
        )
        self.observations.policy.applied_torque = ObsTerm(
            func=mdp.applied_torque,
            params={
                "asset_cfg": SceneEntityCfg("robot", joint_names=self.joint_names, preserve_order=True),
                "scale": 80.0,
                "clip": 2.0,
            },
            clip=(-2.0, 2.0),
            scale=1.0,
        )
        # ------------------------------Actions------------------------------
        # reduce action scale
        self.actions.joint_pos.scale = 0.25
        self.actions.joint_vel.scale = 5.0
        self.actions.joint_pos.clip = {".*": (-100.0, 100.0)}
        self.actions.joint_vel.clip = {".*": (-100.0, 100.0)}
        self.actions.joint_pos.joint_names = self.leg_joint_names
        self.actions.joint_vel.joint_names = self.wheel_joint_names
        # ------------------------------Events------------------------------
        self.events.randomize_reset_base.params = {
            "pose_range": {
                "x": (-1.0, 1.0),
                "y": (-1.0, 1.0),
                "z": (0.0, 0.0),
                "roll": (-0.3, 0.3),
                "pitch": (-0.3, 0.3),
                "yaw": (-3.14, 3.14),
            },
            "velocity_range": {
                "x": (-0.2, 0.2),
                "y": (-0.2, 0.2),
                "z": (-0.2, 0.2),
                "roll": (-0.05, 0.05),
                "pitch": (-0.05, 0.05),
                "yaw": (-0.0, 0.0),
            },
        }
        self.events.randomize_rigid_body_mass_base.params["asset_cfg"].body_names = [self.base_link_name]
        self.events.randomize_rigid_body_mass.params["asset_cfg"].body_names = [
            f"^(?!.*{self.base_link_name}).*"
        ]
        self.events.randomize_com_positions.params["asset_cfg"].body_names = [self.base_link_name]
        self.events.randomize_apply_external_force_torque.params["asset_cfg"].body_names = [self.base_link_name]
        self.scene.terrain.max_init_terrain_level = 2
        self.scene.terrain.terrain_generator.sub_terrains["boxes"].grid_height_range = (0.02, 0.14)
        self.scene.terrain.terrain_generator.sub_terrains["pyramid_stairs"].step_height_range = (0.03, 0.16)
        self.scene.terrain.terrain_generator.sub_terrains["pyramid_stairs_inv"].step_height_range = (0.03, 0.16)
        self.scene.terrain.terrain_generator.sub_terrains["random_rough"].noise_range = (0.01, 0.10)
        self.scene.terrain.terrain_generator.sub_terrains["random_rough"].noise_step = 0.01
        # only randomize actuator gains for actuated joints (knees + wheels), exclude hipx/hipy completely
        self.events.randomize_actuator_gains.params["asset_cfg"].joint_names = self.joint_names
        # ------------------------------Rewards------------------------------
        # General
        self.rewards.is_terminated.weight = 0
        # Root penalties
        self.rewards.lin_vel_z_l2.weight = -2.0
        self.rewards.ang_vel_xy_l2.weight = -0.02
        self.rewards.flat_orientation_l2.weight = -0.5
        self.rewards.base_height_l2.weight = -0.5
        self.rewards.base_height_l2.params["target_height"] = 0.45
        self.rewards.base_height_l2.params["asset_cfg"].body_names = [self.base_link_name]
        self.rewards.body_lin_acc_l2.weight = 0
        self.rewards.body_lin_acc_l2.params["asset_cfg"].body_names = [self.base_link_name]
        # Joint penalties
        self.rewards.joint_torques_l2.weight = -2.5e-5
        self.rewards.joint_torques_l2.params["asset_cfg"].joint_names = self.leg_joint_names
        self.rewards.joint_torques_wheel_l2.weight = 0
        self.rewards.joint_torques_wheel_l2.params["asset_cfg"].joint_names = self.wheel_joint_names
        self.rewards.joint_vel_l2.weight = 0
        self.rewards.joint_vel_l2.params["asset_cfg"].joint_names = self.leg_joint_names
        self.rewards.joint_vel_wheel_l2.weight = 0
        self.rewards.joint_vel_wheel_l2.params["asset_cfg"].joint_names = self.wheel_joint_names
        self.rewards.joint_acc_l2.weight = -2e-7
        self.rewards.joint_acc_l2.params["asset_cfg"].joint_names = self.leg_joint_names
        self.rewards.joint_acc_wheel_l2.weight = -1e-7
        self.rewards.joint_acc_wheel_l2.params["asset_cfg"].joint_names = self.wheel_joint_names
        # self.rewards.create_joint_deviation_l1_rewterm("joint_deviation_hip_l1", -0.2, [".*_hip_joint"])
        self.rewards.joint_pos_limits.weight = -5.0
        self.rewards.joint_pos_limits.params["asset_cfg"].joint_names = self.leg_joint_names
        self.rewards.joint_vel_limits.weight = 0
        self.rewards.joint_vel_limits.params["asset_cfg"].joint_names = self.wheel_joint_names
        self.rewards.joint_power.weight = -2e-5
        self.rewards.joint_power.params["asset_cfg"].joint_names = self.leg_joint_names
        self.rewards.stand_still.weight = -2.0
        self.rewards.stand_still.params["asset_cfg"].joint_names = self.leg_joint_names
        self.rewards.knee_joint_pos_penalty.weight = -0.1
        self.rewards.knee_joint_pos_penalty.params["asset_cfg"].joint_names = self.knee_joint_names
        self.rewards.wheel_vel_penalty.weight = 0
        self.rewards.wheel_vel_penalty.params["sensor_cfg"].body_names = self.foot_link_name
        self.rewards.wheel_vel_penalty.params["asset_cfg"].joint_names = self.wheel_joint_names
        self.rewards.joint_mirror.weight = -0.03
        self.rewards.joint_mirror.params["mirror_joints"] = [
            ["fl_knee_joint", "hr_knee_joint"],
            ["fr_knee_joint", "hl_knee_joint"],
        ]
        # Action penalties
        self.rewards.action_rate_l2.weight = -0.01
        # Contact sensor
        self.rewards.undesired_contacts.weight = -1.0
        self.rewards.undesired_contacts.params["sensor_cfg"].body_names = [f"^(?!.*{self.foot_link_name}).*"]
        self.rewards.undesired_contacts.params["threshold"] = 1.0
        self.rewards.contact_forces.weight = -1.5e-4
        self.rewards.contact_forces.params["sensor_cfg"].body_names = [self.foot_link_name]
        self.rewards.support_contact_progress.weight = 0.0
        self.rewards.support_contact_progress.params["sensor_cfg"].body_names = self.support_link_names
        self.rewards.support_contact_progress.params["wheel_sensor_cfg"].body_names = [self.foot_link_name]
        self.rewards.support_contact_progress.params["wheel_asset_cfg"].joint_names = self.wheel_joint_names
        self.rewards.support_contact_excess.weight = 0.0
        self.rewards.support_contact_excess.params["sensor_cfg"].body_names = self.support_link_names
        # Velocity-tracking rewards
        self.rewards.track_lin_vel_xy_exp.weight = 2.0 # 1.8
        self.rewards.track_ang_vel_z_exp.weight = 1.5 # 1.2
        self.rewards.track_ang_vel_z_exp.params["std"] = 0.5
        self.rewards.track_diff_wheel_yaw_exp.weight = 1.0
        self.rewards.track_diff_wheel_yaw_exp.params["asset_cfg"].joint_names = self.wheel_joint_names
        self.rewards.lin_vel_xy_yaw_command_l2.weight = -0.4
        # Others
        self.rewards.feet_air_time.weight = 0
        self.rewards.feet_air_time.params["threshold"] = 0.5
        self.rewards.feet_air_time.params["sensor_cfg"].body_names = [self.foot_link_name]
        self.rewards.feet_contact.weight = 0
        self.rewards.feet_contact.params["sensor_cfg"].body_names = [self.foot_link_name]
        self.rewards.feet_contact_without_cmd.weight = 0.1
        self.rewards.feet_contact_without_cmd.params["sensor_cfg"].body_names = [self.foot_link_name]
        self.rewards.feet_stumble.weight = 0
        self.rewards.feet_stumble.params["sensor_cfg"].body_names = [self.foot_link_name]
        self.rewards.feet_slide.weight = 0
        self.rewards.feet_slide.params["sensor_cfg"].body_names = [self.foot_link_name]
        self.rewards.feet_slide.params["asset_cfg"].body_names = [self.foot_link_name]
        self.rewards.feet_height.weight = 0
        self.rewards.feet_height.params["target_height"] = 0.1
        self.rewards.feet_height.params["asset_cfg"].body_names = [self.foot_link_name]
        self.rewards.feet_height_body.weight = 0
        self.rewards.feet_height_body.params["target_height"] = -0.4
        self.rewards.feet_height_body.params["asset_cfg"].body_names = [self.foot_link_name]
        self.rewards.feet_gait.weight = 0
        self.rewards.feet_gait.params["synced_feet_pair_names"] = (("fl_wheel", "hr_wheel"), ("fr_wheel", "hl_wheel"))
        self.rewards.upward.weight = 0.08
        # If the weight of rewards is 0, set rewards to None
        if self.__class__.__name__ == "DeeproboticsM20RoughEnvCfg":
            self.disable_zero_weight_rewards()
        # ------------------------------Terminations------------------------------
        # self.terminations.illegal_contact.params["sensor_cfg"].body_names = [self.base_link_name]
        self.terminations.illegal_contact = None
        self.terminations.bad_orientation_2 = None
        # ------------------------------Curriculums------------------------------
        # self.curriculum.command_levels.params["range_multiplier"] = (0.2, 1.0)
        self.curriculum.command_levels = None
        # ------------------------------Commands------------------------------
        self.commands.base_velocity.debug_vis = False
        self.commands.base_velocity.ranges.lin_vel_x = (-2.0, 2.0)
        self.commands.base_velocity.ranges.lin_vel_y = (0.0, 0.0)
        self.commands.base_velocity.ranges.ang_vel_z = (-1.5, 1.5)
@configclass
 class DeeproboticsM20SupportRoughEnvCfg(DeeproboticsM20RoughEnvCfg):
    """Stage 2: resume from Stage 1 and enable brief body/knee support for obstacle crossing."""
    def __post_init__(self):
        super().__post_init__()
        # Keep the same policy/action dimensions as Stage 1 for checkpoint resume.
        # Only rewards, command distribution, and terrain curriculum are changed here.
        self.scene.terrain.max_init_terrain_level = 1
        self.scene.terrain.terrain_generator.sub_terrains["boxes"].proportion = 0.25
        self.scene.terrain.terrain_generator.sub_terrains["pyramid_stairs"].proportion = 0.25
        self.scene.terrain.terrain_generator.sub_terrains["pyramid_stairs_inv"].proportion = 0.15
        self.scene.terrain.terrain_generator.sub_terrains["random_rough"].proportion = 0.15
        self.scene.terrain.terrain_generator.sub_terrains["hf_pyramid_slope"].proportion = 0.10
        self.scene.terrain.terrain_generator.sub_terrains["hf_pyramid_slope_inv"].proportion = 0.10
        self.scene.terrain.terrain_generator.sub_terrains["boxes"].grid_height_range = (0.03, 0.22)
        self.scene.terrain.terrain_generator.sub_terrains["pyramid_stairs"].step_height_range = (0.04, 0.24)
        self.scene.terrain.terrain_generator.sub_terrains["pyramid_stairs_inv"].step_height_range = (0.04, 0.20)
        self.scene.terrain.terrain_generator.sub_terrains["random_rough"].noise_range = (0.01, 0.14)
        # Keep privileged terrain information for the critic only. The actor remains blind/tactile.
        self.observations.policy.height_scan = None
        self.observations.critic.height_scan.params["sensor_cfg"] = SceneEntityCfg("height_scanner")
        self.observations.critic.height_scan.clip = (-1.0, 1.0)
        self.observations.critic.height_scan.scale = 1.0
        # Body/knee contact is allowed only as short support during forward progress.
        self.rewards.undesired_contacts.weight = -0.10
        self.rewards.undesired_contacts.params["sensor_cfg"].body_names = [f"^(?!.*{self.foot_link_name}).*"]
        self.rewards.undesired_contacts.params["threshold"] = 35.0
        self.rewards.command_progress.weight = 0.8
        self.rewards.support_contact_progress.weight = 1.0
        self.rewards.support_contact_progress.params["sensor_cfg"].body_names = self.support_link_names
        self.rewards.support_contact_progress.params["wheel_sensor_cfg"].body_names = [self.foot_link_name]
        self.rewards.support_contact_progress.params["wheel_asset_cfg"].joint_names = self.wheel_joint_names
        self.rewards.support_contact_excess.weight = -0.60
        self.rewards.support_contact_excess.params["sensor_cfg"].body_names = self.support_link_names
        self.rewards.support_contact_excess.params["force_threshold"] = 100.0
        self.rewards.support_contact_excess.params["max_contact_time"] = 0.45
        # Stage 2 should still avoid the bad solution of crawling on the body.
        self.rewards.track_lin_vel_xy_exp.weight = 1.5
        self.rewards.track_ang_vel_z_exp.weight = 1.0
        self.rewards.flat_orientation_l2.weight = -0.45
        self.rewards.base_height_l2.weight = -0.35
        self.rewards.base_height_l2.params["target_height"] = 0.29
        # Stage 2 first focuses on forward obstacle crossing, then expands the speed range.
        self.commands.base_velocity.heading_command = False
        self.commands.base_velocity.rel_standing_envs = 0.0
        self.commands.base_velocity.rel_heading_envs = 0.0
        self.commands.base_velocity.ranges.lin_vel_x = (0.25, 1.2)
        self.commands.base_velocity.ranges.lin_vel_y = (0.0, 0.0)
        self.commands.base_velocity.ranges.ang_vel_z = (-0.6, 0.6)
        self.curriculum.command_levels = CurrTerm(
            func=mdp.command_levels_positive_x,
            params={
                "reward_term_name": "track_lin_vel_xy_exp",
                "initial_max_ratio": 0.35,
                "increment": 0.1,
            },
        )
        self.disable_zero_weight_rewards()
@@ -0,0 +1,21 @@
 # Copyright (c) 2025 Deep Robotics
 # SPDX-License-Identifier: BSD 3-Clause
 # 
 # # Copyright (c) 2024-2025 Ziqi Fan
 # SPDX-License-Identifier: Apache-2.0
 # Copyright (c) 2022-2025, The Isaac Lab Project Developers.
 # All rights reserved.
 #
 # SPDX-License-Identifier: BSD-3-Clause
 """This sub-module contains the functions that are specific to the locomotion environments."""
 from isaaclab.envs.mdp import *  # noqa: F401, F403
 from isaaclab_tasks.manager_based.locomotion.velocity.mdp import *  # noqa: F401, F403
 from .commands import *  # noqa: F401, F403
 from .curriculums import *  # noqa: F401, F403
 from .events import *  # noqa: F401, F403
 from .observations import *  # noqa: F401, F403
 from .rewards import *  # noqa: F401, F403
@@ -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]
    """
@@ -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)
@@ -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)
@@ -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)
@@ -0,0 +1,985 @@
 # Copyright (c) 2025 Deep Robotics
 # SPDX-License-Identifier: BSD 3-Clause
 # Copyright (c) 2024-2025 Ziqi Fan
 # SPDX-License-Identifier: Apache-2.0
 from __future__ import annotations
 import torch
 from typing import TYPE_CHECKING
 import isaaclab.utils.math as math_utils
 from isaaclab.assets import Articulation, RigidObject
 from isaaclab.managers import ManagerTermBase
 from isaaclab.managers import RewardTermCfg as RewTerm
 from isaaclab.managers import SceneEntityCfg
 from isaaclab.sensors import ContactSensor, RayCaster
 from isaaclab.utils.math import quat_apply_inverse, yaw_quat
 if TYPE_CHECKING:
    from isaaclab.envs import ManagerBasedRLEnv
 def track_lin_vel_xy_exp(
    env: ManagerBasedRLEnv, std: float, command_name: str, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")
 ) -> torch.Tensor:
    """Reward tracking of linear velocity commands (xy axes) using exponential kernel."""
    # extract the used quantities (to enable type-hinting)
    asset: RigidObject = env.scene[asset_cfg.name]
    # compute the error
    lin_vel_error = torch.sum(
        torch.square(env.command_manager.get_command(command_name)[:, :2] - asset.data.root_lin_vel_b[:, :2]),
        dim=1,
    )
    reward = torch.exp(-lin_vel_error / std**2)
    # reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
    return reward
 def track_ang_vel_z_exp(
    env: ManagerBasedRLEnv, std: float, command_name: str, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")
 ) -> torch.Tensor:
    """Reward tracking of angular velocity commands (yaw) using exponential kernel."""
    # extract the used quantities (to enable type-hinting)
    asset: RigidObject = env.scene[asset_cfg.name]
    # compute the error
    ang_vel_error = torch.square(env.command_manager.get_command(command_name)[:, 2] - asset.data.root_ang_vel_b[:, 2])
    reward = torch.exp(-ang_vel_error / std**2)
    # reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
    return reward
 def track_diff_wheel_yaw_exp(
    env: ManagerBasedRLEnv,
    std: float,
    command_name: str,
    wheel_radius: float,
    track_width: float,
    left_wheel_indices: tuple[int, ...],
    right_wheel_indices: tuple[int, ...],
    command_threshold: float,
    asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"),
 ) -> torch.Tensor:
    """Reward differential wheel speed matching the commanded yaw rate."""
    asset: Articulation = env.scene[asset_cfg.name]
    joint_vel = asset.data.joint_vel[:, asset_cfg.joint_ids]
    left_wheel_vel = torch.mean(joint_vel[:, left_wheel_indices], dim=1)
    right_wheel_vel = torch.mean(joint_vel[:, right_wheel_indices], dim=1)
    estimated_yaw_rate = wheel_radius * (right_wheel_vel - left_wheel_vel) / track_width
    yaw_command = env.command_manager.get_command(command_name)[:, 2]
    yaw_error = torch.square(yaw_command - estimated_yaw_rate)
    reward = torch.exp(-yaw_error / std**2)
    return torch.where(torch.abs(yaw_command) > command_threshold, reward, torch.zeros_like(reward))
 def lin_vel_xy_yaw_command_l2(
    env: ManagerBasedRLEnv,
    command_name: str,
    yaw_command_threshold: float,
    lin_command_threshold: float,
    asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"),
 ) -> torch.Tensor:
    """Penalize linear drift when the command asks mostly for in-place yaw."""
    asset: Articulation = env.scene[asset_cfg.name]
    command = env.command_manager.get_command(command_name)
    is_turn_in_place = (
        torch.abs(command[:, 2]) > yaw_command_threshold
    ) & (torch.linalg.norm(command[:, :2], dim=1) < lin_command_threshold)
    lin_vel_error = torch.sum(torch.square(asset.data.root_lin_vel_b[:, :2]), dim=1)
    return torch.where(is_turn_in_place, lin_vel_error, torch.zeros_like(lin_vel_error))
 def command_progress(
    env: ManagerBasedRLEnv,
    command_name: str,
    min_command: float,
    max_reward_speed: float,
    asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"),
 ) -> torch.Tensor:
    """Reward positive progress along the commanded planar direction."""
    asset: Articulation = env.scene[asset_cfg.name]
    command_xy = env.command_manager.get_command(command_name)[:, :2]
    command_norm = torch.linalg.norm(command_xy, dim=1)
    command_dir = command_xy / torch.clamp(command_norm.unsqueeze(1), min=1e-6)
    progress_speed = torch.sum(asset.data.root_lin_vel_b[:, :2] * command_dir, dim=1)
    reward = torch.clamp(progress_speed, min=0.0, max=max_reward_speed) / max_reward_speed
    return torch.where(command_norm > min_command, reward, torch.zeros_like(reward))
 def support_contact_excess(
    env: ManagerBasedRLEnv,
    sensor_cfg: SceneEntityCfg,
    force_threshold: float = 120.0,
    max_contact_time: float = 0.6,
 ) -> torch.Tensor:
    """Penalize dragging or high-force support contacts."""
    contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name]
    force_norm = torch.linalg.norm(contact_sensor.data.net_forces_w[:, sensor_cfg.body_ids, :], dim=-1)
    contact_time = contact_sensor.data.current_contact_time[:, sensor_cfg.body_ids]
    force_excess = torch.clamp(force_norm - force_threshold, min=0.0) / force_threshold
    time_excess = torch.clamp(contact_time - max_contact_time, min=0.0)
    return torch.sum(force_excess + time_excess, dim=1)
 class ObstacleSupportEscapeReward(ManagerTermBase):
    """Reward a blocked -> brief support contact -> recovered progress sequence.
    The state machine prevents ordinary forward motion with incidental support contact from receiving the full reward.
    Obstacle evidence is inferred from low commanded progress together with either horizontal wheel force or wheel slip,
    so the actor can remain blind while the training reward uses contact information already available on the robot.
    """
    def __init__(self, cfg: RewTerm, env: ManagerBasedRLEnv):
        super().__init__(cfg, env)
        self._blocked_memory = torch.zeros(self.num_envs, device=self.device)
        self._support_window = torch.zeros(self.num_envs, device=self.device)
        self._recovery_reward_window = torch.zeros(self.num_envs, device=self.device)
        self._cooldown = torch.zeros(self.num_envs, device=self.device)
        self._previous_support_active = torch.zeros(self.num_envs, dtype=torch.bool, device=self.device)
    def reset(self, env_ids=None) -> None:
        if env_ids is None:
            env_ids = slice(None)
        self._blocked_memory[env_ids] = 0.0
        self._support_window[env_ids] = 0.0
        self._recovery_reward_window[env_ids] = 0.0
        self._cooldown[env_ids] = 0.0
        self._previous_support_active[env_ids] = False
    def __call__(
        self,
        env: ManagerBasedRLEnv,
        command_name: str,
        sensor_cfg: SceneEntityCfg,
        wheel_sensor_cfg: SceneEntityCfg,
        asset_cfg: SceneEntityCfg,
        wheel_asset_cfg: SceneEntityCfg,
        wheel_radius: float,
        support_force_threshold: float,
        wheel_obstacle_force_threshold: float,
        max_contact_time: float,
        min_command: float,
        stall_speed_ratio: float,
        min_stall_speed: float,
        slip_threshold: float,
        obstacle_memory_time: float,
        support_window_time: float,
        recovery_speed_ratio: float,
        min_recovery_speed: float,
        recovery_reward_time: float,
        cooldown_time: float,
        recovery_bonus: float,
        max_reward_speed: float,
    ) -> torch.Tensor:
        asset: Articulation = env.scene[asset_cfg.name]
        contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name]
        dt = float(env.step_dt)
        command_x = env.command_manager.get_command(command_name)[:, 0]
        command_speed = torch.abs(command_x)
        direction = torch.where(command_x >= 0.0, torch.ones_like(command_x), -torch.ones_like(command_x))
        progress_speed = direction * asset.data.root_lin_vel_b[:, 0]
        commanded_motion = command_speed > min_command
        wheel_forces = contact_sensor.data.net_forces_w[:, wheel_sensor_cfg.body_ids, :]
        max_horizontal_wheel_force = torch.linalg.norm(wheel_forces[..., :2], dim=-1).amax(dim=1)
        wheel_joint_vel = asset.data.joint_vel[:, wheel_asset_cfg.joint_ids]
        mean_wheel_surface_speed = torch.mean(torch.abs(wheel_joint_vel), dim=1) * wheel_radius
        wheel_slip = torch.clamp(mean_wheel_surface_speed - torch.clamp(progress_speed, min=0.0), min=0.0)
        stall_limit = torch.maximum(stall_speed_ratio * command_speed, torch.full_like(command_speed, min_stall_speed))
        obstacle_evidence = commanded_motion & (progress_speed < stall_limit) & (
            (max_horizontal_wheel_force > wheel_obstacle_force_threshold) | (wheel_slip > slip_threshold)
        )
        self._blocked_memory = torch.clamp(self._blocked_memory - dt, min=0.0)
        self._support_window = torch.clamp(self._support_window - dt, min=0.0)
        self._recovery_reward_window = torch.clamp(self._recovery_reward_window - dt, min=0.0)
        self._cooldown = torch.clamp(self._cooldown - dt, min=0.0)
        self._blocked_memory = torch.where(
            obstacle_evidence, torch.full_like(self._blocked_memory, obstacle_memory_time), self._blocked_memory
        )
        support_force = torch.linalg.norm(contact_sensor.data.net_forces_w[:, sensor_cfg.body_ids, :], dim=-1)
        support_time = contact_sensor.data.current_contact_time[:, sensor_cfg.body_ids]
        support_active = torch.any(support_force > support_force_threshold, dim=1)
        brief_support = torch.any(
            (support_force > support_force_threshold) & (support_time <= max_contact_time), dim=1
        )
        support_onset = brief_support & ~self._previous_support_active
        valid_support_onset = support_onset & (self._blocked_memory > 0.0) & (self._cooldown <= 0.0)
        was_waiting_for_recovery = self._support_window > 0.0
        self._support_window = torch.where(
            valid_support_onset, torch.full_like(self._support_window, support_window_time), self._support_window
        )
        recovery_threshold = torch.maximum(
            recovery_speed_ratio * command_speed, torch.full_like(command_speed, min_recovery_speed)
        )
        excessive_support = torch.any(
            (support_force > support_force_threshold) & (support_time > max_contact_time), dim=1
        )
        recovered = (
            was_waiting_for_recovery
            & commanded_motion
            & (progress_speed >= recovery_threshold)
            & ~excessive_support
        )
        self._recovery_reward_window = torch.where(
            recovered,
            torch.full_like(self._recovery_reward_window, recovery_reward_time),
            self._recovery_reward_window,
        )
        self._support_window = torch.where(recovered, torch.zeros_like(self._support_window), self._support_window)
        self._cooldown = torch.where(recovered, torch.full_like(self._cooldown, cooldown_time), self._cooldown)
        recovery_progress = torch.clamp(progress_speed, min=0.0, max=max_reward_speed) / max_reward_speed
        # Delay the discrete bonus until recovery is observed. This makes support contact alone worthless and
        # prevents the policy from farming reward by repeatedly tapping the body against an obstacle.
        reward = recovery_bonus * recovered.float() / dt
        reward += torch.where(
            self._recovery_reward_window > 0.0, recovery_progress, torch.zeros_like(recovery_progress)
        )
        self._previous_support_active = support_active
        return reward
 def track_lin_vel_xy_yaw_frame_exp(
    env, std: float, command_name: str, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")
 ) -> torch.Tensor:
    """Reward tracking of linear velocity commands (xy axes) in the gravity aligned robot frame using exponential kernel."""
    # extract the used quantities (to enable type-hinting)
    asset = env.scene[asset_cfg.name]
    vel_yaw = quat_apply_inverse(yaw_quat(asset.data.root_quat_w), asset.data.root_lin_vel_w[:, :3])
    lin_vel_error = torch.sum(
        torch.square(env.command_manager.get_command(command_name)[:, :2] - vel_yaw[:, :2]), dim=1
    )
    reward = torch.exp(-lin_vel_error / std**2)
    # reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
    return reward
 def track_ang_vel_z_world_exp(
    env, command_name: str, std: float, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")
 ) -> torch.Tensor:
    """Reward tracking of angular velocity commands (yaw) in world frame using exponential kernel."""
    # extract the used quantities (to enable type-hinting)
    asset = env.scene[asset_cfg.name]
    ang_vel_error = torch.square(env.command_manager.get_command(command_name)[:, 2] - asset.data.root_ang_vel_w[:, 2])
    reward = torch.exp(-ang_vel_error / std**2)
    reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
    return reward
 def joint_power(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
    """Reward joint_power"""
    # extract the used quantities (to enable type-hinting)
    asset: Articulation = env.scene[asset_cfg.name]
    # compute the reward
    reward = torch.sum(
        torch.abs(asset.data.joint_vel[:, asset_cfg.joint_ids] * asset.data.applied_torque[:, asset_cfg.joint_ids]),
        dim=1,
    )
    return reward
 def stand_still_without_cmd(
    env: ManagerBasedRLEnv,
    command_name: str,
    command_threshold: float,
    asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"),
 ) -> torch.Tensor:
    """Penalize joint positions that deviate from the default one when no command."""
    # extract the used quantities (to enable type-hinting)
    asset: Articulation = env.scene[asset_cfg.name]
    # compute out of limits constraints
    diff_angle = asset.data.joint_pos[:, asset_cfg.joint_ids] - asset.data.default_joint_pos[:, asset_cfg.joint_ids]
    reward = torch.sum(torch.abs(diff_angle), dim=1)
    reward *= torch.linalg.norm(env.command_manager.get_command(command_name), dim=1) < command_threshold
    # reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
    return reward
 def joint_pos_penalty(
    env: ManagerBasedRLEnv,
    command_name: str,
    asset_cfg: SceneEntityCfg,
    stand_still_scale: float,
    velocity_threshold: float,
    command_threshold: float,
 ) -> torch.Tensor:
    """Penalize joint position error from default on the articulation."""
    # extract the used quantities (to enable type-hinting)
    asset: Articulation = env.scene[asset_cfg.name]
    cmd = torch.linalg.norm(env.command_manager.get_command(command_name), dim=1)
    body_vel = torch.linalg.norm(asset.data.root_lin_vel_b[:, :2], dim=1)
    running_reward = torch.linalg.norm(
        (asset.data.joint_pos[:, asset_cfg.joint_ids] - asset.data.default_joint_pos[:, asset_cfg.joint_ids]), dim=1
    )
    reward = torch.where(
        torch.logical_or(cmd > command_threshold, body_vel > velocity_threshold),
        running_reward,
        stand_still_scale * running_reward,
    )
    # reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
    return reward
 def wheel_vel_penalty(
    env: ManagerBasedRLEnv,
    sensor_cfg: SceneEntityCfg,
    command_name: str,
    velocity_threshold: float,
    command_threshold: float,
    asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"),
 ) -> torch.Tensor:
    asset: Articulation = env.scene[asset_cfg.name]
    cmd = torch.linalg.norm(env.command_manager.get_command(command_name), dim=1)
    body_vel = torch.linalg.norm(asset.data.root_lin_vel_b[:, :2], dim=1)
    joint_vel = torch.abs(asset.data.joint_vel[:, asset_cfg.joint_ids])
    contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name]
    in_air = contact_sensor.compute_first_air(env.step_dt)[:, sensor_cfg.body_ids]
    running_reward = torch.sum(in_air * joint_vel, dim=1)
    standing_reward = torch.sum(joint_vel, dim=1)
    reward = torch.where(
        torch.logical_or(cmd > command_threshold, body_vel > velocity_threshold),
        running_reward,
        standing_reward,
    )
    return reward
 class GaitReward(ManagerTermBase):
    """Gait enforcing reward term for quadrupeds.
    This reward penalizes contact timing differences between selected foot pairs defined in :attr:`synced_feet_pair_names`
    to bias the policy towards a desired gait, i.e trotting, bounding, or pacing. Note that this reward is only for
    quadrupedal gaits with two pairs of synchronized feet.
    """
    def __init__(self, cfg: RewTerm, env: ManagerBasedRLEnv):
        """Initialize the term.
        Args:
            cfg: The configuration of the reward.
            env: The RL environment instance.
        """
        super().__init__(cfg, env)
        self.std: float = cfg.params["std"]
        self.command_name: str = cfg.params["command_name"]
        self.max_err: float = cfg.params["max_err"]
        self.velocity_threshold: float = cfg.params["velocity_threshold"]
        self.command_threshold: float = cfg.params["command_threshold"]
        self.contact_sensor: ContactSensor = env.scene.sensors[cfg.params["sensor_cfg"].name]
        self.asset: Articulation = env.scene[cfg.params["asset_cfg"].name]
        # match foot body names with corresponding foot body ids
        synced_feet_pair_names = cfg.params["synced_feet_pair_names"]
        if (
            len(synced_feet_pair_names) != 2
            or len(synced_feet_pair_names[0]) != 2
            or len(synced_feet_pair_names[1]) != 2
        ):
            raise ValueError("This reward only supports gaits with two pairs of synchronized feet, like trotting.")
        synced_feet_pair_0 = self.contact_sensor.find_bodies(synced_feet_pair_names[0])[0]
        synced_feet_pair_1 = self.contact_sensor.find_bodies(synced_feet_pair_names[1])[0]
        self.synced_feet_pairs = [synced_feet_pair_0, synced_feet_pair_1]
    def __call__(
        self,
        env: ManagerBasedRLEnv,
        std: float,
        command_name: str,
        max_err: float,
        velocity_threshold: float,
        command_threshold: float,
        synced_feet_pair_names,
        asset_cfg: SceneEntityCfg,
        sensor_cfg: SceneEntityCfg,
    ) -> torch.Tensor:
        """Compute the reward.
        This reward is defined as a multiplication between six terms where two of them enforce pair feet
        being in sync and the other four rewards if all the other remaining pairs are out of sync
        Args:
            env: The RL environment instance.
        Returns:
            The reward value.
        """
        # for synchronous feet, the contact (air) times of two feet should match
        sync_reward_0 = self._sync_reward_func(self.synced_feet_pairs[0][0], self.synced_feet_pairs[0][1])
        sync_reward_1 = self._sync_reward_func(self.synced_feet_pairs[1][0], self.synced_feet_pairs[1][1])
        sync_reward = sync_reward_0 * sync_reward_1
        # for asynchronous feet, the contact time of one foot should match the air time of the other one
        async_reward_0 = self._async_reward_func(self.synced_feet_pairs[0][0], self.synced_feet_pairs[1][0])
        async_reward_1 = self._async_reward_func(self.synced_feet_pairs[0][1], self.synced_feet_pairs[1][1])
        async_reward_2 = self._async_reward_func(self.synced_feet_pairs[0][0], self.synced_feet_pairs[1][1])
        async_reward_3 = self._async_reward_func(self.synced_feet_pairs[1][0], self.synced_feet_pairs[0][1])
        async_reward = async_reward_0 * async_reward_1 * async_reward_2 * async_reward_3
        # only enforce gait if cmd > 0
        cmd = torch.linalg.norm(env.command_manager.get_command(self.command_name), dim=1)
        body_vel = torch.linalg.norm(self.asset.data.root_com_lin_vel_b[:, :2], dim=1)
        reward = torch.where(
            torch.logical_or(cmd > self.command_threshold, body_vel > self.velocity_threshold),
            sync_reward * async_reward,
            0.0,
        )
        # reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
        return reward
    """
    Helper functions.
    """
    def _sync_reward_func(self, foot_0: int, foot_1: int) -> torch.Tensor:
        """Reward synchronization of two feet."""
        air_time = self.contact_sensor.data.current_air_time
        contact_time = self.contact_sensor.data.current_contact_time
        # penalize the difference between the most recent air time and contact time of synced feet pairs.
        se_air = torch.clip(torch.square(air_time[:, foot_0] - air_time[:, foot_1]), max=self.max_err**2)
        se_contact = torch.clip(torch.square(contact_time[:, foot_0] - contact_time[:, foot_1]), max=self.max_err**2)
        return torch.exp(-(se_air + se_contact) / self.std)
    def _async_reward_func(self, foot_0: int, foot_1: int) -> torch.Tensor:
        """Reward anti-synchronization of two feet."""
        air_time = self.contact_sensor.data.current_air_time
        contact_time = self.contact_sensor.data.current_contact_time
        # penalize the difference between opposing contact modes air time of feet 1 to contact time of feet 2
        # and contact time of feet 1 to air time of feet 2) of feet pairs that are not in sync with each other.
        se_act_0 = torch.clip(torch.square(air_time[:, foot_0] - contact_time[:, foot_1]), max=self.max_err**2)
        se_act_1 = torch.clip(torch.square(contact_time[:, foot_0] - air_time[:, foot_1]), max=self.max_err**2)
        return torch.exp(-(se_act_0 + se_act_1) / self.std)
 def joint_mirror(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg, mirror_joints: list[list[str]]) -> torch.Tensor:
    # extract the used quantities (to enable type-hinting)
    asset: Articulation = env.scene[asset_cfg.name]
    if not hasattr(env, "joint_mirror_joints_cache") or env.joint_mirror_joints_cache is None:
        # Cache joint positions for all pairs
        env.joint_mirror_joints_cache = [
            [asset.find_joints(joint_name) for joint_name in joint_pair] for joint_pair in mirror_joints
        ]
    reward = torch.zeros(env.num_envs, device=env.device)
    # Iterate over all joint pairs
    for joint_pair in env.joint_mirror_joints_cache:
        # Calculate the difference for each pair and add to the total reward
        diff = torch.sum(
            torch.square(asset.data.joint_pos[:, joint_pair[0][0]] - asset.data.joint_pos[:, joint_pair[1][0]]),
            dim=-1,
        )
        reward += diff
    reward *= 1 / len(mirror_joints) if len(mirror_joints) > 0 else 0
    reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
    return reward
 def action_mirror(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg, mirror_joints: list[list[str]]) -> torch.Tensor:
    # extract the used quantities (to enable type-hinting)
    asset: Articulation = env.scene[asset_cfg.name]
    if not hasattr(env, "action_mirror_joints_cache") or env.action_mirror_joints_cache is None:
        # Cache joint positions for all pairs
        env.action_mirror_joints_cache = [
            [asset.find_joints(joint_name) for joint_name in joint_pair] for joint_pair in mirror_joints
        ]
    reward = torch.zeros(env.num_envs, device=env.device)
    # Iterate over all joint pairs
    for joint_pair in env.action_mirror_joints_cache:
        # Calculate the difference for each pair and add to the total reward
        diff = torch.sum(
            torch.square(
                torch.abs(env.action_manager.action[:, joint_pair[0][0]])
                - torch.abs(env.action_manager.action[:, joint_pair[1][0]])
            ),
            dim=-1,
        )
        reward += diff
    reward *= 1 / len(mirror_joints) if len(mirror_joints) > 0 else 0
    # reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
    return reward
 def action_sync(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg, joint_groups: list[list[str]]) -> torch.Tensor:
    # extract the used quantities (to enable type-hinting)
    asset: Articulation = env.scene[asset_cfg.name]
    # Cache joint indices if not already done
    if not hasattr(env, "action_sync_joint_cache") or env.action_sync_joint_cache is None:
        env.action_sync_joint_cache = [
            [asset.find_joints(joint_name) for joint_name in joint_group] for joint_group in joint_groups
        ]
    reward = torch.zeros(env.num_envs, device=env.device)
    # Iterate over each joint group
    for joint_group in env.action_sync_joint_cache:
        if len(joint_group) < 2:
            continue  # need at least 2 joints to compare
        # Get absolute actions for all joints in this group
        actions = torch.stack(
            [torch.abs(env.action_manager.action[:, joint[0]]) for joint in joint_group], dim=1
        )  # shape: (num_envs, num_joints_in_group)
        # Calculate mean action for each environment
        mean_actions = torch.mean(actions, dim=1, keepdim=True)
        # Calculate variance from mean for each joint
        variance = torch.mean(torch.square(actions - mean_actions), dim=1)
        # Add to reward (we want to minimize this variance)
        reward += variance.squeeze()
    reward *= 1 / len(joint_groups) if len(joint_groups) > 0 else 0
    # reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
    return reward
 # def feet_air_time(
 #     env: ManagerBasedRLEnv, command_name: str, sensor_cfg: SceneEntityCfg, threshold: float
 # ) -> torch.Tensor:
 #     """Reward long steps taken by the feet using L2-kernel.
 #     This function rewards the agent for taking steps that are longer than a threshold. This helps ensure
 #     that the robot lifts its feet off the ground and takes steps. The reward is computed as the sum of
 #     the time for which the feet are in the air.
 #     If the commands are small (i.e. the agent is not supposed to take a step), then the reward is zero.
 #     """
 #     # extract the used quantities (to enable type-hinting)
 #     contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name]
 #     # compute the reward
 #     first_contact = contact_sensor.compute_first_contact(env.step_dt)[:, sensor_cfg.body_ids]
 #     last_air_time = contact_sensor.data.last_air_time[:, sensor_cfg.body_ids]
 #     reward = torch.sum((last_air_time - threshold) * first_contact, dim=1)
 #     # no reward for zero command
 #     reward *= torch.norm(env.command_manager.get_command(command_name)[:, :2], dim=1) > 0.1
 #     # print(torch.norm(env.command_manager.get_command(command_name)[:, :2], dim=1), "command norm")
 #     reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
    # return reward
 # def feet_air_time(
 #     env: ManagerBasedRLEnv,
 #     asset_cfg: SceneEntityCfg,
 #     sensor_cfg: SceneEntityCfg,
 #     mode_time: float,
 #     velocity_threshold: float,
 # ) -> torch.Tensor:
 #     """Reward longer feet air and contact time."""
 #     # extract the used quantities (to enable type-hinting)
 #     contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name]
 #     asset: Articulation = env.scene[asset_cfg.name]
 #     if contact_sensor.cfg.track_air_time is False:
 #         raise RuntimeError("Activate ContactSensor's track_air_time!")
 #     # compute the reward
 #     current_air_time = contact_sensor.data.current_air_time[:, sensor_cfg.body_ids]
 #     current_contact_time = contact_sensor.data.current_contact_time[:, sensor_cfg.body_ids]
 #     t_max = torch.max(current_air_time, current_contact_time)
 #     t_min = torch.clip(t_max, max=mode_time)
 #     stance_cmd_reward = torch.clip(current_contact_time - current_air_time, -mode_time, mode_time)
 #     cmd = torch.norm(env.command_manager.get_command("base_velocity"), dim=1).unsqueeze(dim=1).expand(-1, 4)
 #     body_vel = torch.linalg.norm(asset.data.root_lin_vel_b[:, :2], dim=1).unsqueeze(dim=1).expand(-1, 4)
 #     reward = torch.where(
 #         torch.logical_or(cmd > 0.0, body_vel > velocity_threshold),
 #         torch.where(t_max < mode_time, t_min, 0),
 #         stance_cmd_reward,
 #     )
 #     return torch.sum(reward, dim=1)
 def feet_air_time_positive_biped(env, command_name: str, threshold: float, sensor_cfg: SceneEntityCfg) -> torch.Tensor:
    """Reward long steps taken by the feet for bipeds.
    This function rewards the agent for taking steps up to a specified threshold and also keep one foot at
    a time in the air.
    If the commands are small (i.e. the agent is not supposed to take a step), then the reward is zero.
    """
    contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name]
    # compute the reward
    air_time = contact_sensor.data.current_air_time[:, sensor_cfg.body_ids]
    contact_time = contact_sensor.data.current_contact_time[:, sensor_cfg.body_ids]
    in_contact = contact_time > 0.0
    in_mode_time = torch.where(in_contact, contact_time, air_time)
    single_stance = torch.sum(in_contact.int(), dim=1) == 1
    reward = torch.min(torch.where(single_stance.unsqueeze(-1), in_mode_time, 0.0), dim=1)[0]
    reward = torch.clamp(reward, max=threshold)
    # no reward for zero command
    reward *= torch.norm(env.command_manager.get_command(command_name)[:, :2], dim=1) > 0.1
    # reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
    return reward
 def feet_air_time_variance_penalty(env: ManagerBasedRLEnv, sensor_cfg: SceneEntityCfg) -> torch.Tensor:
    """Penalize variance in the amount of time each foot spends in the air/on the ground relative to each other"""
    # extract the used quantities (to enable type-hinting)
    contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name]
    # compute the reward
    last_air_time = contact_sensor.data.last_air_time[:, sensor_cfg.body_ids]
    last_contact_time = contact_sensor.data.last_contact_time[:, sensor_cfg.body_ids]
    reward = torch.var(torch.clip(last_air_time, max=0.5), dim=1) + torch.var(
        torch.clip(last_contact_time, max=0.5), dim=1)
    # print(last_air_time, "last air time")
    # print(last_contact_time, "last contact time")
    # reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
    return reward
 def feet_contact(
    env: ManagerBasedRLEnv, command_name: str, expect_contact_num: int, sensor_cfg: SceneEntityCfg
 ) -> torch.Tensor:
    """Reward feet contact"""
    # extract the used quantities (to enable type-hinting)
    contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name]
    # compute the reward
    contact = contact_sensor.compute_first_contact(env.step_dt)[:, sensor_cfg.body_ids]
    contact_num = torch.sum(contact, dim=1)
    reward = (contact_num != expect_contact_num).float()
    # no reward for zero command
    reward *= torch.linalg.norm(env.command_manager.get_command(command_name), dim=1) > 0.5
    reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
    return reward
 def feet_contact_without_cmd(env: ManagerBasedRLEnv, command_name: str, sensor_cfg: SceneEntityCfg) -> torch.Tensor:
    """Reward feet contact"""
    # extract the used quantities (to enable type-hinting)
    contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name]
    # compute the reward
    contact = contact_sensor.compute_first_contact(env.step_dt)[:, sensor_cfg.body_ids]
    # print(contact, "contact")
    reward = torch.sum(contact, dim=-1).float()
    # print(reward, "reward after sum")
    reward *= torch.linalg.norm(env.command_manager.get_command(command_name), dim=1) < 0.5
    # print(env.command_manager.get_command(command_name), "env.command_manager.get_command(command_name)")
    # print(reward, "reward after multiply")
    # reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
    return reward
 def feet_stumble(env: ManagerBasedRLEnv, sensor_cfg: SceneEntityCfg) -> torch.Tensor:
    # extract the used quantities (to enable type-hinting)
    contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name]
    forces_z = torch.abs(contact_sensor.data.net_forces_w[:, sensor_cfg.body_ids, 2])
    forces_xy = torch.linalg.norm(contact_sensor.data.net_forces_w[:, sensor_cfg.body_ids, :2], dim=2)
    # Penalize feet hitting vertical surfaces
    reward = torch.any(forces_xy > 4 * forces_z, dim=1).float()
    # reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
    return reward
 def feet_distance_y_exp(
    env: ManagerBasedRLEnv, stance_width: float, std: float, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")
 ) -> torch.Tensor:
    asset: RigidObject = env.scene[asset_cfg.name]
    cur_footsteps_translated = asset.data.body_link_pos_w[:, asset_cfg.body_ids, :] - asset.data.root_link_pos_w[
        :, :
    ].unsqueeze(1)
    n_feet = len(asset_cfg.body_ids)
    footsteps_in_body_frame = torch.zeros(env.num_envs, n_feet, 3, device=env.device)
    for i in range(n_feet):
        footsteps_in_body_frame[:, i, :] = math_utils.quat_apply(
            math_utils.quat_conjugate(asset.data.root_link_quat_w), cur_footsteps_translated[:, i, :]
        )
    side_sign = torch.tensor(
        [1.0 if i % 2 == 0 else -1.0 for i in range(n_feet)],
        device=env.device,
    )
    stance_width_tensor = stance_width * torch.ones([env.num_envs, 1], device=env.device)
    desired_ys = stance_width_tensor / 2 * side_sign.unsqueeze(0)
    stance_diff = torch.square(desired_ys - footsteps_in_body_frame[:, :, 1])
    reward = torch.exp(-torch.sum(stance_diff, dim=1) / (std**2))
    # reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
    return reward
 def feet_distance_xy_exp(
    env: ManagerBasedRLEnv,
    stance_width: float,
    stance_length: float,
    std: float,
    asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"),
 ) -> torch.Tensor:
    asset: RigidObject = env.scene[asset_cfg.name]
    # Compute the current footstep positions relative to the root
    cur_footsteps_translated = asset.data.body_link_pos_w[:, asset_cfg.body_ids, :] - asset.data.root_link_pos_w[
        :, :
    ].unsqueeze(1)
    footsteps_in_body_frame = torch.zeros(env.num_envs, 4, 3, device=env.device)
    for i in range(4):
        footsteps_in_body_frame[:, i, :] = math_utils.quat_apply(
            math_utils.quat_conjugate(asset.data.root_link_quat_w), cur_footsteps_translated[:, i, :]
        )
    # Desired x and y positions for each foot
    stance_width_tensor = stance_width * torch.ones([env.num_envs, 1], device=env.device)
    stance_length_tensor = stance_length * torch.ones([env.num_envs, 1], device=env.device)
    desired_xs = torch.cat(
        [stance_length_tensor / 2, stance_length_tensor / 2, -stance_length_tensor / 2, -stance_length_tensor / 2],
        dim=1,
    )
    desired_ys = torch.cat(
        [stance_width_tensor / 2, -stance_width_tensor / 2, stance_width_tensor / 2, -stance_width_tensor / 2], dim=1
    )
    # Compute differences in x and y
    stance_diff_x = torch.square(desired_xs - footsteps_in_body_frame[:, :, 0])
    stance_diff_y = torch.square(desired_ys - footsteps_in_body_frame[:, :, 1])
    # Combine x and y differences and compute the exponential penalty
    stance_diff = stance_diff_x + stance_diff_y
    reward = torch.exp(-torch.sum(stance_diff, dim=1) / std**2)
    # reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
    return reward
 def feet_height(
    env: ManagerBasedRLEnv,
    command_name: str,
    asset_cfg: SceneEntityCfg,
    target_height: float,
    tanh_mult: float,
 ) -> torch.Tensor:
    """Reward the swinging feet for clearing a specified height off the ground"""
    asset: RigidObject = env.scene[asset_cfg.name]
    foot_z_target_error = torch.square(asset.data.body_pos_w[:, asset_cfg.body_ids, 2] - target_height)
    # foot_velocity_tanh = torch.tanh(
    #     tanh_mult * torch.linalg.norm(asset.data.body_lin_vel_w[:, asset_cfg.body_ids, :2], dim=2)
    # )
    # reward = torch.sum(foot_z_target_error * foot_velocity_tanh, dim=1)
    reward = torch.sum(foot_z_target_error, dim=1)
    # print(foot_z_target_error, "foot_z_target_error")
    # no reward for zero command
    reward *= torch.linalg.norm(env.command_manager.get_command(command_name), dim=1) > 0.2
    # reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
    return reward
 def feet_height_body(
    env: ManagerBasedRLEnv,
    command_name: str,
    asset_cfg: SceneEntityCfg,
    target_height: float,
    tanh_mult: float,
 ) -> torch.Tensor:
    """Reward the swinging feet for clearing a specified height off the ground"""
    asset: RigidObject = env.scene[asset_cfg.name]
    cur_footpos_translated = asset.data.body_pos_w[:, asset_cfg.body_ids, :] - asset.data.root_pos_w[:, :].unsqueeze(1)
    footpos_in_body_frame = torch.zeros(env.num_envs, len(asset_cfg.body_ids), 3, device=env.device)
    cur_footvel_translated = asset.data.body_lin_vel_w[:, asset_cfg.body_ids, :] - asset.data.root_lin_vel_w[
        :, :
    ].unsqueeze(1)
    footvel_in_body_frame = torch.zeros(env.num_envs, len(asset_cfg.body_ids), 3, device=env.device)
    for i in range(len(asset_cfg.body_ids)):
        footpos_in_body_frame[:, i, :] = math_utils.quat_apply_inverse(
            asset.data.root_quat_w, cur_footpos_translated[:, i, :]
        )
        footvel_in_body_frame[:, i, :] = math_utils.quat_apply_inverse(
            asset.data.root_quat_w, cur_footvel_translated[:, i, :]
        )
    foot_z_target_error = torch.square(footpos_in_body_frame[:, :, 2] - target_height).view(env.num_envs, -1)
    foot_velocity_tanh = torch.tanh(tanh_mult * torch.norm(footvel_in_body_frame[:, :, :2], dim=2))
    reward = torch.sum(foot_z_target_error * foot_velocity_tanh, dim=1)
    reward *= torch.linalg.norm(env.command_manager.get_command(command_name), dim=1) > 0.1
    # reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
    return reward
 def feet_slide(
    env: ManagerBasedRLEnv, sensor_cfg: SceneEntityCfg, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")
 ) -> torch.Tensor:
    """Penalize feet sliding.
    This function penalizes the agent for sliding its feet on the ground. The reward is computed as the
    norm of the linear velocity of the feet multiplied by a binary contact sensor. This ensures that the
    agent is penalized only when the feet are in contact with the ground.
    """
    # Penalize feet sliding
    contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name]
    contacts = contact_sensor.data.net_forces_w_history[:, :, sensor_cfg.body_ids, :].norm(dim=-1).max(dim=1)[0] > 1.0
    asset: RigidObject = env.scene[asset_cfg.name]
    # feet_vel = asset.data.body_lin_vel_w[:, asset_cfg.body_ids, :2]
    # reward = torch.sum(feet_vel.norm(dim=-1) * contacts, dim=1)
    cur_footvel_translated = asset.data.body_lin_vel_w[:, asset_cfg.body_ids, :] - asset.data.root_lin_vel_w[
        :, :
    ].unsqueeze(1)
    footvel_in_body_frame = torch.zeros(env.num_envs, len(asset_cfg.body_ids), 3, device=env.device)
    for i in range(len(asset_cfg.body_ids)):
        footvel_in_body_frame[:, i, :] = math_utils.quat_apply_inverse(
            asset.data.root_quat_w, cur_footvel_translated[:, i, :]
        )
    foot_leteral_vel = torch.sqrt(torch.sum(torch.square(footvel_in_body_frame[:, :, :2]), dim=2)).view(
        env.num_envs, -1
    )
    reward = torch.sum(foot_leteral_vel * contacts, dim=1)
    # reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
    return reward
 # def stand_still_joint_deviation_l1(
 #     env, command_name: str, command_threshold: float = 0.06, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")
 # ) -> torch.Tensor:
 #     """Penalize offsets from the default joint positions when the command is very small."""
    # command = env.command_manager.get_command(command_name)
 #     # Penalize motion when command is nearly zero.
 #     return joint_deviation_l1(env, asset_cfg) * (torch.norm(command[:, :], dim=1) < command_threshold)
 # def joint_deviation_l1(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
 #     """Penalize joint positions that deviate from the default one."""
 #     # extract the used quantities (to enable type-hinting)
 #     asset: Articulation = env.scene[asset_cfg.name]
 #     # compute out of limits constraints
 #     angle = asset.data.joint_pos[:, asset_cfg.joint_ids] - asset.data.default_joint_pos[:, asset_cfg.joint_ids]
 #     return torch.sum(torch.abs(angle), dim=1)
 # def smoothness_1(env: ManagerBasedRLEnv) -> torch.Tensor:
 #     # Penalize changes in actions
 #     diff = torch.square(env.action_manager.action - env.action_manager.prev_action)
 #     diff = diff * (env.action_manager.prev_action[:, :] != 0)  # ignore first step
 #     return torch.sum(diff, dim=1)
 # def joint_acc_l2_new(env: ManagerBasedRLEnv) -> torch.Tensor:
 # def smoothness_2(env: ManagerBasedRLEnv) -> torch.Tensor:
 #     # Penalize changes in actions
 #     diff = torch.square(env.action_manager.action - 2 * env.action_manager.prev_action + env.action_manager.prev_prev_action)
 #     diff = diff * (env.action_manager.prev_action[:, :] != 0)  # ignore first step
 #     diff = diff * (env.action_manager.prev_prev_action[:, :] != 0)  # ignore second step
 #     # print(torch.sum(diff, dim=1), "smoothness l2")
 #     return torch.sum(diff, dim=1)
 def upward(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
    """Penalize z-axis base linear velocity using L2 squared kernel."""
    # extract the used quantities (to enable type-hinting)
    asset: RigidObject = env.scene[asset_cfg.name]
    reward = torch.square(1 - asset.data.projected_gravity_b[:, 2])
    return reward
 def base_height_l2(
    env: ManagerBasedRLEnv,
    target_height: float,
    asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"),
    sensor_cfg: SceneEntityCfg | None = None,
 ) -> torch.Tensor:
    """Penalize asset height from its target using L2 squared kernel.
    Note:
        For flat terrain, target height is in the world frame. For rough terrain,
        sensor readings can adjust the target height to account for the terrain.
    """
    # extract the used quantities (to enable type-hinting)
    asset: RigidObject = env.scene[asset_cfg.name]
    if sensor_cfg is not None:
        sensor: RayCaster = env.scene[sensor_cfg.name]
        # Adjust the target height using the sensor data
        ray_hits = sensor.data.ray_hits_w[..., 2]
        if torch.isnan(ray_hits).any() or torch.isinf(ray_hits).any() or torch.max(torch.abs(ray_hits)) > 1e6:
            adjusted_target_height = asset.data.root_link_pos_w[:, 2]
        else:
            adjusted_target_height = target_height + torch.mean(ray_hits, dim=1)
    else:
        # Use the provided target height directly for flat terrain
        adjusted_target_height = target_height
    # Compute the L2 squared penalty
    reward = torch.square(asset.data.root_pos_w[:, 2] - adjusted_target_height)
    # reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
    return reward
 def lin_vel_z_l2(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
    """Penalize z-axis base linear velocity using L2 squared kernel."""
    # extract the used quantities (to enable type-hinting)
    asset: RigidObject = env.scene[asset_cfg.name]
    reward = torch.square(asset.data.root_lin_vel_b[:, 2])
    # reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
    return reward
 def ang_vel_xy_l2(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
    """Penalize xy-axis base angular velocity using L2 squared kernel."""
    # extract the used quantities (to enable type-hinting)
    asset: RigidObject = env.scene[asset_cfg.name]
    reward = torch.sum(torch.square(asset.data.root_ang_vel_b[:, :2]), dim=1)
    # reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
    return reward
 def undesired_contacts(env: ManagerBasedRLEnv, threshold: float, sensor_cfg: SceneEntityCfg) -> torch.Tensor:
    """Penalize undesired contacts as the number of violations that are above a threshold."""
    # extract the used quantities (to enable type-hinting)
    contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name]
    # check if contact force is above threshold
    net_contact_forces = contact_sensor.data.net_forces_w_history
    is_contact = torch.max(torch.norm(net_contact_forces[:, :, sensor_cfg.body_ids], dim=-1), dim=1)[0] > threshold
    # sum over contacts for each environment
    reward = torch.sum(is_contact, dim=1).float()
    # reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
    return reward
 def flat_orientation_l2(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
    """Penalize non-flat base orientation using L2 squared kernel.
    This is computed by penalizing the xy-components of the projected gravity vector.
    """
    # extract the used quantities (to enable type-hinting)
    asset: RigidObject = env.scene[asset_cfg.name]
    reward = torch.sum(torch.square(asset.data.projected_gravity_b[:, :2]), dim=1)
    # reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
    return reward
 def feet_air_time_including_ang_z(
    env: ManagerBasedRLEnv, command_name: str, sensor_cfg: SceneEntityCfg, threshold: float
 ) -> torch.Tensor:
    """Reward long steps taken by the feet using L2-kernel.
    This function rewards the agent for taking steps that are longer than a threshold. This helps ensure
    that the robot lifts its feet off the ground and takes steps. The reward is computed as the sum of
    the time for which the feet are in the air.
    If the commands are small (i.e. the agent is not supposed to take a step), then the reward is zero.
    """
    # extract the used quantities (to enable type-hinting)
    contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name]
    # compute the reward
    first_contact = contact_sensor.compute_first_contact(env.step_dt)[:, sensor_cfg.body_ids]
    last_air_time = contact_sensor.data.last_air_time[:, sensor_cfg.body_ids]
    reward = torch.sum((last_air_time - threshold) * first_contact, dim=1)
    # no reward for zero command
    reward *= torch.norm(env.command_manager.get_command(command_name), dim=1) > 0.1
    # reward *= torch.norm(env.command_manager.get_command(command_name)[:, :3], dim=1) > 0.1
    return reward
 def lin_vel_xy_l2_with_ang_z_command(
    env: ManagerBasedRLEnv,
    command_name: str,
    command_threshold: float,
    asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"),
    ) -> torch.Tensor:
    """Penalize xy-axis base linear velocity using L2 squared kernel if command is ang_vel_z."""
    # extract the used quantities (to enable type-hinting)
    asset: RigidObject = env.scene[asset_cfg.name]
    # reward = torch.square(asset.data.root_lin_vel_b[:, 2])
    reward = torch.sum(torch.square(asset.data.root_lin_vel_b[:, :2]), dim=1)
    command = env.command_manager.get_command(command_name)
    reward *= (torch.sum(torch.square(command[:, 2:]), dim=1) > command_threshold) & \
            (torch.sum(torch.square(command[:, :2]), dim=1) < command_threshold)
    # reward *= torch.sum(torch.square(env.command_manager.get_command(command_name)[:, 2:]), dim=1) > command_threshold
    # reward *= torch.clamp(-env.scene["robot"].data.projected_gravity_b[:, 2], 0, 0.7) / 0.7
    return reward
@@ -0,0 +1,866 @@
 # Copyright (c) 2025 Deep Robotics
 # SPDX-License-Identifier: BSD 3-Clause
 # Copyright (c) 2024-2025 Ziqi Fan
 # SPDX-License-Identifier: Apache-2.0
 # Copyright (c) 2022-2025, The Isaac Lab Project Developers.
 # All rights reserved.
 #
 # SPDX-License-Identifier: BSD-3-Clause
 import inspect
 import math
 import sys
 from dataclasses import MISSING
 import isaaclab.sim as sim_utils
 from isaaclab.assets import ArticulationCfg, AssetBaseCfg
 from isaaclab.envs import ManagerBasedRLEnvCfg
 from isaaclab.managers import CurriculumTermCfg as CurrTerm
 from isaaclab.managers import EventTermCfg as EventTerm
 from isaaclab.managers import ObservationGroupCfg as ObsGroup
 from isaaclab.managers import ObservationTermCfg as ObsTerm
 from isaaclab.managers import RewardTermCfg as RewTerm
 from isaaclab.managers import SceneEntityCfg
 from isaaclab.managers import TerminationTermCfg as DoneTerm
 from isaaclab.scene import InteractiveSceneCfg
 from isaaclab.sensors import ContactSensorCfg, RayCasterCfg, patterns
 from isaaclab.terrains import TerrainImporterCfg
 from isaaclab.utils import configclass
 from isaaclab.utils.noise import AdditiveUniformNoiseCfg as Unoise
 import rl_training.tasks.manager_based.locomotion.velocity.mdp as mdp
 ##
 # Pre-defined configs
 ##
 from isaaclab.terrains.config.rough import ROUGH_TERRAINS_CFG  # isort: skip
 ##
 # Scene definition
 ##
@configclass
 class MySceneCfg(InteractiveSceneCfg):
    """Configuration for the terrain scene with a legged robot."""
    # ground terrain
    terrain = TerrainImporterCfg(
        prim_path="/World/ground",
        terrain_type="generator",
        terrain_generator=ROUGH_TERRAINS_CFG,
        max_init_terrain_level=5,
        collision_group=-1,
        physics_material=sim_utils.RigidBodyMaterialCfg(
            friction_combine_mode="multiply",
            restitution_combine_mode="multiply",
            static_friction=1.0,
            dynamic_friction=1.0,
            restitution=1.0,
        ),
        # Keep scene startup independent of NVIDIA's online asset server.
        # The robot visual geometry is still loaded from the STL meshes in its URDF.
        visual_material=sim_utils.PreviewSurfaceCfg(
            diffuse_color=(0.18, 0.20, 0.22),
            roughness=0.9,
        ),
        debug_vis=False,
    )
    # robots
    robot: ArticulationCfg = MISSING
    # sensors
    height_scanner = RayCasterCfg(
        prim_path="{ENV_REGEX_NS}/Robot/base",
        offset=RayCasterCfg.OffsetCfg(pos=(0.0, 0.0, 20.0)),
        ray_alignment="yaw",
        pattern_cfg=patterns.GridPatternCfg(resolution=0.1, size=[1.6, 1.0]),
        debug_vis=False,
        mesh_prim_paths=["/World/ground"],
    )
    height_scanner_base = RayCasterCfg(
        prim_path="{ENV_REGEX_NS}/Robot/base",
        offset=RayCasterCfg.OffsetCfg(pos=(0.0, 0.0, 20.0)),
        ray_alignment="yaw",
        pattern_cfg=patterns.GridPatternCfg(resolution=0.05, size=(0.1, 0.1)),
        debug_vis=False,
        mesh_prim_paths=["/World/ground"],
    )
    contact_forces = ContactSensorCfg(prim_path="{ENV_REGEX_NS}/Robot/.*", history_length=3, track_air_time=True)
    # lights
    sky_light = AssetBaseCfg(
        prim_path="/World/skyLight",
        spawn=sim_utils.DomeLightCfg(
            intensity=1200.0,
            color=(0.75, 0.75, 0.75),
        ),
    )
 ##
 # MDP settings
 ##
@configclass
 class CommandsCfg:
    """Command specifications for the MDP."""
    base_velocity = mdp.UniformThresholdVelocityCommandCfg(
        asset_name="robot",
        resampling_time_range=(10.0, 10.0),
        rel_standing_envs=0.02,
        rel_heading_envs=1.0,
        heading_command=True,
        heading_control_stiffness=0.5,
        debug_vis=True,
        ranges=mdp.UniformThresholdVelocityCommandCfg.Ranges(
            lin_vel_x=(-1.0, 1.0), lin_vel_y=(-1.0, 1.0), ang_vel_z=(-1.0, 1.0), heading=(-math.pi, math.pi)
        ),
    )
@configclass
 class ActionsCfg:
    """Action specifications for the MDP."""
    joint_pos = mdp.JointPositionActionCfg(
        asset_name="robot", joint_names=[".*"], scale=0.5, use_default_offset=True, clip=None, preserve_order=True
    )
@configclass
 class ObservationsCfg:
    """Observation specifications for the MDP."""
    @configclass
    class PolicyCfg(ObsGroup):
        """Observations for policy group."""
        # observation terms (order preserved)
        base_lin_vel = ObsTerm(
            func=mdp.base_lin_vel,
            noise=Unoise(n_min=-0.1, n_max=0.1),
            clip=(-100.0, 100.0),
            scale=1.0,
        )
        base_ang_vel = ObsTerm(
            func=mdp.base_ang_vel,
            noise=Unoise(n_min=-0.2, n_max=0.2),
            clip=(-100.0, 100.0),
            scale=1.0,
        )
        projected_gravity = ObsTerm(
            func=mdp.projected_gravity,
            noise=Unoise(n_min=-0.05, n_max=0.05),
            clip=(-100.0, 100.0),
            scale=1.0,
        )
        velocity_commands = ObsTerm(
            func=mdp.generated_commands,
            params={"command_name": "base_velocity"},
            clip=(-100.0, 100.0),
            scale=1.0,
        )
        joint_pos = ObsTerm(
            func=mdp.joint_pos_rel,
            params={"asset_cfg": SceneEntityCfg("robot", joint_names=".*", preserve_order=True)},
            noise=Unoise(n_min=-0.01, n_max=0.01),
            clip=(-100.0, 100.0),
            scale=1.0,
        )
        joint_vel = ObsTerm(
            func=mdp.joint_vel_rel,
            params={"asset_cfg": SceneEntityCfg("robot", joint_names=".*", preserve_order=True)},
            noise=Unoise(n_min=-1.5, n_max=1.5),
            clip=(-100.0, 100.0),
            scale=1.0,
        )
        actions = ObsTerm(
            func=mdp.last_action,
            clip=(-100.0, 100.0),
            scale=1.0,
        )
        height_scan = ObsTerm(
            func=mdp.height_scan,
            params={"sensor_cfg": SceneEntityCfg("height_scanner")},
            noise=Unoise(n_min=-0.1, n_max=0.1),
            clip=(-1.0, 1.0),
            scale=1.0,
        )
        def __post_init__(self):
            self.enable_corruption = True
            self.concatenate_terms = True
    @configclass
    class CriticCfg(ObsGroup):
        """Observations for critic group."""
        # observation terms (order preserved)
        base_lin_vel = ObsTerm(
            func=mdp.base_lin_vel,
            clip=(-100.0, 100.0),
            scale=1.0,
        )
        base_ang_vel = ObsTerm(
            func=mdp.base_ang_vel,
            clip=(-100.0, 100.0),
            scale=1.0,
        )
        projected_gravity = ObsTerm(
            func=mdp.projected_gravity,
            clip=(-100.0, 100.0),
            scale=1.0,
        )
        velocity_commands = ObsTerm(
            func=mdp.generated_commands,
            params={"command_name": "base_velocity"},
            clip=(-100.0, 100.0),
            scale=1.0,
        )
        joint_pos = ObsTerm(
            func=mdp.joint_pos_rel,
            params={"asset_cfg": SceneEntityCfg("robot", joint_names=".*", preserve_order=True)},
            clip=(-100.0, 100.0),
            scale=1.0,
        )
        joint_vel = ObsTerm(
            func=mdp.joint_vel_rel,
            params={"asset_cfg": SceneEntityCfg("robot", joint_names=".*", preserve_order=True)},
            clip=(-100.0, 100.0),
            scale=1.0,
        )
        actions = ObsTerm(
            func=mdp.last_action,
            clip=(-100.0, 100.0),
            scale=1.0,
        )
        height_scan = ObsTerm(
            func=mdp.height_scan,
            params={"sensor_cfg": SceneEntityCfg("height_scanner")},
            clip=(-1.0, 1.0),
            scale=1.0,
        )
        # joint_effort = ObsTerm(
        #     func=mdp.joint_effort,
        #     clip=(-100, 100),
        #     scale=0.01,
        # )
        def __post_init__(self):
            self.enable_corruption = False
            self.concatenate_terms = True
    # observation groups
    policy: PolicyCfg = PolicyCfg()
    critic: CriticCfg = CriticCfg()
@configclass
 class EventCfg:
    """Configuration for events."""
    # startup
    randomize_rigid_body_material = EventTerm(
        func=mdp.randomize_rigid_body_material,
        mode="startup",
        params={
            "asset_cfg": SceneEntityCfg("robot", body_names=".*"),
            "static_friction_range": (0.3, 1.0),
            "dynamic_friction_range": (0.3, 0.8),
            "restitution_range": (0.0, 0.4),
            "num_buckets": 1024,
        },
    )
    randomize_rigid_body_mass = EventTerm(
        func=mdp.randomize_rigid_body_mass,
        mode="startup",
        params={
            "asset_cfg": SceneEntityCfg("robot", body_names=""),
            "mass_distribution_params": (0.85, 1.15),
            "operation": "scale",
            "recompute_inertia": True,
        },
    )
    randomize_rigid_body_mass_base = EventTerm(
        func=mdp.randomize_rigid_body_mass,
        mode="startup",
        params={
            "asset_cfg": SceneEntityCfg("robot", body_names=""),
            "mass_distribution_params": (-1.0, 3.0),
            "operation": "add",
            "recompute_inertia": True,
        },
    )
    randomize_rigid_body_inertia = EventTerm(
        func=mdp.randomize_rigid_body_inertia,
        mode="startup",
        params={
            "asset_cfg": SceneEntityCfg("robot", body_names=".*"),
            "inertia_distribution_params": (0.85, 1.15),
            "operation": "scale",
        },
    )
    randomize_com_positions = EventTerm(
        func=mdp.randomize_rigid_body_com,
        mode="startup",
        params={
            "asset_cfg": SceneEntityCfg("robot", body_names=".*"),
            "com_range": {"x": (-0.03, 0.03), "y": (-0.03, 0.03), "z": (-0.02, 0.02)},
        },
    )
    # reset
    randomize_apply_external_force_torque = EventTerm(
        func=mdp.apply_external_force_torque,
        mode="reset",
        params={
            "asset_cfg": SceneEntityCfg("robot", body_names=""),
            "force_range": (-10.0, 10.0),
            "torque_range": (-10.0, 10.0),
        },
    )
    randomize_reset_joints = EventTerm(
        func=mdp.reset_joints_by_scale,
        # func=mdp.reset_joints_by_offset,
        mode="reset",
        params={
            "position_range": (1.0, 1.0),
            "velocity_range": (0.0, 0.0),
        },
    )
    randomize_actuator_gains = EventTerm(
        func=mdp.randomize_actuator_gains,
        mode="reset",
        params={
            "asset_cfg": SceneEntityCfg("robot", joint_names=".*"),
            "stiffness_distribution_params": (0.85, 1.15),
            "damping_distribution_params": (0.85, 1.15),
            "operation": "scale",
            "distribution": "uniform",
        },
    )
    randomize_reset_base = EventTerm(
        func=mdp.reset_root_state_uniform,
        mode="reset",
        params={
            "pose_range": {"x": (-0.5, 0.5), "y": (-0.5, 0.5), "yaw": (-3.14, 3.14)},
            "velocity_range": {
                "x": (-0.5, 0.5),
                "y": (-0.5, 0.5),
                "z": (-0.5, 0.5),
                "roll": (-0.5, 0.5),
                "pitch": (-0.5, 0.5),
                "yaw": (-0.5, 0.5),
            },
        },
    )
    # interval
    randomize_push_robot = EventTerm(
        func=mdp.push_by_setting_velocity,
        mode="interval",
        interval_range_s=(10.0, 15.0),
        params={"velocity_range": {"x": (-0.5, 0.5), "y": (-0.5, 0.5)}},
    )
@configclass
 class RewardsCfg:
    """Reward terms for the MDP."""
    # General
    is_terminated = RewTerm(func=mdp.is_terminated, weight=0.0)
    # Root penalties
    lin_vel_z_l2 = RewTerm(func=mdp.lin_vel_z_l2, weight=0.0)
    ang_vel_xy_l2 = RewTerm(func=mdp.ang_vel_xy_l2, weight=0.0)
    flat_orientation_l2 = RewTerm(func=mdp.flat_orientation_l2, weight=0.0)
    base_height_l2 = RewTerm(
        func=mdp.base_height_l2,
        weight=0.0,
        params={
            "asset_cfg": SceneEntityCfg("robot", body_names=""),
            "sensor_cfg": SceneEntityCfg("height_scanner_base"),
            "target_height": 0.0,
        },
    )
    body_lin_acc_l2 = RewTerm(
        func=mdp.body_lin_acc_l2,
        weight=0.0,
        params={"asset_cfg": SceneEntityCfg("robot", body_names="")},
    )
    # Joint penalties
    joint_torques_l2 = RewTerm(
        func=mdp.joint_torques_l2, weight=0.0, params={"asset_cfg": SceneEntityCfg("robot", joint_names=".*")}
    )
    joint_vel_l2 = RewTerm(
        func=mdp.joint_vel_l2, weight=0.0, params={"asset_cfg": SceneEntityCfg("robot", joint_names=".*")}
    )
    joint_acc_l2 = RewTerm(
        func=mdp.joint_acc_l2, weight=0.0, params={"asset_cfg": SceneEntityCfg("robot", joint_names=".*")}
    )
    joint_deviation_l1 = RewTerm(
            func=mdp.joint_deviation_l1,
            weight=0.0,
            params={"asset_cfg": SceneEntityCfg("robot", joint_names=".*")},
        )
    def create_joint_deviation_l1_rewterm(self, attr_name, weight, joint_names_pattern):
        rew_term = RewTerm(
            func=mdp.joint_deviation_l1,
            weight=weight,
            params={"asset_cfg": SceneEntityCfg("robot", joint_names=joint_names_pattern)},
        )
        setattr(self, attr_name, rew_term)
    joint_pos_limits = RewTerm(
        func=mdp.joint_pos_limits, weight=0.0, params={"asset_cfg": SceneEntityCfg("robot", joint_names=".*")}
    )
    joint_vel_limits = RewTerm(
        func=mdp.joint_vel_limits,
        weight=0.0,
        params={"asset_cfg": SceneEntityCfg("robot", joint_names=".*"), "soft_ratio": 1.0},
    )
    joint_power = RewTerm(
        func=mdp.joint_power,
        weight=0.0,
        params={
            "asset_cfg": SceneEntityCfg("robot", joint_names=".*"),
        },
    )
    stand_still_without_cmd = RewTerm(
        func=mdp.stand_still_without_cmd,
        weight=0.0,
        params={
            "command_name": "base_velocity",
            "command_threshold": 0.1,
            "asset_cfg": SceneEntityCfg("robot", joint_names=".*"),
        },
    )
    joint_pos_penalty = RewTerm(
        func=mdp.joint_pos_penalty,
        weight=0.0,
        params={
            "command_name": "base_velocity",
            "asset_cfg": SceneEntityCfg("robot", joint_names=".*"),
            "stand_still_scale": 5.0,
            "velocity_threshold": 0.5,
            "command_threshold": 0.1,
        },
    )
    hipx_joint_pos_penalty = RewTerm(
        func=mdp.joint_pos_penalty,
        weight=0.0,
        params={
            "command_name": "base_velocity",
            "asset_cfg": SceneEntityCfg("robot", joint_names=".*"),
            "stand_still_scale": 5.0,
            "velocity_threshold": 0.5,
            "command_threshold": 0.1,
        },
    )
    hipy_joint_pos_penalty = RewTerm(
        func=mdp.joint_pos_penalty,
        weight=0.0,
        params={
            "command_name": "base_velocity",
            "asset_cfg": SceneEntityCfg("robot", joint_names=".*"),
            "stand_still_scale": 5.0,
            "velocity_threshold": 0.5,
            "command_threshold": 0.1,
        },
    )
    knee_joint_pos_penalty = RewTerm(
        func=mdp.joint_pos_penalty,
        weight=0.0,
        params={
            "command_name": "base_velocity",
            "asset_cfg": SceneEntityCfg("robot", joint_names=".*"),
            "stand_still_scale": 5.0,
            "velocity_threshold": 0.5,
            "command_threshold": 0.1,
        },
    )
    wheel_vel_penalty = RewTerm(
        func=mdp.wheel_vel_penalty,
        weight=0.0,
        params={
            "asset_cfg": SceneEntityCfg("robot", joint_names=""),
            "sensor_cfg": SceneEntityCfg("contact_forces", body_names=""),
            "command_name": "base_velocity",
            "velocity_threshold": 0.5,
            "command_threshold": 0.1,
        },
    )
    joint_mirror = RewTerm(
        func=mdp.joint_mirror,
        weight=0.0,
        params={
            "asset_cfg": SceneEntityCfg("robot"),
            "mirror_joints": [["FR.*", "RL.*"], ["FL.*", "RR.*"]],
        },
    )
    action_mirror = RewTerm(
        func=mdp.action_mirror,
        weight=0.0,
        params={
            "asset_cfg": SceneEntityCfg("robot"),
            "mirror_joints": [["FR.*", "RL.*"], ["FL.*", "RR.*"]],
        },
    )
    action_sync = RewTerm(
        func=mdp.action_sync,
        weight=0.0,
        params={
            "asset_cfg": SceneEntityCfg("robot"),
            "joint_groups": [
                ["FR_hip_joint", "FL_hip_joint", "RL_hip_joint", "RR_hip_joint"],
                ["FR_thigh_joint", "FL_thigh_joint", "RL_thigh_joint", "RR_thigh_joint"],
                ["FR_calf_joint", "FL_calf_joint", "RL_calf_joint", "RR_calf_joint"],
            ],
        },
    )
    # Action penalties
    applied_torque_limits = RewTerm(
        func=mdp.applied_torque_limits,
        weight=0.0,
        params={"asset_cfg": SceneEntityCfg("robot", joint_names=".*")},
    )
    action_rate_l2 = RewTerm(func=mdp.action_rate_l2, weight=0.0)
    # smoothness_1 = RewTerm(func=mdp.smoothness_1, weight=0.0)  # Same as action_rate_l2
    # smoothness_2 = RewTerm(func=mdp.smoothness_2, weight=0.0)  # Unvaliable now
    # Contact sensor
    undesired_contacts = RewTerm(
        func=mdp.undesired_contacts,
        weight=0.0,
        params={
            "sensor_cfg": SceneEntityCfg("contact_forces", body_names=""),
            "threshold": 1.0,
        },
    )
    contact_forces = RewTerm(
        func=mdp.contact_forces,
        weight=0.0,
        params={"sensor_cfg": SceneEntityCfg("contact_forces", body_names=""), "threshold": 100.0},
    )
    # Velocity-tracking rewards
    track_lin_vel_xy_exp = RewTerm(
        func=mdp.track_lin_vel_xy_exp, weight=0.0, params={"command_name": "base_velocity", "std": math.sqrt(0.5)}
    )
    track_ang_vel_z_exp = RewTerm(
        func=mdp.track_ang_vel_z_exp, weight=0.0, params={"command_name": "base_velocity", "std": math.sqrt(0.5)}
    )
    # Others
    # feet_air_time = RewTerm(
    #     func=mdp.feet_air_time,
    #     weight=0.0,
    #     params={
    #         "command_name": "base_velocity",
    #         "threshold": 0.5,
    #         "sensor_cfg": SceneEntityCfg("contact_forces", body_names=""),
    #     },
    # )
    feet_air_time = RewTerm(
        func=mdp.feet_air_time_including_ang_z,
        weight=0.0,
        params={
            "command_name": "base_velocity",
            "threshold": 0.5,
            "sensor_cfg": SceneEntityCfg("contact_forces", body_names=""),
        },
    )
    feet_air_time_variance = RewTerm(
        func=mdp.feet_air_time_variance_penalty,
        weight=0,
        params={"sensor_cfg": SceneEntityCfg("contact_forces", body_names="")},
    )
    feet_gait = RewTerm(
        func=mdp.GaitReward,
        weight=0.0,
        params={
            "std": math.sqrt(0.5),
            "command_name": "base_velocity",
            "max_err": 0.2,
            "velocity_threshold": 0.5,
            "command_threshold": 0.1,
            "synced_feet_pair_names": (("", ""), ("", "")),
            "asset_cfg": SceneEntityCfg("robot"),
            "sensor_cfg": SceneEntityCfg("contact_forces"),
        },
    )
    feet_contact = RewTerm(
        func=mdp.feet_contact,
        weight=0.0,
        params={
            "sensor_cfg": SceneEntityCfg("contact_forces", body_names=""),
            "command_name": "base_velocity",
            "expect_contact_num": 2,
        },
    )
    feet_contact_without_cmd = RewTerm(
        func=mdp.feet_contact_without_cmd,
        weight=0.0,
        params={
            "sensor_cfg": SceneEntityCfg("contact_forces", body_names=""),
            "command_name": "base_velocity",
        },
    )
    feet_stumble = RewTerm(
        func=mdp.feet_stumble,
        weight=0.0,
        params={
            "sensor_cfg": SceneEntityCfg("contact_forces", body_names=""),
        },
    )
    feet_slide = RewTerm(
        func=mdp.feet_slide,
        weight=0.0,
        params={
            "sensor_cfg": SceneEntityCfg("contact_forces", body_names=""),
            "asset_cfg": SceneEntityCfg("robot", body_names=""),
        },
    )
    stand_still = RewTerm(
        func=mdp.stand_still_joint_deviation_l1,
        weight=0,
        params={
            "command_name": "base_velocity",
            "asset_cfg": SceneEntityCfg("robot", joint_names=""),
        },
    ) # negetive
    feet_height = RewTerm(
        func=mdp.feet_height,
        weight=0.0,
        params={
            "asset_cfg": SceneEntityCfg("robot", body_names=""),
            "tanh_mult": 2.0,
            "target_height": 0.05,
            "command_name": "base_velocity",
        },
    )
    feet_height_body = RewTerm(
        func=mdp.feet_height_body,
        weight=0.0,
        params={
            "asset_cfg": SceneEntityCfg("robot", body_names=""),
            "tanh_mult": 2.0,
            "target_height": -0.3,
            "command_name": "base_velocity",
        },
    )
    feet_distance_y_exp = RewTerm(
        func=mdp.feet_distance_y_exp,
        weight=0.0,
        params={
            "std": math.sqrt(0.25),
            "asset_cfg": SceneEntityCfg("robot", body_names=""),
            "stance_width": float,
        },
    )
    # feet_distance_xy_exp = RewTerm(
    #     func=mdp.feet_distance_xy_exp,
    #     weight=0.0,
    #     params={
    #         "std": math.sqrt(0.25),
    #         "asset_cfg": SceneEntityCfg("robot", body_names=""),
    #         "stance_length": float,
    #         "stance_width": float,
    #     },
    # )
    upward = RewTerm(func=mdp.upward, weight=0.0)
    # lin_vel_xy_l2_with_ang_z_command = RewTerm(
    #     func=mdp.lin_vel_xy_l2_with_ang_z_command,
    #     weight=0,
    #     params={
    #         "command_name": "base_velocity",
    #         "command_threshold": 0.1,
    #     },
    # ) # negetive
@configclass
 class TerminationsCfg:
    """Termination terms for the MDP."""
    # MDP terminations
    time_out = DoneTerm(func=mdp.time_out, time_out=True)
    # command_resample
    terrain_out_of_bounds = DoneTerm(
        func=mdp.terrain_out_of_bounds,
        params={"asset_cfg": SceneEntityCfg("robot"), "distance_buffer": 3.0},
        time_out=True,
    )
    # Contact sensor
    illegal_contact = DoneTerm(
        func=mdp.illegal_contact,
        params={"sensor_cfg": SceneEntityCfg("contact_forces", body_names=""), "threshold": 1.0},
    )
    bad_orientation_2 = DoneTerm(func=mdp.bad_orientation_2)
@configclass
 class CurriculumCfg:
    """Curriculum terms for the MDP."""
    terrain_levels = CurrTerm(func=mdp.terrain_levels_vel)
    command_levels = CurrTerm(
        func=mdp.command_levels_vel,
        params={
            "reward_term_name": "track_lin_vel_xy_exp",
            "range_multiplier": (0.1, 1.0),
        },
    )
 ##
 # Environment configuration
 ##
@configclass
 class LocomotionVelocityRoughEnvCfg(ManagerBasedRLEnvCfg):
    """Configuration for the locomotion velocity-tracking environment."""
    # Scene settings
    scene: MySceneCfg = MySceneCfg(num_envs=4096, env_spacing=2.5)
    # Basic settings
    observations: ObservationsCfg = ObservationsCfg()
    actions: ActionsCfg = ActionsCfg()
    commands: CommandsCfg = CommandsCfg()
    # MDP settings
    rewards: RewardsCfg = RewardsCfg()
    terminations: TerminationsCfg = TerminationsCfg()
    events: EventCfg = EventCfg()
    curriculum: CurriculumCfg = CurriculumCfg()
    def __post_init__(self):
        """Post initialization."""
        # general settings
        self.decimation = 4
        self.episode_length_s = 20.0
        # simulation settings
        self.sim.dt = 0.005
        self.sim.render_interval = self.decimation
        self.sim.physics_material = self.scene.terrain.physics_material
        self.sim.physx.gpu_max_rigid_patch_count = 10 * 2**15
        self.sim.physx.max_position_iteration_count = 4
        self.sim.physx.max_velocity_iteration_count = 1
        # update sensor update periods
        # we tick all the sensors based on the smallest update period (physics update period)
        if self.scene.height_scanner is not None:
            self.scene.height_scanner.update_period = self.decimation * self.sim.dt
        if self.scene.contact_forces is not None:
            self.scene.contact_forces.update_period = self.sim.dt
        # check if terrain levels curriculum is enabled - if so, enable curriculum for terrain generator
        # this generates terrains with increasing difficulty and is useful for training
        if getattr(self.curriculum, "terrain_levels", None) is not None:
            if self.scene.terrain.terrain_generator is not None:
                self.scene.terrain.terrain_generator.curriculum = True
        else:
            if self.scene.terrain.terrain_generator is not None:
                self.scene.terrain.terrain_generator.curriculum = False
    def disable_zero_weight_rewards(self):
        """If the weight of rewards is 0, set rewards to None"""
        for attr in dir(self.rewards):
            if not attr.startswith("__"):
                reward_attr = getattr(self.rewards, attr)
                if not callable(reward_attr) and reward_attr.weight == 0:
                    setattr(self.rewards, attr, None)
 def create_obsgroup_class(class_name, terms, enable_corruption=False, concatenate_terms=True):
    """
    Dynamically create and register a ObsGroup class based on the given configuration terms.
    :param class_name: Name of the configuration class.
    :param terms: Configuration terms, a dictionary where keys are term names and values are term content.
    :param enable_corruption: Whether to enable corruption for the observation group. Defaults to False.
    :param concatenate_terms: Whether to concatenate the observation terms in the group. Defaults to True.
    :return: The dynamically created class.
    """
    # Dynamically determine the module name
    module_name = inspect.getmodule(inspect.currentframe()).__name__
    # Define the post-init function
    def post_init_wrapper(self):
        setattr(self, "enable_corruption", enable_corruption)
        setattr(self, "concatenate_terms", concatenate_terms)
    # Dynamically create the class using ObsGroup as the base class
    terms["__post_init__"] = post_init_wrapper
    dynamic_class = configclass(type(class_name, (ObsGroup,), terms))
    # Custom serialization and deserialization
    def __getstate__(self):
        state = self.__dict__.copy()
        return state
    def __setstate__(self, state):
        self.__dict__.update(state)
    # Add custom serialization methods to the class
    dynamic_class.__getstate__ = __getstate__
    dynamic_class.__setstate__ = __setstate__
    # Place the class in the global namespace for accessibility
    globals()[class_name] = dynamic_class
    # Register the dynamic class in the module's dictionary
    if module_name in sys.modules:
        sys.modules[module_name].__dict__[class_name] = dynamic_class
    else:
        raise ImportError(f"Module {module_name} not found.")
    # Return the class for external instantiation
    return dynamic_class
@@ -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")
@@ -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,
 )